Generate linear solver using linearisation of equations (#649)

Author: tommbendall

examples/shallow_water/williamson_5.py

@@ -73,13 +73,14 @@ def williamson_5(
     rsq = min_value(R0**2, (lamda - lamda_c)**2 + (phi - phi_c)**2)
     r = sqrt(rsq)
     tpexpr = mountain_height * (1 - r/R0)
-    eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr,
-                                 topog_expr=tpexpr)
+    eqns = ShallowWaterEquations(
+        domain, parameters, fexpr=fexpr, topog_expr=tpexpr
+    )
 
     # I/O
     output = OutputParameters(
-        dirname=dirname, dumplist_latlon=['D'], dumpfreq=dumpfreq,
-        dump_vtus=True, dump_nc=False, dumplist=['D', 'topography']
+        dirname=dirname, dumpfreq=dumpfreq,
+        dump_vtus=True, dump_nc=True, dumplist=['D', 'topography']
     )
     diagnostic_fields = [Sum('D', 'topography'), RelativeVorticity(),
                          MeridionalComponent('u'), ZonalComponent('u')]
@@ -101,7 +102,7 @@ def williamson_5(
 
     # Time stepper
     stepper = SemiImplicitQuasiNewton(
-        eqns, io, transported_fields, transport_methods
+        eqns, io, transported_fields, transport_methods, tau_values={'D': 1.0}
     )
 
     # ------------------------------------------------------------------------ #

figures/shallow_water/williamson_5_final.png

@@ -1,15 +1,16 @@
 """Defines the Boussinesq equations."""
 
 from firedrake import inner, dx, div, cross, split, as_vector
-from firedrake.fml import subject
+from firedrake.fml import subject, all_terms
 from gusto.core.labels import (
-    time_derivative, transport, prognostic, linearisation,
+    prognostic, linearisation,
     pressure_gradient, coriolis, divergence, gravity, incompressible
 )
 from gusto.equations.common_forms import (
     advection_form, vector_invariant_form,
     kinetic_energy_form, advection_equation_circulation_form,
-    linear_advection_form
+    linear_advection_form, linear_circulation_form,
+    linear_vector_invariant_form
 )
 from gusto.equations.prognostic_equations import PrognosticEquationSet
 
@@ -41,7 +42,7 @@ class BoussinesqEquations(PrognosticEquationSet):
     def __init__(self, domain, parameters,
                  compressible=True,
                  space_names=None,
-                 linearisation_map='default',
+                 linearisation_map=all_terms,
                  u_transport_option="vector_invariant_form",
                  no_normal_flow_bc_ids=None,
                  active_tracers=None):
@@ -59,9 +60,8 @@ def __init__(self, domain, parameters,
                 in which case the spaces are taken from the de Rham complex.
             linearisation_map (func, optional): a function specifying which
                 terms in the equation set to linearise. If None is specified
-                then no terms are linearised. Defaults to the string 'default',
-                in which case the linearisation includes time derivatives and
-                scalar transport terms.
+                then no terms are linearised. Defaults to the FML `all_terms`
+                function.
             u_transport_option (str, optional): specifies the transport term
                 used for the velocity equation. Supported options are:
                 'vector_invariant_form', 'vector_advection_form' and
@@ -89,16 +89,6 @@ def __init__(self, domain, parameters,
         if active_tracers is None:
             active_tracers = []
 
-        if linearisation_map == 'default':
-            # Default linearisation is time derivatives, scalar transport,
-            # pressure gradient, gravity and divergence terms
-            # Don't include active tracers
-            linearisation_map = lambda t: \
-                t.get(prognostic) in ['u', 'p', 'b'] \
-                and (any(t.has_label(time_derivative, pressure_gradient,
-                                     divergence, gravity))
-                     or (t.get(prognostic) not in ['u', 'p'] and t.has_label(transport)))
-
         super().__init__(field_names, domain, space_names,
                          linearisation_map=linearisation_map,
                          no_normal_flow_bc_ids=no_normal_flow_bc_ids,
@@ -109,8 +99,8 @@ def __init__(self, domain, parameters,
 
         w, phi, gamma = self.tests[0:3]
         u, p, b = split(self.X)
-        u_trial, p_trial, _ = split(self.trials)
-        _, p_bar, b_bar = split(self.X_ref)
+        u_trial, p_trial, b_trial = split(self.trials)[0:3]
+        u_bar, p_bar, b_bar = split(self.X_ref)[0:3]
 
         # -------------------------------------------------------------------- #
         # Time Derivative Terms
@@ -122,28 +112,51 @@ def __init__(self, domain, parameters,
         # -------------------------------------------------------------------- #
         # Velocity transport term -- depends on formulation
         if u_transport_option == "vector_invariant_form":
-            u_adv = prognostic(vector_invariant_form(domain, w, u, u), 'u')
-        elif u_transport_option == "vector_advection_form":
-            u_adv = prognostic(advection_form(w, u, u), 'u')
+            u_adv = prognostic(vector_invariant_form(self.domain, w, u, u), 'u')
+            # Manually add linearisation, as linearisation cannot handle the
+            # perp function on the plane / vertical slice
+            if self.linearisation_map(u_adv.terms[0]):
+                linear_u_adv = linear_vector_invariant_form(self.domain, w, u_trial, u_bar)
+                u_adv = linearisation(u_adv, linear_u_adv)
+
         elif u_transport_option == "circulation_form":
+            # This is different to vector invariant form as the K.E. form
+            # doesn't have a variable marked as "transporting velocity"
             ke_form = prognostic(kinetic_energy_form(w, u, u), 'u')
-            u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), 'u') + ke_form
+            circ_form = prognostic(advection_equation_circulation_form(self.domain, w, u, u), 'u')
+            # Manually add linearisation, as linearisation cannot handle the
+            # perp function on the plane / vertical slice
+            if self.linearisation_map(circ_form.terms[0]):
+                linear_circ_form = linear_circulation_form(self.domain, w, u_trial, u_bar)
+                circ_form = linearisation(circ_form, linear_circ_form)
+            u_adv = circ_form + ke_form
+
+        elif u_transport_option == "vector_advection_form":
+            u_adv = prognostic(advection_form(w, u, u), 'u')
+
         else:
-            raise ValueError("Invalid u_transport_option: %s" % u_transport_option)
+            raise ValueError("Invalid u_transport_option: %s" % self.u_transport_option)
 
         # Buoyancy transport
         b_adv = prognostic(advection_form(gamma, b, u), 'b')
+
+        # TODO #651: we should remove this hand-coded linearisation
+        # currently REXI can't handle generated transport linearisations
         if self.linearisation_map(b_adv.terms[0]):
-            linear_b_adv = linear_advection_form(gamma, b_bar, u_trial)
+            linear_b_adv = linear_advection_form(gamma, b_trial, u_trial, b_bar, u_bar)
             b_adv = linearisation(b_adv, linear_b_adv)
 
         if compressible:
             # Pressure transport
             p_adv = prognostic(advection_form(phi, p, u), 'p')
+
+            # TODO #651: we should remove this hand-coded linearisation
+            # currently REXI can't handle generated transport linearisations
             if self.linearisation_map(p_adv.terms[0]):
-                linear_p_adv = linear_advection_form(phi, p_bar, u_trial)
+                linear_p_adv = linear_advection_form(phi, p_trial, u_trial, p_bar, u_bar)
                 p_adv = linearisation(p_adv, linear_p_adv)
             adv_form = subject(u_adv + p_adv + b_adv, self.X)
+
         else:
             adv_form = subject(u_adv + b_adv, self.X)
 
@@ -193,11 +206,11 @@ def __init__(self, domain, parameters,
         # Extra Terms (Coriolis)
         # -------------------------------------------------------------------- #
         if self.parameters.Omega is not None:
-            # TODO: add linearisation
             Omega = as_vector((0, 0, self.parameters.Omega))
             coriolis_form = coriolis(subject(prognostic(
                 inner(w, cross(2*Omega, u))*dx, 'u'), self.X))
             residual += coriolis_form
+
         # -------------------------------------------------------------------- #
         # Linearise equations
         # -------------------------------------------------------------------- #
@@ -230,7 +243,7 @@ class LinearBoussinesqEquations(BoussinesqEquations):
     def __init__(self, domain, parameters,
                  compressible=True,
                  space_names=None,
-                 linearisation_map='default',
+                 linearisation_map=all_terms,
                  u_transport_option="vector_invariant_form",
                  no_normal_flow_bc_ids=None,
                  active_tracers=None):
@@ -249,8 +262,8 @@ def __init__(self, domain, parameters,
             linearisation_map (func, optional): a function specifying which
                 terms in the equation set to linearise. If None is specified
                 then no terms are linearised. Defaults to the string 'default',
-                in which case the linearisation includes time derivatives and
-                scalar transport terms.
+                in which case the linearisation drops terms for any active
+                tracers.
             u_transport_option (str, optional): specifies the transport term
                 used for the velocity equation. Supported options are:
                 'vector_invariant_form', 'vector_advection_form' and
@@ -267,13 +280,6 @@ def __init__(self, domain, parameters,
             NotImplementedError: active tracers are not implemented.
         """
 
-        if linearisation_map == 'default':
-            # Default linearisation is time derivatives, pressure gradient,
-            # Coriolis and transport term from depth equation
-            linearisation_map = lambda t: \
-                (any(t.has_label(time_derivative, pressure_gradient, coriolis,
-                                 gravity, divergence, incompressible))
-                 or (t.get(prognostic) in ['p', 'b'] and t.has_label(transport)))
         super().__init__(domain=domain,
                          parameters=parameters,
                          compressible=compressible,

gusto/equations/boussinesq_equations.py

@@ -12,7 +12,9 @@
 
 __all__ = ["advection_form", "advection_form_1d", "continuity_form",
            "continuity_form_1d", "vector_invariant_form",
+           "linear_vector_invariant_form",
            "kinetic_energy_form", "advection_equation_circulation_form",
+           "linear_circulation_form",
            "diffusion_form", "diffusion_form_1d",
            "linear_advection_form", "linear_continuity_form",
            "split_continuity_form", "tracer_conservative_form", "split_hv_advective_form"]
@@ -60,21 +62,25 @@ def advection_form_1d(test, q, ubar):
     return transport(form, TransportEquationType.advective)
 
 
-def linear_advection_form(test, qbar, ubar):
+def linear_advection_form(test, q, u, qbar, ubar):
     """
     The form corresponding to the linearised advective transport operator.
 
     Args:
         test (:class:`TestFunction`): the test function.
-        qbar (:class:`ufl.Expr`): the variable to be transported.
-        ubar (:class:`ufl.Expr`): the transporting velocity.
+        q (:class:`ufl.Expr`): the perturbation variable to be transported.
+        u (:class:`ufl.Expr`): the perturbation transporting velocity.
+        qbar (:class:`ufl.Expr`): the mean variable to be transported.
+        ubar (:class:`ufl.Expr`): the mean transporting velocity.
 
     Returns:
         :class:`LabelledForm`: a labelled transport form.
     """
 
-    L = test*dot(ubar, grad(qbar))*dx
-    form = transporting_velocity(L, ubar)
+    form = (
+        transporting_velocity(test*dot(ubar, grad(q))*dx, ubar)
+        + transporting_velocity(test*dot(u, grad(qbar))*dx, u)
+    )
 
     return transport(form, TransportEquationType.advective)
 
@@ -121,21 +127,25 @@ def continuity_form_1d(test, q, ubar):
     return transport(form, TransportEquationType.conservative)
 
 
-def linear_continuity_form(test, qbar, ubar):
+def linear_continuity_form(test, q, u, qbar, ubar):
     """
     The form corresponding to the linearised continuity transport operator.
 
     Args:
         test (:class:`TestFunction`): the test function.
-        qbar (:class:`ufl.Expr`): the variable to be transported.
-        ubar (:class:`ufl.Expr`): the transporting velocity.
+        q (:class:`ufl.Expr`): the perturbation variable to be transported.
+        u (:class:`ufl.Expr`): the perturbation transporting velocity.
+        qbar (:class:`ufl.Expr`): the mean variable to be transported.
+        ubar (:class:`ufl.Expr`): the mean transporting velocity.
 
     Returns:
         :class:`LabelledForm`: a labelled transport form.
     """
 
-    L = test*div(qbar*ubar)*dx
-    form = transporting_velocity(L, ubar)
+    form = (
+        transporting_velocity(test*div(q*ubar)*dx, ubar)
+        + transporting_velocity(test*div(qbar*u)*dx, u)
+    )
 
     return transport(form, TransportEquationType.conservative)
 
@@ -172,6 +182,34 @@ def vector_invariant_form(domain, test, q, ubar):
     return transport(form, TransportEquationType.vector_invariant)
 
 
+def linear_vector_invariant_form(domain, test, q, ubar):
+    u"""
+    The linear form corresponding to the vector invariant transport operator.
+
+    The vector invariant transport operator is: (∇×q)×u + (1/2)∇(u.q)
+    and its linearised form is:
+    (∇×q')×u_bar + (∇×q_bar)×u' + (1/2)∇(u_bar.q') + (1/2)∇(u'.q_bar)
+
+    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.
+        ubar (:class:`ufl.Expr`): the transporting velocity.
+
+    Returns:
+        class:`LabelledForm`: a labelled transport form.
+    """
+
+    L = linear_circulation_form(domain, test, q, ubar).form
+
+    # Add K.E. term
+    L -= div(test)*inner(q, ubar)*dx
+    form = transporting_velocity(L, ubar)
+
+    return transport(form, TransportEquationType.vector_invariant)
+
+
 def kinetic_energy_form(test, q, ubar):
     u"""
     The form corresponding to the kinetic energy term.
@@ -210,6 +248,7 @@ def advection_equation_circulation_form(domain, test, q, ubar):
     term.
 
     Args:
+        domain (:class:`Domain`): the model's domain object.
         test (:class:`TestFunction`): the test function.
         q (:class:`ufl.Expr`): the variable to be transported.
         ubar (:class:`ufl.Expr`): the transporting velocity.
@@ -230,6 +269,26 @@ def advection_equation_circulation_form(domain, test, q, ubar):
     return transport(form, TransportEquationType.circulation)
 
 
+def linear_circulation_form(domain, test, q, ubar):
+    """
+    The linear circulation term in the transport of a vector-valued field.
+
+    Args:
+        test (:class:`TestFunction`): the test function.
+        q (:class:`ufl.Expr`): the variable to be transported.
+        ubar (:class:`ufl.Expr`): the transporting velocity.
+
+    Returns:
+        class:`LabelledForm`: a labelled transport form.
+    """
+
+    form = (
+        advection_equation_circulation_form(domain, test, q, ubar)
+        + advection_equation_circulation_form(domain, test, ubar, q)
+    )
+    return form
+
+
 def diffusion_form(test, q, kappa):
     u"""
     The diffusion form, -∇.(κ∇q) for diffusivity κ and variable q.