Improve tracer density

gusto/active_tracers.py

@@ -100,7 +100,8 @@ class WaterVapour(ActiveTracer):
     """An object encoding the details of water vapour as a tracer."""
     def __init__(self, name='water_vapour', space='theta',
                  variable_type=TracerVariableType.mixing_ratio,
-                 transport_eqn=TransportEquationType.advective):
+                 transport_eqn=TransportEquationType.advective,
+                 density_name=None):
         """
         Args:
             name (str, optional): the variable's name. Defaults to
@@ -113,16 +114,22 @@ def __init__(self, name='water_vapour', space='theta',
             transport_eqn (:class:`TransportEquationType`, optional): enumerator
                 indicating the type of transport equation to be solved (e.g.
                 advective). Defaults to `TransportEquationType.advective`.
+            density_name (str): the name of the associated density for a mixing
+                ratio when using the tracer_conservative transport. Defaults to None,
+                as this argument is not needed for other transport types.
         """
         super().__init__(f'{name}', space, variable_type,
-                         transport_eqn=transport_eqn, phase=Phases.gas, chemical='H2O')
+                         transport_eqn=transport_eqn,
+                         density_name=density_name,
+                         phase=Phases.gas, chemical='H2O')
 
 
 class CloudWater(ActiveTracer):
     """An object encoding the details of cloud water as a tracer."""
     def __init__(self, name='cloud_water', space='theta',
                  variable_type=TracerVariableType.mixing_ratio,
-                 transport_eqn=TransportEquationType.advective):
+                 transport_eqn=TransportEquationType.advective,
+                 density_name=None):
         """
         Args:
             name (str, optional): the variable name. Default is 'cloud_water'.
@@ -134,16 +141,22 @@ def __init__(self, name='cloud_water', space='theta',
             transport_eqn (:class:`TransportEquationType`, optional): enumerator
                 indicating the type of transport equation to be solved (e.g.
                 advective). Defaults to `TransportEquationType.advective`.
+            density_name (str): the name of the associated density for a mixing
+                ratio when using the tracer_conservative transport. Defaults to None,
+                as this argument is not needed for other transport types.
         """
         super().__init__(f'{name}', space, variable_type,
-                         transport_eqn=transport_eqn, phase=Phases.liquid, chemical='H2O')
+                         transport_eqn=transport_eqn,
+                         density_name=density_name,
+                         phase=Phases.liquid, chemical='H2O')
 
 
 class Rain(ActiveTracer):
     """An object encoding the details of rain as a tracer."""
     def __init__(self, name='rain', space='theta',
                  variable_type=TracerVariableType.mixing_ratio,
-                 transport_eqn=TransportEquationType.advective):
+                 transport_eqn=TransportEquationType.advective,
+                 density_name=None):
         """
         Args:
             name (str, optional): the name for the variable. Defaults to 'rain'.
@@ -155,6 +168,11 @@ def __init__(self, name='rain', space='theta',
             transport_eqn (:class:`TransportEquationType`, optional): enumerator
                 indicating the type of transport equation to be solved (e.g.
                 advective). Defaults to `TransportEquationType.advective`.
+            density_name (str): the name of the associated density for a mixing
+                ratio when using the tracer_conservative transport. Defaults to None,
+                as this argument is not needed for other transport types.
         """
         super().__init__(f'{name}', space, variable_type,
-                         transport_eqn=transport_eqn, phase=Phases.liquid, chemical='H2O')
+                         transport_eqn=transport_eqn,
+                         density_name=density_name,
+                         phase=Phases.liquid, chemical='H2O')

gusto/diagnostics.py

@@ -5,7 +5,8 @@
     LinearVariationalProblem, LinearVariationalSolver, FacetNormal, \
     ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, avg, jump, pi, \
     TensorFunctionSpace, SpatialCoordinate, as_vector, \
-    Projector, Interpolator, FunctionSpace
+    Projector, Interpolator, FunctionSpace, FiniteElement, \
+    TensorProductElement
 from firedrake.assign import Assigner
 
 from abc import ABCMeta, abstractmethod, abstractproperty
@@ -1691,7 +1692,12 @@ class TracerDensity(DiagnosticField):
     """Diagnostic for computing the density of a tracer. This is
     computed as the product of a mixing ratio and dry density"""
 
-    name = "TracerDensity"
+    @property
+    def name(self):
+        """Gives the name of this diagnostic field. This records
+        the mixing ratio and density names, in case multiple tracer
+        densities are used."""
+        return "TracerDensity_"+self.mixing_ratio_name+'_'+self.density_name
 
     def __init__(self, mixing_ratio_name, density_name, space=None, method='interpolate'):
         """
@@ -1700,12 +1706,15 @@ def __init__(self, mixing_ratio_name, density_name, space=None, method='interpol
             density_name (str): the name of the tracer density variable
             space (:class:`FunctionSpace`, optional): the function space to
                 evaluate the diagnostic field in. Defaults to None, in which
-                case a default space will be chosen for this diagnostic.
+                case a new space will be constructed for this diagnostic. This
+                space will have enough a high enough degree to accurately compute
+                the product of the mixing ratio and density.
             method (str, optional): a string specifying the method of evaluation
                 for this diagnostic. Valid options are 'interpolate', 'project' and
                 'assign'. Defaults to 'interpolate'.
         """
-        super().__init__(method=method, required_fields=(mixing_ratio_name, density_name))
+        super().__init__(space=space, method=method, required_fields=(mixing_ratio_name, density_name))
+
         self.mixing_ratio_name = mixing_ratio_name
         self.density_name = density_name
 
@@ -1720,4 +1729,40 @@ def setup(self, domain, state_fields):
         m_X = state_fields(self.mixing_ratio_name)
         rho_d = state_fields(self.density_name)
         self.expr = m_X*rho_d
-        super().setup(domain, state_fields)
+
+        if self.space is None:
+            # Construct a space for the diagnostic that has enough
+            # degrees to accurately capture the tracer density. This
+            # will be the sum of the degrees of the individual mixing ratio
+            # and density function spaces.
+            m_X_space = m_X.function_space()
+            rho_d_space = rho_d.function_space()
+
+            if domain.spaces.extruded_mesh:
+                # Extract the base horizontal and vertical elements
+                # for the mixing ratio and density.
+                m_X_horiz = m_X_space.ufl_element().sub_elements[0]
+                m_X_vert = m_X_space.ufl_element().sub_elements[1]
+                rho_d_horiz = rho_d_space.ufl_element().sub_elements[0]
+                rho_d_vert = rho_d_space.ufl_element().sub_elements[1]
+
+                horiz_degree = m_X_horiz.degree() + rho_d_horiz.degree()
+                vert_degree = m_X_vert.degree() + rho_d_vert.degree()
+
+                cell = domain.mesh._base_mesh.ufl_cell().cellname()
+                horiz_elt = FiniteElement('DG', cell, horiz_degree)
+                vert_elt = FiniteElement('DG', cell, vert_degree)
+                elt = TensorProductElement(horiz_elt, vert_elt)
+            else:
+                m_X_degree = m_X_space.ufl_element().degree()
+                rho_d_degree = rho_d_space.ufl_element().degree()
+                degree = m_X_degree + rho_d_degree
+
+                cell = domain.mesh.ufl_cell().cellname()
+                elt = FiniteElement('DG', cell, degree)
+
+            tracer_density_space = FunctionSpace(domain.mesh, elt, name='tracer_density_space')
+            super().setup(domain, state_fields, space=tracer_density_space)
+
+        else:
+            super().setup(domain, state_fields)