Merge branch 'JDBetteridge/parallel_tests' of github.com:firedrakepro…

…ject/gusto into JDBetteridge/parallel_tests

docs/Gemfile.lock

@@ -222,8 +222,8 @@ GEM
     rb-fsevent (0.11.2)
     rb-inotify (0.10.1)
       ffi (~> 1.0)
-    rexml (3.2.8)
-      strscan (>= 3.0.9)
+    rexml (3.3.2)
+      strscan
     rouge (3.26.0)
     ruby2_keywords (0.0.5)
     rubyzip (2.3.2)

gusto/core/labels.py

@@ -95,6 +95,7 @@ def __call__(self, target, value=None):
 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)
 linearisation = Label("linearisation", validator=lambda value: type(value) in [LabelledForm, Term])
+mass_weighted = Label("mass_weighted", validator=lambda value: type(value) in [LabelledForm, Term])
 ibp_label = Label("ibp", validator=lambda value: type(value) == IntegrateByParts)
 
 # labels for terms in the equations

gusto/diagnostics/compressible_euler_diagnostics.py

@@ -5,17 +5,20 @@
                        LinearVariationalSolver, FacetNormal, ds_b, dS_v, div,
                        avg, jump, SpatialCoordinate)
 
-from gusto.diagnostics.diagnostics import DiagnosticField, Energy
+from gusto.diagnostics.diagnostics import (
+    DiagnosticField, Energy, IterativeDiagnosticField
+)
 from gusto.equations import CompressibleEulerEquations
 import gusto.equations.thermodynamics as tde
 from gusto.recovery import Recoverer, BoundaryMethod
 from gusto.equations import TracerVariableType, Phases
 
 __all__ = ["RichardsonNumber", "Entropy", "PhysicalEntropy", "DynamicEntropy",
            "CompressibleKineticEnergy", "Exner", "Theta_e", "InternalEnergy",
-           "PotentialEnergy", "ThermodynamicKineticEnergy", "Dewpoint",
+           "PotentialEnergy", "ThermodynamicKineticEnergy",
            "Temperature", "Theta_d", "RelativeHumidity", "Pressure", "Exner_Vt",
-           "HydrostaticImbalance", "Precipitation", "BruntVaisalaFrequencySquared"]
+           "HydrostaticImbalance", "Precipitation", "BruntVaisalaFrequencySquared",
+           "WetBulbTemperature", "DewpointTemperature"]
 
 
 class RichardsonNumber(DiagnosticField):
@@ -416,21 +419,168 @@ def setup(self, domain, state_fields):
         super().setup(domain, state_fields, space=domain.spaces("DG"))
 
 
-class Dewpoint(ThermodynamicDiagnostic):
-    """The dewpoint temperature diagnostic field."""
+class DewpointTemperature(IterativeDiagnosticField):
+    """
+    The dewpoint temperature diagnostic field. The temperature to which air
+    must be cooled in order to become saturated.
+
+    Note: this will not give sensible answers in the absence of water vapour.
+    """
     name = "Dewpoint"
 
-    def setup(self, domain, state_fields):
+    def __init__(self, equations, space=None, method='interpolate',
+                 num_iterations=3, gamma=1.0):
         """
-        Sets up the :class:`Function` for the diagnostic field.
+        Args:
+            equations (:class:`PrognosticEquationSet`): the equation set being
+                solved by the model.
+            space (:class:`FunctionSpace`, optional): the function space to
+                evaluate the diagnostic field in. Defaults to None, in which
+                case a default space will be chosen for this diagnostic.
+            method (str, optional): a string specifying the method of evaluation
+                for this diagnostic. Valid options are 'interpolate', 'project',
+                'assign' and 'solve'. Defaults to 'interpolate'.
+            num_iterations (integer, optional): number of times to iteratively
+                evaluate the expression. Defaults to 3.
+            gamma (float, optional): weight given to previous guess, which is
+                used to avoid numerical instabilities. Defaults to 1.0.
+        """
+
+        self.parameters = equations.parameters
+        super().__init__(space=space, method=method,
+                         num_iterations=num_iterations, gamma=gamma)
+
+    def implicit_expr(self, domain, state_fields):
+        """
+        The implicit UFL expression for the diagnostic, which should depend
+        on self.field
 
         Args:
             domain (:class:`Domain`): the model's domain object.
             state_fields (:class:`StateFields`): the model's field container.
         """
-        self._setup_thermodynamics(domain, state_fields)
-        self.expr = tde.T_dew(self.parameters, self.p, self.r_v)
-        super().setup(domain, state_fields, space=self.Vtheta)
+
+        theta = state_fields('theta')
+        rho = state_fields('rho')
+        if 'water_vapour' in state_fields._field_names:
+            r_v = state_fields('water_vapour')
+        else:
+            raise RuntimeError('Dewpoint temperature diagnostic should only'
+                               + 'be used with water vapour')
+
+        exner = tde.exner_pressure(self.parameters, rho, theta)
+        pressure = tde.p(self.parameters, exner)
+        temperature = tde.T(self.parameters, theta, exner, r_v=r_v)
+        r_sat = tde.r_sat(self.parameters, self.field, pressure)
+
+        return self.field - temperature*(r_sat - r_v)
+
+    def set_first_guess(self, domain, state_fields):
+        """
+        The first guess of the diagnostic, set to be the dry temperature.
+
+        Args:
+            domain (:class:`Domain`): the model's domain object.
+            state_fields (:class:`StateFields`): the model's field container.
+        """
+        theta = state_fields('theta')
+        rho = state_fields('rho')
+        if 'water_vapour' in state_fields._field_names:
+            r_v = state_fields('water_vapour')
+        else:
+            raise RuntimeError('Dewpoint temperature diagnostic should only'
+                               + 'be used with water vapour')
+
+        exner = tde.exner_pressure(self.parameters, rho, theta)
+        temperature = tde.T(self.parameters, theta, exner, r_v=r_v)
+
+        return temperature
+
+
+class WetBulbTemperature(IterativeDiagnosticField):
+    """
+    The wet-bulb temperature diagnostic field. The temperature of air cooled to
+    saturation by the evaporation of water.
+    """
+    name = "WetBulb"
+
+    def __init__(self, equations, space=None, method='interpolate',
+                 num_iterations=3, gamma=0.5):
+        """
+        Args:
+            equations (:class:`PrognosticEquationSet`): the equation set being
+                solved by the model.
+            space (:class:`FunctionSpace`, optional): the function space to
+                evaluate the diagnostic field in. Defaults to None, in which
+                case a default space will be chosen for this diagnostic.
+            method (str, optional): a string specifying the method of evaluation
+                for this diagnostic. Valid options are 'interpolate', 'project',
+                'assign' and 'solve'. Defaults to 'interpolate'.
+            num_iterations (integer, optional): number of times to iteratively
+                evaluate the expression. Defaults to 3.
+            gamma (float, optional): weight given to previous guess, which is
+                used to avoid numerical instabilities. Defaults to 0.8.
+        """
+
+        self.parameters = equations.parameters
+        super().__init__(space=space, method=method,
+                         num_iterations=num_iterations, gamma=gamma)
+
+    def implicit_expr(self, domain, state_fields):
+        """
+        The implicit UFL expression for the diagnostic, which should depend
+        on self.field
+
+        Args:
+            domain (:class:`Domain`): the model's domain object.
+            state_fields (:class:`StateFields`): the model's field container.
+        """
+
+        theta = state_fields('theta')
+        rho = state_fields('rho')
+        if 'water_vapour' in state_fields._field_names:
+            r_v = state_fields('water_vapour')
+        else:
+            r_v = 0.0*theta  # zero expression
+
+        exner = tde.exner_pressure(self.parameters, rho, theta)
+        pressure = tde.p(self.parameters, exner)
+        temperature = tde.T(self.parameters, theta, exner, r_v=r_v)
+        r_sat = tde.r_sat(self.parameters, self.field, pressure)
+
+        # In the comments, preserve a simpler expression:
+        # L_v0 = self.parameters.L_v0
+        # R_v = self.parameters.R_v
+        # c_v = self.parameters.cv
+        # return L_v0 / R_v + (R_v*temperature - L_v0)/R_v * exp(R_v/c_v*(r_sat - r_v))
+
+        # Reduce verbosity by introducing intermediate variables
+        b = -self.parameters.L_v0 - (self.parameters.c_pl - self.parameters.c_pv)*self.parameters.T_0
+        a = self.parameters.R_v + self.parameters.c_pl - self.parameters.c_pv
+        A = self.parameters.c_vv
+        B = self.parameters.cv
+
+        return - b / a + (a*temperature + b) / a * ((A*r_sat + B) / (A*r_v + B))**(a/A)
+
+    def set_first_guess(self, domain, state_fields):
+        """
+        The first guess of the diagnostic, set to be the dry temperature.
+
+        Args:
+            domain (:class:`Domain`): the model's domain object.
+            state_fields (:class:`StateFields`): the model's field container.
+        """
+        theta = state_fields('theta')
+        rho = state_fields('rho')
+        if 'water_vapour' in state_fields._field_names:
+            r_v = state_fields('water_vapour')
+        else:
+            r_v = 0.0*theta  # zero expression
+
+        exner = tde.exner_pressure(self.parameters, rho, theta)
+        temperature = tde.T(self.parameters, theta, exner, r_v=r_v)
+
+        return temperature
 
 
 class Temperature(ThermodynamicDiagnostic):

gusto/diagnostics/diagnostics.py

@@ -21,7 +21,7 @@
            "XComponent", "YComponent", "ZComponent", "MeridionalComponent",
            "ZonalComponent", "RadialComponent", "Energy", "KineticEnergy",
            "Sum", "Difference", "SteadyStateError", "Perturbation",
-           "Divergence", "TracerDensity"]
+           "Divergence", "TracerDensity", "IterativeDiagnosticField"]
 
 
 class Diagnostics(object):
@@ -220,6 +220,123 @@ def __call__(self):
         return self.field
 
 
+class IterativeDiagnosticField(DiagnosticField):
+    """
+    Iterative evaluation of a diagnostic expression for diagnostics with
+    no explicit definition but an implicit definition. This uses a form of
+    damped Picard iteration.
+    """
+    def __init__(self, space=None, method='interpolate', required_fields=(),
+                 num_iterations=3, gamma=0.8):
+        """
+        Args:
+            space (:class:`FunctionSpace`, optional): the function space to
+                evaluate the diagnostic field in. Defaults to None, in which
+                case a default space will be chosen for this diagnostic.
+            method (str, optional): a string specifying the method of evaluation
+                for this diagnostic. Valid options are 'interpolate', 'project',
+                'assign' and 'solve'. Defaults to 'interpolate'.
+            required_fields (tuple, optional): tuple of names of the fields that
+                are required for the computation of this diagnostic field.
+                Defaults to ().
+            num_iterations (integer, optional): number of times to iteratively
+                evaluate the expression. Defaults to 3.
+            gamma (float, optional): weight given to previous guess, which is
+                used to avoid numerical instabilities.
+        """
+
+        super().__init__(space=space, method=method, required_fields=required_fields)
+        self.num_iterations = num_iterations
+        self.gamma = gamma
+
+    def setup(self, domain, state_fields, space=None):
+        """
+        Sets up the :class:`Function` for the diagnostic field.
+
+        Args:
+            domain (:class:`Domain`): the model's domain object.
+            state_fields (:class:`StateFields`): the model's field container.
+            space (:class:`FunctionSpace`, optional): the function space for the
+                diagnostic field to be computed in. Defaults to None, in which
+                case the space will be DG0.
+        """
+
+        # NB: child classes of this need to set up:
+        # (a) self.implicit_expr: the implicit UFL expression for the diagnostic
+        #                         which must depend on self.field
+        # (b) self.first_guess: the UFL expression for the first guess
+
+        if not self._initialised:
+            if self.space is None:
+                if space is None:
+                    if not hasattr(domain.spaces, "DG0"):
+                        space = domain.spaces.create_space("DG0", "DG", 0)
+                    else:
+                        space = domain.spaces("DG0")
+                self.space = space
+            else:
+                space = self.space
+
+            # Add space to domain
+            assert space.name is not None, \
+                f'Diagnostics {self.name} is using a function space which does not have a name'
+            if not hasattr(domain.spaces, space.name):
+                domain.spaces.add_space(space.name, space)
+
+            self.field = state_fields(self.name, space=space, dump=self.to_dump, pick_up=False)
+            self.expr = (1.0 - self.gamma)*self.field + self.gamma*self.implicit_expr(domain, state_fields)
+            self.first_guess = self.set_first_guess(domain, state_fields)
+
+            if self.method != 'solve':
+                assert self.expr is not None, \
+                    f"The expression for diagnostic {self.name} has not been specified"
+
+            # Solve method must be declared in diagnostic's own setup routine
+            if self.method == 'interpolate':
+                self.evaluator = Interpolator(self.expr, self.field)
+            elif self.method == 'project':
+                self.evaluator = Projector(self.expr, self.field)
+            elif self.method == 'assign':
+                self.evaluator = Assigner(self.field, self.expr)
+
+            self._initialised = True
+
+    @property
+    @abstractmethod
+    def implicit_expr(self, domain, state_fields):
+        """
+        The implicit UFL expression for the diagnostic, which should depend
+        on self.field
+
+        Args:
+            domain (:class:`Domain`): the model's domain object.
+            state_fields (:class:`StateFields`): the model's field container.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def set_first_guess(self, domain, state_fields):
+        """
+        The first guess of the diagnostic
+
+        Args:
+            domain (:class:`Domain`): the model's domain object.
+            state_fields (:class:`StateFields`): the model's field container.
+        """
+        pass
+
+    def compute(self):
+        """Compute the diagnostic field from the current state."""
+
+        # Set first guess
+        self.field.interpolate(self.first_guess)
+
+        # Iterate
+        for _ in range(self.num_iterations):
+            super().compute()
+
+
 class CourantNumber(DiagnosticField):
     """Dimensionless Courant number diagnostic field."""
     name = "CourantNumber"