Conservative projection (#508)

Co-authored-by: Tim Andrews <[email protected]>

gusto/core/__init__.py

@@ -1,11 +1,12 @@
-from gusto.core.configuration import *           # noqa
-from gusto.core.coordinates import *             # noqa
-from gusto.core.coord_transforms import *        # noqa
-from gusto.core.domain import *                  # noqa
-from gusto.core.fields import *                  # noqa
-from gusto.core.function_spaces import *         # noqa
-from gusto.core.io import *                      # noqa
-from gusto.core.kernels import *                 # noqa
-from gusto.core.labels import *                  # noqa
-from gusto.core.logging import *                 # noqa
-from gusto.core.meshes import *                  # noqa
+from gusto.core.configuration import *             # noqa
+from gusto.core.conservative_projection import *   # noqa
+from gusto.core.coordinates import *               # noqa
+from gusto.core.coord_transforms import *          # noqa
+from gusto.core.domain import *                    # noqa
+from gusto.core.fields import *                    # noqa
+from gusto.core.function_spaces import *           # noqa
+from gusto.core.io import *                        # noqa
+from gusto.core.kernels import *                   # noqa
+from gusto.core.labels import *                    # noqa
+from gusto.core.logging import *                   # noqa
+from gusto.core.meshes import *                    # noqa

gusto/core/configuration.py

@@ -8,7 +8,8 @@
     "IntegrateByParts", "TransportEquationType", "OutputParameters",
     "BoussinesqParameters", "CompressibleParameters",
     "ShallowWaterParameters",
-    "EmbeddedDGOptions", "RecoveryOptions", "SUPGOptions", "MixedFSOptions",
+    "EmbeddedDGOptions", "ConservativeEmbeddedDGOptions", "RecoveryOptions",
+    "ConservativeRecoveryOptions", "SUPGOptions", "MixedFSOptions",
     "SpongeLayerParameters", "DiffusionParameters", "BoundaryLayerParameters"
 ]
 
@@ -164,6 +165,14 @@ class EmbeddedDGOptions(WrapperOptions):
     embedding_space = None
 
 
+class ConservativeEmbeddedDGOptions(EmbeddedDGOptions):
+    """Specifies options for a conservative embedded DG method."""
+
+    project_back_method = 'conservative_project'
+    rho_name = None
+    orig_rho_space = None
+
+
 class RecoveryOptions(WrapperOptions):
     """Specifies options for a recovery wrapper method."""
 
@@ -177,6 +186,15 @@ class RecoveryOptions(WrapperOptions):
     broken_method = 'interpolate'
 
 
+class ConservativeRecoveryOptions(RecoveryOptions):
+    """Specifies options for a conservative recovery wrapper method."""
+
+    rho_name = None
+    orig_rho_space = None
+    project_high_method = 'conservative_project'
+    project_low_method = 'conservative_project'
+
+
 class SUPGOptions(WrapperOptions):
     """Specifies options for an SUPG scheme."""
 

gusto/core/conservative_projection.py

@@ -0,0 +1,93 @@
+"""
+This provides an operator for perform a conservative projection.
+
+The :class:`ConservativeProjector` provided in this module is an operator that
+projects a field such as a mixing ratio from one function space to another,
+weighted by a density field to ensure that mass is conserved by the projection.
+"""
+
+from firedrake import (Function, TestFunction, TrialFunction, lhs, rhs, inner,
+                       dx, LinearVariationalProblem, LinearVariationalSolver,
+                       Constant, assemble)
+import ufl
+
+__all__ = ["ConservativeProjector"]
+
+
+class ConservativeProjector(object):
+    """
+    Projects a field such that mass is conserved.
+
+    This object is designed for projecting fields such as mixing ratios of
+    tracer species from one function space to another, but weighted by density
+    such that mass is conserved by the projection.
+    """
+
+    def __init__(self, rho_source, rho_target, m_source, m_target,
+                 subtract_mean=False):
+        """
+        Args:
+            rho_source (:class:`Function`): the density to use for weighting the
+                source mixing ratio field. Can also be a :class:`ufl.Expr`.
+            rho_target (:class:`Function`): the density to use for weighting the
+                target mixing ratio field. Can also be a :class:`ufl.Expr`.
+            m_source (:class:`Function`): the source mixing ratio field. Can
+                also be a :class:`ufl.Expr`.
+            m_target (:class:`Function`): the target mixing ratio field to
+                compute.
+            subtract_mean (bool, optional): whether to solve the projection by
+                subtracting the mean value of m for both sides. This is more
+                expensive as it involves calculating the mean, but will ensure
+                preservation of a constant when projecting to a continuous
+                space. Default to False.
+
+        Raises:
+            RuntimeError: the geometric shape of the two rho fields must be equal.
+            RuntimeError: the geometric shape of the two m fields must be equal.
+        """
+
+        self.subtract_mean = subtract_mean
+
+        if not isinstance(rho_source, (ufl.core.expr.Expr, Function)):
+            raise ValueError("Can only recover UFL expression or Functions not '%s'" % type(rho_source))
+
+        if not isinstance(rho_target, (ufl.core.expr.Expr, Function)):
+            raise ValueError("Can only recover UFL expression or Functions not '%s'" % type(rho_target))
+
+        if not isinstance(m_source, (ufl.core.expr.Expr, Function)):
+            raise ValueError("Can only recover UFL expression or Functions not '%s'" % type(m_source))
+
+        # Check shape values
+        if m_source.ufl_shape != m_target.ufl_shape:
+            raise RuntimeError('Shape mismatch between source %s and target function spaces %s in project' % (m_source.ufl_shape, m_target.ufl_shape))
+
+        if rho_source.ufl_shape != rho_target.ufl_shape:
+            raise RuntimeError('Shape mismatch between source %s and target function spaces %s in project' % (rho_source.ufl_shape, rho_target.ufl_shape))
+
+        self.m_source = m_source
+        self.m_target = m_target
+
+        V = self.m_target.function_space()
+        mesh = V.mesh()
+
+        self.m_mean = Constant(0.0, domain=mesh)
+        self.volume = assemble(Constant(1.0, domain=mesh)*dx)
+
+        test = TestFunction(V)
+        m_trial = TrialFunction(V)
+        eqn = (rho_source*inner(test, m_source - self.m_mean)*dx
+               - rho_target*inner(test, m_trial - self.m_mean)*dx)
+        problem = LinearVariationalProblem(lhs(eqn), rhs(eqn), self.m_target)
+        self.solver = LinearVariationalSolver(problem)
+
+    def project(self):
+        """Apply the projection."""
+
+        # Compute mean value
+        if self.subtract_mean:
+            self.m_mean.assign(assemble(self.m_source*dx) / self.volume)
+
+        # Solve projection
+        self.solver.solve()
+
+        return self.m_target

gusto/equations/prognostic_equations.py

@@ -305,6 +305,17 @@ def add_tracers_to_prognostics(self, domain, active_tracers):
                 name of the active tracer.
         """
 
+        # Check if there are any conservatively transported tracers.
+        # If so, ensure that the reference density is indexed before this tracer.
+        for i in range(len(active_tracers) - 1):
+            tracer = active_tracers[i]
+            if tracer.transport_eqn == TransportEquationType.tracer_conservative:
+                ref_density = next(x for x in active_tracers if x.name == tracer.density_name)
+                j = active_tracers.index(ref_density)
+                if j > i:
+                    # Swap the indices of the tracer and the reference density
+                    active_tracers[i], active_tracers[j] = active_tracers[j], active_tracers[i]
+
         # Loop through tracer fields and add field names and spaces
         for tracer in active_tracers:
             if isinstance(tracer, ActiveTracer):

gusto/recovery/reversible_recovery.py

@@ -3,20 +3,24 @@
 higher-order function space.
 """
 
+from gusto.core.conservative_projection import ConservativeProjector
 from firedrake import (Projector, Function, Interpolator)
 from .recovery import Recoverer
 
+__all__ = ["ReversibleRecoverer", "ConservativeRecoverer"]
+
 
 class ReversibleRecoverer(object):
     """
     An object for performing a reconstruction of a low-order discontinuous
     field into a higher-order discontinuous space. This uses the recovery
     operator, but with further adjustments to ensure reversibility.
 
-    :arg source_field:      the source field.
-    :arg target_field:      the target_field.
-    :arg reconstruct_opts:  an object containing the various options for the
-                               reconstruction.
+    Args:
+        source_field (:class:`Function`): the source field.
+        target_field (:class:`Function`): the target field.
+        reconstruct_opts (:class:`RecoveryOptions`): an object containing the
+            various options for the reconstruction.
     """
     def __init__(self, source_field, target_field, reconstruct_opts):
 
@@ -92,3 +96,63 @@ def project(self):
         self.q_corr_low.assign(self.q_low - self.q_corr_low)
         self.injector.interpolate() if self.interp_inj else self.injector.project()
         self.q_high.assign(self.q_corr_high + self.q_rec_high)
+
+
+class ConservativeRecoverer(object):
+    """
+    An object for performing a reconstruction of a low-order discontinuous
+    field into a higher-order discontinuous space, but such that mass is
+    conserved. This uses the recovery operator, but with further adjustments to
+    ensure both reversibility and mass conservation.
+
+    Args:
+        source_field (:class:`Function`): the source field.
+        target_field (:class:`Function`): the target field.
+        source_density (:class:`Function`): the source density field.
+        target_density (:class:`Function`): the target density field.
+        reconstruct_opts (:class:`RecoveryOptions`): an object containing the
+            various options for the reconstruction.
+    """
+    def __init__(self, source_field, target_field, source_density,
+                 target_density, reconstruct_opts):
+
+        self.opts = reconstruct_opts
+
+        # Declare the fields used by the reconstructor
+        self.q_low = source_field
+        self.q_high = target_field
+        self.q_recovered = Function(self.opts.recovered_space)
+        self.q_corr_low = Function(source_field.function_space())
+        self.q_corr_high = Function(target_field.function_space())
+        self.q_rec_high = Function(target_field.function_space())
+
+        # -------------------------------------------------------------------- #
+        # Set up the operators for different transformations
+        # -------------------------------------------------------------------- #
+
+        # Does recovery by first projecting into broken space then averaging
+        self.recoverer = Recoverer(self.q_low, self.q_recovered,
+                                   method=self.opts.broken_method,
+                                   boundary_method=self.opts.boundary_method)
+
+        # Obtain the recovered field in the higher order space
+        self.projector_high = Projector(self.q_recovered, self.q_rec_high)
+
+        # Obtain the correction in the lower order space
+        # Swap density arguments!
+        self.projector_low = ConservativeProjector(target_density, source_density,
+                                                   self.q_rec_high, self.q_corr_low,
+                                                   subtract_mean=True)
+
+        # Final injection operator
+        # Should identify low order field in higher order space
+        self.injector = ConservativeProjector(source_density, target_density,
+                                              self.q_corr_low, self.q_corr_high)
+
+    def project(self):
+        self.recoverer.project()
+        self.projector_high.project()
+        self.projector_low.project()
+        self.q_corr_low.assign(self.q_low - self.q_corr_low)
+        self.injector.project()
+        self.q_high.assign(self.q_corr_high + self.q_rec_high)

gusto/time_discretisation/time_discretisation.py

@@ -94,7 +94,8 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                             'Time discretisation: suboption SUPG is currently not implemented within MixedOptions')
                     else:
                         raise RuntimeError(
-                            f'Time discretisation: suboption wrapper {self.wrapper_name} not implemented')
+                            f'Time discretisation: suboption wrapper {suboption.name} not implemented')
+
             elif self.wrapper_name == "embedded_dg":
                 self.wrapper = EmbeddedDGWrapper(self, options)
             elif self.wrapper_name == "recovered":