Merge remote-tracking branch 'origin/main' into TBendall/HydrostaticT…

…ests

gusto/core/configuration.py

@@ -269,6 +269,19 @@ class BoundaryLayerParameters(Configuration):
     mu = 100.                   # Interior penalty coefficient for vertical diffusion
 
 
+class HeldSuarezParameters(Configuration):
+    """
+    Parameters used in the default configuration for the Held Suarez test case.
+    """
+    T0stra = 200               # Stratosphere temp
+    T0surf = 315               # Surface temperature at equator
+    T0horiz = 60               # Equator to pole temperature difference
+    T0vert = 10                # Stability parameter
+    sigmab = 0.7               # Height of the boundary layer
+    tau_d = 40 * 24 * 60 * 60  # 40 day time scale
+    tau_u = 4 * 24 * 60 * 60   # 4 day timescale
+
+
 class SubcyclingOptions(Configuration):
     """
     Describes the process of subcycling a time discretisation, by dividing the

gusto/diagnostics/diagnostics.py

@@ -1,14 +1,15 @@
 """Common diagnostic fields."""
 
 
-from firedrake import (assemble, dot, dx, Function, sqrt, TestFunction,
+from firedrake import (dot, dx, Function, sqrt, TestFunction,
                        TrialFunction, Constant, grad, inner, FacetNormal,
                        LinearVariationalProblem, LinearVariationalSolver,
                        ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, avg, pi,
                        TensorFunctionSpace, SpatialCoordinate, as_vector,
-                       Projector, Interpolator, FunctionSpace, FiniteElement,
+                       Projector, assemble, FunctionSpace, FiniteElement,
                        TensorProductElement)
 from firedrake.assign import Assigner
+from firedrake.__future__ import interpolate
 from ufl.domain import extract_unique_domain
 
 from abc import ABCMeta, abstractmethod, abstractproperty
@@ -193,7 +194,7 @@ def setup(self, domain, state_fields, space=None):
 
             # Solve method must be declared in diagnostic's own setup routine
             if self.method == 'interpolate':
-                self.evaluator = Interpolator(self.expr, self.field)
+                self.evaluator = interpolate(self.expr, self.space)
             elif self.method == 'project':
                 self.evaluator = Projector(self.expr, self.field)
             elif self.method == 'assign':
@@ -207,7 +208,7 @@ def compute(self):
         logger.debug(f'Computing diagnostic {self.name} with {self.method} method')
 
         if self.method == 'interpolate':
-            self.evaluator.interpolate()
+            self.field.assign(assemble(self.evaluator))
         elif self.method == 'assign':
             self.evaluator.assign()
         elif self.method == 'project':
@@ -294,7 +295,7 @@ def setup(self, domain, state_fields, space=None):
 
             # Solve method must be declared in diagnostic's own setup routine
             if self.method == 'interpolate':
-                self.evaluator = Interpolator(self.expr, self.field)
+                self.evaluator = interpolate(self.expr, self.space)
             elif self.method == 'project':
                 self.evaluator = Projector(self.expr, self.field)
             elif self.method == 'assign':

gusto/diagnostics/shallow_water_diagnostics.py

@@ -2,9 +2,10 @@
 
 
 from firedrake import (
-    dx, TestFunction, TrialFunction, grad, inner, curl, Function, Interpolator,
+    dx, TestFunction, TrialFunction, grad, inner, curl, Function, assemble,
     LinearVariationalProblem, LinearVariationalSolver, conditional
 )
+from firedrake.__future__ import interpolate
 from gusto.diagnostics.diagnostics import DiagnosticField, Energy
 
 __all__ = ["ShallowWaterKineticEnergy", "ShallowWaterPotentialEnergy",
@@ -321,14 +322,14 @@ def setup(self, domain, state_fields):
 
         qsat_expr = self.equation.compute_saturation(state_fields.X(
             self.equation.field_name))
-        self.qsat_interpolator = Interpolator(qsat_expr, self.qsat_func)
+        self.qsat_interpolate = interpolate(qsat_expr, space)
         self.expr = conditional(q_t < self.qsat_func, q_t, self.qsat_func)
 
         super().setup(domain, state_fields, space=space)
 
     def compute(self):
         """Performs the computation of the diagnostic field."""
-        self.qsat_interpolator.interpolate()
+        self.qsat_func.assign(assemble(self.qsat_interpolate))
         super().compute()
 
 
@@ -371,13 +372,13 @@ def setup(self, domain, state_fields):
 
         qsat_expr = self.equation.compute_saturation(state_fields.X(
             self.equation.field_name))
-        self.qsat_interpolator = Interpolator(qsat_expr, self.qsat_func)
+        self.qsat_interpolate = interpolate(qsat_expr, space)
         vapour = conditional(q_t < self.qsat_func, q_t, self.qsat_func)
         self.expr = q_t - vapour
 
         super().setup(domain, state_fields, space=space)
 
     def compute(self):
         """Performs the computation of the diagnostic field."""
-        self.qsat_interpolator.interpolate()
+        self.qsat_func.assign(assemble(self.qsat_interpolate))
         super().compute()

gusto/physics/__init__.py

@@ -2,4 +2,5 @@
 from gusto.physics.chemistry import *                       # noqa
 from gusto.physics.boundary_and_turbulence import *         # noqa
 from gusto.physics.microphysics import *                    # noqa
-from gusto.physics.shallow_water_microphysics import *      # noqa
+from gusto.physics.shallow_water_microphysics import *      # noqa
+from gusto.physics.held_suarez_forcing import *             # noqa

gusto/physics/held_suarez_forcing.py

@@ -0,0 +1,188 @@
+import numpy as np
+from firedrake import (Interpolator, Function, dx, pi, SpatialCoordinate,
+                       split, conditional, ge, sin, dot, ln, cos, inner, Projector)
+from firedrake.fml import subject
+from gusto.core.coord_transforms import lonlatr_from_xyz
+from gusto.recovery import Recoverer, BoundaryMethod
+from gusto.physics.physics_parametrisation import PhysicsParametrisation
+from gusto.core.labels import prognostic
+from gusto.equations import thermodynamics
+from gusto.core.configuration import HeldSuarezParameters
+from gusto.core import logger
+
+
+class Relaxation(PhysicsParametrisation):
+    """
+    Relaxation term for Held Suarez
+    """
+
+    def __init__(self, equation, variable_name, parameters, hs_parameters=None):
+        """
+        Args:
+            equation (:class:`PrognosticEquationSet`): the model's equation.
+            variable_name (str): the name of the variable
+            hs_parameters (:class'Configuration'): contains the parameters for the Held-suariez test case
+
+        """
+        label_name = f'relaxation_{variable_name}'
+        if hs_parameters is None:
+            hs_parameters = HeldSuarezParameters()
+            logger.warning('Using default Held-Suarez parameters')
+        super().__init__(equation, label_name, hs_parameters)
+
+        if equation.domain.on_sphere:
+            x, y, z = SpatialCoordinate(equation.domain.mesh)
+            _, lat, _ = lonlatr_from_xyz(x, y, z)
+        else:
+            # TODO: this could be determined some other way
+            # Take a mid-latitude
+            lat = pi / 4
+
+        self.X = Function(equation.X.function_space())
+        X = self.X
+        self.domain = equation.domain
+        theta_idx = equation.field_names.index('theta')
+        self.theta = X.subfunctions[theta_idx]
+        Vt = equation.domain.spaces('theta')
+        rho_idx = equation.field_names.index('rho')
+        rho = split(X)[rho_idx]
+
+        boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None
+        self.rho_averaged = Function(Vt)
+        self.rho_recoverer = Recoverer(rho, self.rho_averaged, boundary_method=boundary_method)
+        self.exner = Function(Vt)
+        self.exner_interpolator = Interpolator(
+            thermodynamics.exner_pressure(equation.parameters,
+                                          self.rho_averaged, self.theta), self.exner)
+        self.sigma = Function(Vt)
+        kappa = equation.parameters.kappa
+
+        T0surf = hs_parameters.T0surf
+        T0horiz = hs_parameters.T0horiz
+        T0vert = hs_parameters.T0vert
+        T0stra = hs_parameters.T0stra
+
+        sigma_b = hs_parameters.sigmab
+        tau_d = hs_parameters.tau_d
+        tau_u = hs_parameters.tau_u
+
+        theta_condition = (T0surf - T0horiz * sin(lat)**2 - (T0vert * ln(self.exner) * cos(lat)**2)/kappa)
+        Theta_eq = conditional(T0stra/self.exner >= theta_condition, T0stra/self.exner, theta_condition)
+
+        # timescale of temperature forcing
+        tau_cond = (self.sigma**(1/kappa) - sigma_b) / (1 - sigma_b)
+        newton_freq = 1 / tau_d + (1/tau_u - 1/tau_d) * conditional(0 >= tau_cond, 0, tau_cond) * cos(lat)**4
+        forcing_expr = newton_freq * (self.theta - Theta_eq)
+
+        # Create source for forcing
+        self.source_relaxation = Function(Vt)
+        self.source_interpolator = Interpolator(forcing_expr, self.source_relaxation)
+
+        # Add relaxation term to residual
+        test = equation.tests[theta_idx]
+        dx_reduced = dx(degree=equation.domain.max_quad_degree)
+        forcing_form = test * self.source_relaxation * dx_reduced
+        equation.residual += self.label(subject(prognostic(forcing_form, 'theta'), X), self.evaluate)
+
+    def evaluate(self, x_in, dt):
+        """
+        Evalutes the source term generated by the physics.
+
+        Args:
+            x_in: (:class:`Function`): the (mixed) field to be evolved.
+            dt: (:class:`Constant`): the timestep, which can be the time
+                interval for the scheme.
+        """
+        self.X.assign(x_in)
+        self.rho_recoverer.project()
+        self.exner_interpolator.interpolate()
+
+        # Determine sigma:= exner / exner_surf
+        exner_columnwise, index_data = self.domain.coords.get_column_data(self.exner, self.domain)
+        sigma_columnwise = np.zeros_like(exner_columnwise)
+        for col in range(len(exner_columnwise[:, 0])):
+            sigma_columnwise[col, :] = exner_columnwise[col, :] / exner_columnwise[col, 0]
+        self.domain.coords.set_field_from_column_data(self.sigma, sigma_columnwise, index_data)
+
+        self.source_interpolator.interpolate()
+
+
+class RayleighFriction(PhysicsParametrisation):
+    """
+    Forcing term on the velocity of the form
+    F_u = -u / a,
+    where a is some friction factor
+    """
+    def __init__(self, equation, hs_parameters=None):
+        """
+         Args:
+            equation (:class:`PrognosticEquationSet`): the model's equation.
+            hs_parameters (:class'Configuration'): contains the parameters for the Held-suariez test case
+        """
+        label_name = 'rayleigh_friction'
+        if hs_parameters is None:
+            hs_parameters = HeldSuarezParameters()
+            logger.warning('Using default Held-Suarez parameters')
+        super().__init__(equation, label_name, hs_parameters)
+
+        self.domain = equation.domain
+        self.X = Function(equation.X.function_space())
+        X = self.X
+        k = equation.domain.k
+        u_idx = equation.field_names.index('u')
+        u = split(X)[u_idx]
+        theta_idx = equation.field_names.index('theta')
+        self.theta = X.subfunctions[theta_idx]
+        rho_idx = equation.field_names.index('rho')
+        rho = split(X)[rho_idx]
+        Vt = equation.domain.spaces('theta')
+        Vu = equation.domain.spaces('HDiv')
+        u_hori = u - k*dot(u, k)
+
+        boundary_method = BoundaryMethod.extruded if self.domain == 0 else None
+        self.rho_averaged = Function(Vt)
+        self.exner = Function(Vt)
+        self.rho_recoverer = Recoverer(rho, self.rho_averaged, boundary_method=boundary_method)
+        self.exner_interpolator = Interpolator(
+            thermodynamics.exner_pressure(equation.parameters,
+                                          self.rho_averaged, self.theta), self.exner)
+
+        self.sigma = Function(Vt)
+        sigmab = hs_parameters.sigmab
+        kappa = equation.parameters.kappa
+        tau_fric = 24 * 60 * 60
+
+        tau_cond = (self.sigma**(1/kappa) - sigmab) / (1 - sigmab)
+        wind_timescale = conditional(ge(0, tau_cond), 0, tau_cond) / tau_fric
+        forcing_expr = u_hori * wind_timescale
+
+        self.source_friction = Function(Vu)
+        self.source_projector = Projector(forcing_expr, self.source_friction)
+
+        tests = equation.tests
+        test = tests[u_idx]
+        dx_reduced = dx(degree=equation.domain.max_quad_degree)
+        source_form = inner(test, self.source_friction) * dx_reduced
+        equation.residual += self.label(subject(prognostic(source_form, 'u'), X), self.evaluate)
+
+    def evaluate(self, x_in, dt):
+        """
+        Evaluates the source term generated by the physics. This does nothing if
+        the implicit formulation is not used.
+
+        Args:
+            x_in: (:class: 'Function'): the (mixed) field to be evolved.
+            dt: (:class: 'Constant'): the timestep, which can be the time
+                interval for the scheme.
+        """
+        self.X.assign(x_in)
+        self.rho_recoverer.project()
+        self.exner_interpolator.interpolate()
+        # Determine sigma:= exner / exner_surf
+        exner_columnwise, index_data = self.domain.coords.get_column_data(self.exner, self.domain)
+        sigma_columnwise = np.zeros_like(exner_columnwise)
+        for col in range(len(exner_columnwise[:, 0])):
+            sigma_columnwise[col, :] = exner_columnwise[col, :] / exner_columnwise[col, 0]
+        self.domain.coords.set_field_from_column_data(self.sigma, sigma_columnwise, index_data)
+
+        self.source_projector.project()

gusto/physics/microphysics.py

@@ -4,9 +4,10 @@
 """
 
 from firedrake import (
-    Interpolator, conditional, Function, dx, min_value, max_value, Constant, pi,
-    Projector
+    conditional, Function, dx, min_value, max_value, Constant, pi,
+    Projector, assemble
 )
+from firedrake.__future__ import interpolate
 from firedrake.fml import identity, Term, subject
 from gusto.equations import Phases, TracerVariableType
 from gusto.recovery import Recoverer, BoundaryMethod
@@ -170,8 +171,7 @@ def __init__(self, equation, vapour_name='water_vapour',
         # Add terms to equations and make interpolators
         # -------------------------------------------------------------------- #
         self.source = [Function(V) for factor in factors]
-        self.source_interpolators = [Interpolator(sat_adj_expr*factor, source)
-                                     for factor, source in zip(factors, self.source)]
+        self.source_interpolate = [interpolate(sat_adj_expr*factor, V) for factor in factors]
 
         tests = [equation.tests[idx] for idx in V_idxs]
 
@@ -195,8 +195,8 @@ def evaluate(self, x_in, dt):
         if isinstance(self.equation, CompressibleEulerEquations):
             self.rho_recoverer.project()
         # Evaluate the source
-        for interpolator in self.source_interpolators:
-            interpolator.interpolate()
+        for interpolator, src in zip(self.source_interpolate, self.source):
+            src.assign(assemble(interpolator))
 
 
 class AdvectedMoments(Enum):
@@ -440,7 +440,7 @@ def __init__(self, equation, cloud_name='cloud_water', rain_name='rain',
                                                         min_value(accu_rate, self.cloud_water / self.dt),
                                                         min_value(accr_rate + accu_rate, self.cloud_water / self.dt))))
 
-        self.source_interpolator = Interpolator(rain_expr, self.source)
+        self.source_interpolate = interpolate(rain_expr, Vt)
 
         # Add term to equation's residual
         test_cl = equation.tests[self.cloud_idx]
@@ -464,7 +464,7 @@ def evaluate(self, x_in, dt):
         self.rain.assign(x_in.subfunctions[self.rain_idx])
         self.cloud_water.assign(x_in.subfunctions[self.cloud_idx])
         # Evaluate the source
-        self.source.assign(self.source_interpolator.interpolate())
+        self.source.assign(assemble(self.source_interpolate))
 
 
 class EvaporationOfRain(PhysicsParametrisation):
@@ -609,8 +609,7 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour',
         # Add terms to equations and make interpolators
         # -------------------------------------------------------------------- #
         self.source = [Function(V) for factor in factors]
-        self.source_interpolators = [Interpolator(evap_rate*factor, source)
-                                     for factor, source in zip(factors, self.source)]
+        self.source_interpolate = [interpolate(evap_rate*factor, V) for factor in factors]
 
         tests = [equation.tests[idx] for idx in V_idxs]
 
@@ -634,5 +633,5 @@ def evaluate(self, x_in, dt):
         if isinstance(self.equation, CompressibleEulerEquations):
             self.rho_recoverer.project()
         # Evaluate the source
-        for interpolator in self.source_interpolators:
-            interpolator.interpolate()
+        for interpolator, src in zip(self.source_interpolate, self.source):
+            src.assign(assemble(interpolator))