tie up final set of loose ends. More clearly separate circulation and…

… vector invariant types

gusto/common_forms.py

@@ -117,12 +117,7 @@ def vector_invariant_form(domain, test, q, ubar):
         class:`LabelledForm`: a labelled transport form.
     """
 
-    if domain.mesh.topological_dimension() == 3:
-        L = inner(test, cross(curl(q), ubar))*dx
-
-    else:
-        perp = domain.perp
-        L = inner(test, div(perp(q))*perp(ubar))*dx
+    L = advection_equation_circulation_form(domain, test, q, ubar).terms[0].form
 
     # Add K.E. term
     L -= 0.5*div(test)*inner(q, ubar)*dx
@@ -145,14 +140,12 @@ def kinetic_energy_form(test, q, ubar):
         ubar (:class:`ufl.Expr`): the transporting velocity.
 
     Returns:
-        class:`LabelledForm`: a labelled transport form.
+        class:`ufl.Form`: the kinetic energy form.
     """
 
-    L = 0.5*div(test)*inner(q, ubar)*dx
-
-    form = transporting_velocity(L, ubar)
+    L = -0.5*div(test)*inner(q, ubar)*dx
 
-    return transport(form, TransportEquationType.vector_invariant)
+    return L
 
 
 def advection_equation_circulation_form(domain, test, q, ubar):
@@ -182,12 +175,16 @@ def advection_equation_circulation_form(domain, test, q, ubar):
         class:`LabelledForm`: a labelled transport form.
     """
 
-    form = (
-        vector_invariant_form(domain, test, q, ubar)
-        - kinetic_energy_form(test, q, ubar)
-    )
+    if domain.mesh.topological_dimension() == 3:
+        L = inner(test, cross(curl(q), ubar))*dx
+
+    else:
+        perp = domain.perp
+        L = inner(test, div(perp(q))*perp(ubar))*dx
+
+    form = transporting_velocity(L, ubar)
 
-    return form
+    return transport(form, TransportEquationType.circulation)
 
 
 def diffusion_form(test, q, kappa):

gusto/configuration.py

@@ -50,12 +50,14 @@ class TransportEquationType(Enum):
     advective: ∂q/∂t + (u.∇)q = 0                                             \n
     conservative: ∂q/∂t + ∇.(u*q) = 0                                         \n
     vector_invariant: ∂q/∂t + (∇×q)×u + (1/2)∇(q.u) + (1/2)[(∇q).u -(∇u).q)] = 0
+    circulation: ∂q/∂t + (∇×q)×u + non-transport terms = 0
     """
 
     no_transport = 702
     advective = 19
     conservative = 291
     vector_invariant = 9081
+    circulation = 512
 
 
 class Configuration(object):

gusto/equations.py

@@ -9,14 +9,15 @@
 from gusto.fields import PrescribedFields
 from gusto.fml.form_manipulation_labelling import Term, all_terms, keep, drop, Label
 from gusto.labels import (subject, time_derivative, transport, prognostic,
-                          transporting_velocity, replace_subject, linearisation,
+                          replace_subject, linearisation,
                           name, pressure_gradient, coriolis, perp,
                           replace_trial_function, hydrostatic)
 from gusto.thermodynamics import exner_pressure
 from gusto.common_forms import (advection_form, continuity_form,
                                 vector_invariant_form, kinetic_energy_form,
                                 advection_equation_circulation_form,
-                                diffusion_form, linear_continuity_form)
+                                diffusion_form, linear_continuity_form,
+                                linear_advection_form)
 from gusto.active_tracers import ActiveTracer, Phases, TracerVariableType
 from gusto.configuration import TransportEquationType
 import ufl
@@ -631,12 +632,6 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None,
             u_adv = prognostic(advection_form(w, u, u), "u")
         elif u_transport_option == "circulation_form":
             ke_form = prognostic(kinetic_energy_form(w, u, u), "u")
-            ke_form = transport.remove(ke_form)
-            ke_form = ke_form.label_map(
-                lambda t: t.has_label(transporting_velocity),
-                lambda t: Term(ufl.replace(
-                    t.form, {t.get(transporting_velocity): u}), t.labels))
-            ke_form = transporting_velocity.remove(ke_form)
             u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), "u") + ke_form
         else:
             raise ValueError("Invalid u_transport_option: %s" % u_transport_option)
@@ -890,12 +885,6 @@ def __init__(self, domain, parameters, Omega=None, sponge=None,
             u_adv = prognostic(advection_form(w, u, u), "u")
         elif u_transport_option == "circulation_form":
             ke_form = prognostic(kinetic_energy_form(w, u, u), "u")
-            ke_form = transport.remove(ke_form)
-            ke_form = ke_form.label_map(
-                lambda t: t.has_label(transporting_velocity),
-                lambda t: Term(ufl.replace(
-                    t.form, {t.get(transporting_velocity): u}), t.labels))
-            ke_form = transporting_velocity.remove(ke_form)
             u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), "u") + ke_form
         else:
             raise ValueError("Invalid u_transport_option: %s" % u_transport_option)
@@ -1223,12 +1212,6 @@ def __init__(self, domain, parameters, Omega=None,
             u_adv = prognostic(advection_form(w, u, u), "u")
         elif u_transport_option == "circulation_form":
             ke_form = prognostic(kinetic_energy_form(w, u, u), "u")
-            ke_form = transport.remove(ke_form)
-            ke_form = ke_form.label_map(
-                lambda t: t.has_label(transporting_velocity),
-                lambda t: Term(ufl.replace(
-                    t.form, {t.get(transporting_velocity): u}), t.labels))
-            ke_form = transporting_velocity.remove(ke_form)
             u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), "u") + ke_form
         else:
             raise ValueError("Invalid u_transport_option: %s" % u_transport_option)

gusto/timeloop.py

@@ -377,7 +377,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
         if kwargs:
             raise ValueError("unexpected kwargs: %s" % list(kwargs.keys()))
 
-        self.spatial_methods = []
+        self.spatial_methods = spatial_methods
 
         if physics_schemes is not None:
             self.physics_schemes = physics_schemes
@@ -397,7 +397,6 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
             for method in spatial_methods:
                 if scheme.field_name == method.variable and method.term_label == transport:
                     method_found = True
-                    self.spatial_methods.append(method)
             assert method_found, f'No transport method found for variable {scheme.field_name}'
 
         self.diffusion_schemes = []
@@ -411,7 +410,6 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
                 for method in spatial_methods:
                     if scheme.field_name == method.variable and method.term_label == diffusion:
                         method_found = True
-                        self.diffusion_methods.append(method)
                 assert method_found, f'No diffusion method found for variable {scheme.field_name}'
 
         if auxiliary_equations_and_schemes is not None:
@@ -420,6 +418,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods,
             self.auxiliary_schemes = [
                 (eqn.field_name, scheme)
                 for eqn, scheme in auxiliary_equations_and_schemes]
+
         else:
             auxiliary_equations_and_schemes = []
             self.auxiliary_schemes = []

gusto/transport_methods.py

@@ -52,19 +52,30 @@ def replace_form(self, equation):
             lambda t: t.has_label(transport) and t.get(prognostic) == self.variable,
             map_if_true=keep, map_if_false=drop
         )
-        original_term = original_form.terms[0]
 
-        # Update transporting velocity
-        new_transporting_velocity = self.form.terms[0].get(transporting_velocity)
-        original_term = transporting_velocity.update_value(original_term, new_transporting_velocity)
+        if len(original_form.terms) == 0:
+            # This is likely not the appropriate equation so skip
+            pass
 
-        # Create new term
-        new_term = Term(self.form.form, original_term.labels)
+        elif len(original_form.terms) == 1:
+            # Replace form
+            original_term = original_form.terms[0]
 
-        # Replace original term with new term
-        equation.residual = equation.residual.label_map(
-            lambda t: t.has_label(transport) and t.get(prognostic) == self.variable,
-            map_if_true=lambda t: new_term)
+            # Update transporting velocity
+            new_transporting_velocity = self.form.terms[0].get(transporting_velocity)
+            original_term = transporting_velocity.update_value(original_term, new_transporting_velocity)
+
+            # Create new term
+            new_term = Term(self.form.form, original_term.labels)
+
+            # Replace original term with new term
+            equation.residual = equation.residual.label_map(
+                lambda t: t.has_label(transport) and t.get(prognostic) == self.variable,
+                map_if_true=lambda t: new_term)
+
+        else:
+            raise RuntimeError('Unable to find single transport term for '
+                               + f'{self.variable}. {len(original_form.terms)} found')
 
 
 # ---------------------------------------------------------------------------- #
@@ -150,6 +161,11 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE,
                 form = upwind_continuity_form(self.domain, self.test, self.field,
                                               ibp=ibp, outflow=outflow)
 
+        elif self.transport_equation_type == TransportEquationType.circulation:
+            if outflow:
+                raise NotImplementedError('Outflow not implemented for upwind circulation')
+            form = upwind_circulation_form(self.domain, self.test, self.field, ibp=ibp)
+
         elif self.transport_equation_type == TransportEquationType.vector_invariant:
             if outflow:
                 raise NotImplementedError('Outflow not implemented for upwind vector invariant')
@@ -358,20 +374,16 @@ def vector_manifold_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE,
     return transport(form)
 
 
-def upwind_vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE):
+def upwind_circulation_form(domain, test, q, ibp=IntegrateByParts.ONCE):
     u"""
-    The form corresponding to the DG upwind vector invariant transport operator.
+    The form corresponding to the DG upwind vector circulation operator.
 
     The self-transporting transport operator for a vector-valued field u can be
     written as circulation and kinetic energy terms:
     (u.∇)u = (∇×u)×u + (1/2)∇u^2
 
-    When the transporting field u and transported field q are similar, we write
-    this as:
-    (u.∇)q = (∇×q)×u + (1/2)∇(u.q)
-
-    This form discretises this final equation, using an upwind discretisation
-    when integrating by parts.
+    This form discretises the first term in this equation, (∇×u)×u, using an
+    upwind discretisation when integrating by parts.
 
     Args:
         domain (:class:`Domain`): the model's domain object, containing the
@@ -437,8 +449,46 @@ def upwind_vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE):
                              perp(ubar))*perp(q), n)*dS_
             )
 
-    L -= 0.5*div(test)*inner(q, ubar)*dx
+    form = transporting_velocity(L, ubar)
+
+    return transport(form, TransportEquationType.circulation)
+
+
+def upwind_vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE):
+    u"""
+    The form corresponding to the DG upwind vector invariant transport operator.
+
+    The self-transporting transport operator for a vector-valued field u can be
+    written as circulation and kinetic energy terms:
+    (u.∇)u = (∇×u)×u + (1/2)∇u^2
+
+    When the transporting field u and transported field q are similar, we write
+    this as:
+    (u.∇)q = (∇×q)×u + (1/2)∇(u.q)
+
+    This form discretises this final equation, using an upwind discretisation
+    when integrating 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`.
+
+    Raises:
+        NotImplementedError: the specified integration by parts is not 'once'.
+
+    Returns:
+        class:`LabelledForm`: a labelled transport form.
+    """
+
+    circulation_form = upwind_circulation_form(domain, test, q, ibp=ibp)
+    ubar = circulation_form.terms[0].get(transporting_velocity)
 
+    L = circulation_form.terms[0].form - 0.5*div(test)*inner(q, ubar)*dx
     form = transporting_velocity(L, ubar)
 
     return transport(form, TransportEquationType.vector_invariant)