Horizontal & vertical splitting of advective transport terms, and a DG_Upwind version

gusto/core/labels.py

@@ -92,6 +92,8 @@ def __call__(self, target, value=None):
 # ---------------------------------------------------------------------------- #
 implicit = Label("implicit")
 explicit = Label("explicit")
+horizontal = Label("horizontal")
+vertical = Label("vertical")
 source_label = Label("source_label")
 transporting_velocity = Label("transporting_velocity", validator=lambda value: type(value) in [Function, ufl.tensors.ListTensor, ufl.indexed.Indexed])
 prognostic = Label("prognostic", validator=lambda value: type(value) == str)

gusto/equations/common_forms.py

@@ -7,14 +7,14 @@
 from firedrake.fml import subject, drop
 from gusto.core.configuration import TransportEquationType
 from gusto.core.labels import (transport, transporting_velocity, diffusion,
-                               prognostic, linearisation)
+                               prognostic, linearisation, horizontal, vertical)
 
 __all__ = ["advection_form", "advection_form_1d", "continuity_form",
            "continuity_form_1d", "vector_invariant_form",
            "kinetic_energy_form", "advection_equation_circulation_form",
            "diffusion_form", "diffusion_form_1d",
            "linear_advection_form", "linear_continuity_form",
-           "split_continuity_form", "tracer_conservative_form"]
+           "split_continuity_form", "tracer_conservative_form", "split_hv_advective_form"]
 
 
 def advection_form(test, q, ubar):
@@ -346,3 +346,69 @@ def tracer_conservative_form(test, q, rho, ubar):
     form = transporting_velocity(L, ubar)
 
     return transport(form, TransportEquationType.tracer_conservative)
+
+
+def split_hv_advective_form(equation, field_name):
+    u"""
+    Splits advective term into horizontal and vertical terms.
+    This describes splitting u.∇(q) terms into u.(∇_h)q and w dq/dz,
+    for transporting velocity u and transported q.
+    Args:
+        equation (:class:`PrognosticEquation`): the model's equation.
+    Returns:
+        :class:`PrognosticEquation`: the model's equation.
+    """
+    k = equation.domain.k   # vertical unit vector
+    for t in equation.residual:
+        if (t.get(transport) == TransportEquationType.advective and t.get(prognostic) == field_name):
+            # Get fields and test functions
+            subj = t.get(subject)
+
+            # u is either a prognostic or prescribed field
+            if (hasattr(equation, "field_names")
+               and 'u' in equation.field_names):
+                idx = equation.field_names.index(field_name)
+                W = equation.function_space
+                test = TestFunctions(W)[idx]
+                q = split(subj)[idx]
+                u_idx = equation.field_names.index('u')
+                uadv = split(equation.X)[u_idx]
+            elif 'u' in equation.prescribed_fields._field_names:
+                uadv = equation.prescribed_fields('u')
+                q = subj
+                W = equation.function_space
+                test = TestFunction(W)
+            else:
+                raise ValueError('Cannot get velocity field')
+
+            # Create new advective and divergence terms
+            u_vertical = k*inner(uadv, k)
+            u_horizontal = uadv - u_vertical
+
+            vertical_adv_term = prognostic(vertical(transport(transporting_velocity(inner(test, dot(u_vertical, grad(q)))*dx, uadv), TransportEquationType.advective)), field_name)
+            horizontal_adv_term = prognostic(horizontal(transport(transporting_velocity(inner(test, dot(u_horizontal, grad(q)))*dx, uadv), TransportEquationType.advective)), field_name)
+
+            # Add linearisations of new terms if required
+            if (t.has_label(linearisation)):
+                u_trial = TrialFunctions(W)[u_idx]
+                u_trial_vert = k*inner(u_trial, k)
+                u_trial_horiz = u_trial - u_trial_vert
+                qbar = split(equation.X_ref)[idx]
+                # Add linearisation to adv_term
+                linear_hori_term = horizontal(transport(transporting_velocity(test*dot(u_trial_horiz, grad(qbar))*dx, u_trial), TransportEquationType.advective))
+                adv_horiz_term = linearisation(horizontal_adv_term, linear_hori_term)
+                # Add linearisation to div_term
+                linear_vert_term = horizontal(transport(transporting_velocity(test*dot(u_trial_vert, grad(qbar))*dx, u_trial), TransportEquationType.advective))
+                adv_vert_term = linearisation(vertical_adv_term, linear_vert_term)
+            else:
+                adv_vert_term = vertical_adv_term
+                adv_horiz_term = horizontal_adv_term
+            # Drop old term
+            equation.residual = equation.residual.label_map(
+                lambda t: t.get(transport) == TransportEquationType.advective and t.get(prognostic) == field_name,
+                map_if_true=drop)
+
+            # Add new terms onto residual
+            equation.residual += subject(adv_horiz_term, subj) + subject(adv_vert_term, subj)
+
+    return equation

gusto/spatial_methods/spatial_methods.py

@@ -44,7 +44,7 @@ def __init__(self, equation, variable, term_label):
             map_if_true=keep, map_if_false=drop)
 
         num_terms = len(self.original_form.terms)
-        assert num_terms == 1, f'Unable to find {term_label.label} term ' \
+        assert num_terms >= 1, f'Unable to find {term_label.label} term ' \
             + f'for {variable}. {num_terms} found'
 
     def replace_form(self, equation):

gusto/spatial_methods/transport_methods.py

@@ -10,12 +10,12 @@
 from gusto.core.configuration import IntegrateByParts, TransportEquationType
 from gusto.core.labels import (
     prognostic, transport, transporting_velocity, ibp_label, mass_weighted,
-    all_but_last
+    all_but_last, horizontal, vertical, explicit
 )
 from gusto.core.logging import logger
 from gusto.spatial_methods.spatial_methods import SpatialMethod
 
-__all__ = ["DefaultTransport", "DGUpwind"]
+__all__ = ["DefaultTransport", "DGUpwind", "Split_DGUpwind"]
 
 
 # ---------------------------------------------------------------------------- #
@@ -111,8 +111,41 @@ def replace_form(self, equation):
                 map_if_true=lambda t: new_term)
 
         else:
-            raise RuntimeError('Found multiple transport terms for '
-                               + f'{self.variable}. {len(original_form.terms)} found')
+            horizontal_form = equation.residual.label_map(
+                lambda t: t.has_label(transport) and t.has_label(horizontal) and t.get(prognostic) == self.variable,
+                map_if_true=keep, map_if_false=drop
+            )
+            vertical_form = equation.residual.label_map(
+                lambda t: t.has_label(transport) and t.has_label(vertical) and t.get(prognostic) == self.variable,
+                map_if_true=keep, map_if_false=drop
+            )
+
+            # Replace form
+            horizontal_term = horizontal_form.terms[0]
+            vertical_term = vertical_form.terms[0]
+
+            # Update transporting velocity
+            new_horizontal_transporting_velocity = self.form_h.terms[0].get(transporting_velocity)
+            new_vertical_transporting_velocity = self.form_v.terms[0].get(transporting_velocity)
+            horizontal_term = transporting_velocity.update_value(horizontal_term, new_horizontal_transporting_velocity)
+            vertical_term = transporting_velocity.update_value(vertical_term, new_vertical_transporting_velocity)
+
+            # Create new term
+            new_horizontal_term = Term(self.form_h.form, horizontal_term.labels)
+            new_vertical_term = Term(self.form_v.form, vertical_term.labels)
+
+            # Check if this is a conservative transport
+            if horizontal_term.has_label(mass_weighted) or vertical_term.has_label(mass_weighted):
+                raise RuntimeError('Mass weighted transport terms not yet supported for multiple terms')
+
+            # Replace original terms with new terms
+            equation.residual = equation.residual.label_map(
+                lambda t: t.has_label(transport) and t.has_label(horizontal) and t.get(prognostic) == self.variable,
+                map_if_true=lambda _: new_horizontal_term)
+
+            equation.residual = equation.residual.label_map(
+                lambda t: t.has_label(transport) and t.has_label(vertical) and t.get(prognostic) == self.variable,
+                map_if_true=lambda _: new_vertical_term)
 
 
 # ---------------------------------------------------------------------------- #
@@ -277,9 +310,133 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE,
         self.form = form
 
 
+class Split_DGUpwind(TransportMethod):
+    """
+    The Discontinuous Galerkin Upwind transport scheme applied separately in the
+    horizontal and vertical directions.
+    Discretises the gradient of a field weakly, taking the upwind value of the
+    transported variable at facets.
+    """
+    def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE,
+                 vector_manifold_correction=False, outflow=False):
+        """
+        Args:
+            equation (:class:`PrognosticEquation`): the equation, which includes
+                a transport term.
+            variable (str): name of the variable to set the transport scheme for
+            ibp (:class:`IntegrateByParts`, optional): an enumerator for how
+                many times to integrate by parts. Defaults to `ONCE`.
+            vector_manifold_correction (bool, optional): whether to include a
+                vector manifold correction term. Defaults to False.
+            outflow (bool, optional): whether to include outflow at the domain
+                boundaries, through exterior facet terms. Defaults to False.
+        """
+
+        super().__init__(equation, variable)
+        self.ibp = ibp
+        self.vector_manifold_correction = vector_manifold_correction
+        self.outflow = outflow
+
+        # -------------------------------------------------------------------- #
+        # Determine appropriate form to use
+        # -------------------------------------------------------------------- #
+        # first check for 1d mesh and scalar velocity space
+        if equation.domain.on_sphere:
+            raise NotImplementedError('Split hv Upwind transport scheme has not been '
+                                      + 'implemented for spherical geometry')
+        if equation.domain.mesh.topological_dimension() == 1 and len(equation.domain.spaces("HDiv").shape) == 0:
+            assert not vector_manifold_correction
+            raise ValueError('You cannot do horizontal and vertical splitting in 1D')
+
+        else:
+            if self.transport_equation_type == TransportEquationType.advective:
+
+                form_h, form_v = split_upwind_advection_form(self.domain, self.test,
+                                                             self.field,
+                                                             ibp=ibp, outflow=outflow)
+
+            else:
+                raise NotImplementedError('Split hv Upwind transport scheme has not been '
+                                          + 'implemented for this transport equation type')
+        self.form_v = form_v
+        self.form_h = form_h
+
+
 # ---------------------------------------------------------------------------- #
 # Forms for DG Upwind transport
 # ---------------------------------------------------------------------------- #
+def split_upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False):
+    u"""
+    The forms corresponding to the DG upwind advective transport operator in
+    the horizontal and vertical directions.
+    This discretises u_h.(∇_h)q and w dq/dz, for transporting velocity u and transported
+    variable q. An upwind discretisation is used for the facet terms when the
+    form is integrated by parts.
+    Args:
+        domain (:class:`Domain`): the model's domain object, containing the
+            mesh and the compatible function spaces.
+        test (:class:`TestFunction`): the test function.
+        q (:class:`ufl.Expr`): the variable to be transported.
+        ibp (:class:`IntegrateByParts`, optional): an enumerator representing
+            the number of times to integrate by parts. Defaults to
+            `IntegrateByParts.ONCE`.
+        outflow (bool, optional): whether to include outflow at the domain
+            boundaries, through exterior facet terms. Defaults to False.
+    Raises:
+        ValueError: Can only use outflow option when the integration by parts
+            option is not "never".
+    Returns:
+        class:`LabelledForm`: a labelled transport form.
+    """
+
+    if outflow and ibp == IntegrateByParts.NEVER:
+        raise ValueError("outflow is True and ibp is None are incompatible options")
+    Vu = domain.spaces("HDiv")
+    k = domain.k
+    dS_ = (dS_v + dS_h) if Vu.extruded else dS
+    ubar = Function(Vu)
+    ubar_v = k*inner(ubar, k)
+    ubar_h = ubar - ubar_v
+
+    if ibp == IntegrateByParts.ONCE:
+        L_h = -inner(div(outer(test, ubar_h)), q)*dx
+        L_v = -inner(div(outer(test, ubar_v)), q)*dx
+    else:
+        L_h = inner(outer(test, ubar_h), grad(q))*dx
+        L_v = inner(outer(test, ubar_v), grad(q))*dx
+
+    if ibp != IntegrateByParts.NEVER:
+        n = FacetNormal(domain.mesh)
+        un_h = 0.5*(dot(ubar_h, n) + abs(dot(ubar_h, n)))
+
+        L_h += dot(jump(test), (un_h('+')*q('+') - un_h('-')*q('-')))*dS_
+
+        un_v = 0.5*(dot(ubar_v, n) + abs(dot(ubar_v, n)))
+
+        L_v += dot(jump(test), (un_v('+')*q('+') - un_v('-')*q('-')))*dS_
+
+        if ibp == IntegrateByParts.TWICE:
+            L_h -= (inner(test('+'), dot(ubar_h('+'), n('+')) * q('+'))
+                    + inner(test('-'), dot(ubar_h('-'), n('-')) * q('-'))) * dS_
+
+            L_v -= (inner(test('+'), dot(ubar_v('+'), n('+')) * q('+'))
+                    + inner(test('-'), dot(ubar_v('-'), n('-')) * q('-'))) * dS_
+
+    if outflow:
+        n = FacetNormal(domain.mesh)
+        un_h = 0.5*(dot(ubar_h, n) + abs(dot(ubar_h, n)))
+        L_h += test*un_h*q*(ds_v + ds_t + ds_b)
+
+        un_v = 0.5*(dot(ubar_v, n) + abs(dot(ubar_v, n)))
+        L_v += test*un_v*q*(ds_v + ds_t + ds_b)
+
+    form_h = transporting_velocity(L_h, ubar)
+    form_v = transporting_velocity(L_v, ubar)
+    labelled_form_h = ibp_label(transport(explicit(form_h), TransportEquationType.advective), ibp)
+    labelled_form_v = ibp_label(transport(form_v, TransportEquationType.advective), ibp)
+    return labelled_form_h, labelled_form_v
+
+
 def upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False):
     u"""
     The form corresponding to the DG upwind advective transport operator.

integration-tests/transport/test_split_dg_transport.py

@@ -0,0 +1,64 @@
+"""
+Tests the Split horizontal and vertical DG upwind transport scheme for
+advective form transport equation. This tests that the
+field is transported to the correct place.
+"""
+
+from firedrake import norm, VectorFunctionSpace, as_vector
+from gusto import *
+
+
+def run(timestepper, tmax, f_end):
+    timestepper.run(0, tmax)
+    return norm(timestepper.fields("f") - f_end) / norm(f_end)
+
+
+def test_split_dg_transport_scalar(tmpdir, tracer_setup):
+    setup = tracer_setup(tmpdir, "slice")
+    domain = setup.domain
+    V = domain.spaces("DG")
+
+    eqn = AdvectionEquation(domain, V, "f")
+    eqn = split_hv_advective_form(eqn, "f")
+
+    transport_method = Split_DGUpwind(eqn, "f")
+    transport_scheme = SSPRK3(domain)
+
+    time_varying_velocity = False
+    timestepper = PrescribedTransport(
+        eqn, transport_scheme, setup.io, time_varying_velocity, transport_method
+    )
+
+    # Initial conditions
+    timestepper.fields("f").interpolate(setup.f_init)
+    timestepper.fields("u").project(setup.uexpr)
+
+    error = run(timestepper, setup.tmax, setup.f_end)
+    assert error < setup.tol, \
+        'The transport error is greater than the permitted tolerance'
+
+
+def test_split_dg_transport_vector(tmpdir, tracer_setup):
+    setup = tracer_setup(tmpdir, "slice")
+    domain = setup.domain
+    gdim = domain.mesh.geometric_dimension()
+    f_init = as_vector([setup.f_init]*gdim)
+    V = VectorFunctionSpace(domain.mesh, "DG", 1)
+    eqn = AdvectionEquation(domain, V, "f")
+    eqn = split_hv_advective_form(eqn, "f")
+
+    transport_scheme = SSPRK3(domain)
+    transport_method = Split_DGUpwind(eqn, "f")
+
+    time_varying_velocity = False
+    timestepper = PrescribedTransport(
+        eqn, transport_scheme, setup.io, time_varying_velocity, transport_method
+    )
+
+    # Initial conditions
+    timestepper.fields("f").interpolate(f_init)
+    timestepper.fields("u").project(setup.uexpr)
+    f_end = as_vector([setup.f_end]*gdim)
+    error = run(timestepper, setup.tmax, f_end)
+    assert error < setup.tol, \
+        'The transport error is greater than the permitted tolerance'