Fix the hydrostatic switch

examples/compressible_euler/skamarock_klemp_nonhydrostatic.py

@@ -26,7 +26,8 @@
     logger, SUPGOptions, Perturbation, CompressibleParameters,
     CompressibleEulerEquations, HydrostaticCompressibleEulerEquations,
     compressible_hydrostatic_balance, RungeKuttaFormulation, CompressibleSolver,
-    SubcyclingOptions, hydrostatic_parameters
+    SubcyclingOptions, dx, TestFunction, TrialFunction, ZComponent,
+    LinearVariationalProblem, LinearVariationalSolver
 )
 
 skamarock_klemp_nonhydrostatic_defaults = {
@@ -67,6 +68,8 @@ def skamarock_klemp_nonhydrostatic(
     # ------------------------------------------------------------------------ #
 
     element_order = 1
+    alpha = 0.5
+    u_eqn_type = 'vector_advection_form'
 
     # ------------------------------------------------------------------------ #
     # Set up model objects
@@ -80,7 +83,9 @@ def skamarock_klemp_nonhydrostatic(
     # Equation
     parameters = CompressibleParameters()
     if hydrostatic:
-        eqns = HydrostaticCompressibleEulerEquations(domain, parameters)
+        eqns = HydrostaticCompressibleEulerEquations(
+            domain, parameters, u_transport_option=u_eqn_type
+        )
     else:
         eqns = CompressibleEulerEquations(domain, parameters)
 
@@ -110,7 +115,7 @@ def skamarock_klemp_nonhydrostatic(
             dump_vtus=False, dump_nc=True,
         )
 
-    diagnostic_fields = [Perturbation('theta')]
+    diagnostic_fields = [Perturbation('theta'), ZComponent('u')]
     io = IO(domain, output, diagnostic_fields=diagnostic_fields)
 
     # Transport schemes
@@ -136,16 +141,15 @@ def skamarock_klemp_nonhydrostatic(
     # Linear solver
     if hydrostatic:
         linear_solver = CompressibleSolver(
-            eqns, solver_parameters=hydrostatic_parameters,
-            overwrite_solver_parameters=True
+            eqns, alpha=alpha
         )
     else:
         linear_solver = CompressibleSolver(eqns)
 
     # Time stepper
     stepper = SemiImplicitQuasiNewton(
         eqns, io, transported_fields, transport_methods,
-        linear_solver=linear_solver
+        linear_solver=linear_solver, alpha=alpha, num_outer=2, num_inner=2
     )
 
     # ------------------------------------------------------------------------ #
@@ -175,12 +179,23 @@ def skamarock_klemp_nonhydrostatic(
     # Calculate hydrostatic exner
     compressible_hydrostatic_balance(eqns, theta_b, rho_b)
 
+    # Define initial theta
     theta_pert = (
         deltaTheta * sin(pi*z/domain_height)
         / (1 + (x - domain_width/2)**2 / pert_width**2)
     )
     theta0.interpolate(theta_b + theta_pert)
-    rho0.assign(rho_b)
+
+    # find perturbed rho
+    gamma = TestFunction(Vr)
+    rho_trial = TrialFunction(Vr)
+    dx_qp = dx(degree=domain.max_quad_degree)
+    lhs = gamma * rho_trial * dx_qp
+    rhs = gamma * (rho_b * theta_b / theta0) * dx_qp
+    rho_problem = LinearVariationalProblem(lhs, rhs, rho0)
+    rho_solver = LinearVariationalSolver(rho_problem)
+    rho_solver.solve()
+
     u0.project(as_vector([wind_initial, 0.0]))
 
     stepper.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])

gusto/core/labels.py

@@ -113,3 +113,5 @@ def __call__(self, target, value=None):
 hydrostatic = DynamicsLabel("hydrostatic")
 incompressible = DynamicsLabel("incompressible")
 sponge = DynamicsLabel("sponge")
+horizontal_prognostic = Label("horizontal_prognostic")
+vertical_prognostic = Label("vertical_prognostic")

gusto/equations/compressible_euler_equations.py

@@ -4,10 +4,11 @@
     sin, pi, inner, dx, div, cross, FunctionSpace, FacetNormal, jump, avg, dS_v,
     conditional, SpatialCoordinate, split, Constant, as_vector
 )
-from firedrake.fml import subject, replace_subject
+from firedrake.fml import subject, drop, keep
 from gusto.core.labels import (
     time_derivative, transport, prognostic, hydrostatic, linearisation,
-    pressure_gradient, coriolis, gravity, sponge
+    pressure_gradient, coriolis, gravity, sponge, implicit,
+    horizontal_prognostic
 )
 from gusto.equations.thermodynamics import exner_pressure
 from gusto.equations.common_forms import (
@@ -336,49 +337,62 @@ def __init__(self, domain, parameters, sponge_options=None,
                          no_normal_flow_bc_ids=no_normal_flow_bc_ids,
                          active_tracers=active_tracers)
 
-        # Replace
-        self.residual = self.residual.label_map(
-            lambda t: t.has_label(time_derivative),
-            map_if_true=lambda t: self.hydrostatic_projection(t, 'u')
-        )
-
-        # Add an extra hydrostatic term
         u_idx = self.field_names.index('u')
         u = split(self.X)[u_idx]
+        w = self.tests[u_idx]
         k = self.domain.k
-        self.residual += hydrostatic(
-            subject(
-                prognostic(
-                    -inner(k, self.tests[u_idx]) * inner(k, u) * dx, "u"),
-                self.X
-            )
+        u_vert = k*inner(k, u)
+        u_hori = u - u_vert
+
+        # -------------------------------------------------------------------- #
+        # Add hydrostatic term
+        # -------------------------------------------------------------------- #
+        # Term to appear in both explicit and implicit forcings
+        # For the explicit forcing, this will cancel out the u^n part of the
+        # time derivative
+        self.residual += hydrostatic(subject(prognostic(
+            inner(k, w)*inner(k, u)/domain.dt*dx, 'u'), self.X)
         )
 
-    def hydrostatic_projection(self, term, field_name):
-        """
-        Performs the 'hydrostatic' projection.
+        # Term that appears only in implicit forcing
+        # For the implicit forcing, in combination with the term above, the
+        # u^{n+1} term will be cancelled out
+        self.residual += implicit(hydrostatic(subject(prognostic(
+            Constant(-1.0)*inner(k, w)*inner(k, u)/domain.dt*dx, 'u'), self.X))
+        )
 
-        Takes a term involving a vector prognostic variable and replaces the
-        prognostic with only its horizontal components. It also adds the
-        'hydrostatic' label to that term.
+        # # Add Euler-Poincare term
+        # self.residual += hydrostatic(subject(prognostic(
+        #     Constant(0.5)*div(w)*inner(u_hori, u)*dx, 'u'), self.X)
+        # )
 
-        Args:
-            term (:class:`Term`): the term to perform the projection upon.
-            field_name (str): the prognostic field to make hydrostatic.
+        # -------------------------------------------------------------------- #
+        # Only transport horizontal wind
+        # -------------------------------------------------------------------- #
+        # Drop wind transport term, it needs replacing with a version that
+        # includes only the horizontal components for the transported field
+        self.residual = self.residual.label_map(
+            lambda t: t.has_label(transport) and t.get(prognostic) == 'u',
+            map_if_true=drop,
+            map_if_false=keep
+        )
 
-        Returns:
-            :class:`LabelledForm`: the labelled form containing the new term.
-        """
+        # u_term = prognostic(
+        #     advection_equation_circulation_form(domain, w, u_hori, u), 'u'
+        # )
 
-        f_idx = self.field_names.index(field_name)
-        k = self.domain.k
-        X = term.get(subject)
-        field = split(X)[f_idx]
-
-        new_subj = field - inner(field, k) * k
-        # In one step:
-        # - set up the replace_subject routine (which returns a function)
-        # - call that function on the supplied `term` argument,
-        #   to replace the subject with the new hydrostatic subject
-        # - add the hydrostatic label
-        return replace_subject(new_subj, old_idx=f_idx)(term)
+        # Velocity transport term -- depends on formulation
+        if u_transport_option == "vector_invariant_form":
+            u_term = prognostic(vector_invariant_form(domain, w, u_hori, u), 'u')
+        elif u_transport_option == "vector_advection_form":
+            u_term = prognostic(advection_form(w, u_hori, u), 'u')
+        elif u_transport_option == "circulation_form":
+            circ_form = prognostic(
+                advection_equation_circulation_form(domain, w, u_hori, u), 'u'
+            )
+            ke_form = prognostic(kinetic_energy_form(w, u_hori, u), 'u')
+            u_term = ke_form + circ_form
+        else:
+            raise ValueError("Invalid u_transport_option: %s" % u_transport_option)
+
+        self.residual += horizontal_prognostic(subject(u_term, self.X))