Merge pull request #425 from firedrakeproject/physics_package

A package of idealised physics routines

examples/compressible/skamarock_klemp_hydrostatic.py

@@ -17,7 +17,7 @@
 dt = 25.
 if '--running-tests' in sys.argv:
     nlayers = 5  # horizontal layers
-    columns = 50  # number of columns
+    columns = 10  # number of columns
     tmax = dt
     dumpfreq = 1
 else:

gusto/common_forms.py

@@ -183,14 +183,14 @@ def advection_equation_circulation_form(domain, test, q, ubar):
 
 def diffusion_form(test, q, kappa):
     u"""
-    The diffusion form, ∇.(κ∇q) for diffusivity κ and variable q.
+    The diffusion form, -∇.(κ∇q) for diffusivity κ and variable q.
 
     Args:
         test (:class:`TestFunction`): the test function.
         q (:class:`ufl.Expr`): the variable to be diffused.
         kappa: (:class:`ufl.Expr`): the diffusivity value.
     """
 
-    form = inner(test, div(kappa*grad(q)))*dx
+    form = -inner(test, div(kappa*grad(q)))*dx
 
     return diffusion(form)

gusto/configuration.py

@@ -8,7 +8,7 @@
     "IntegrateByParts", "TransportEquationType", "OutputParameters",
     "CompressibleParameters", "ShallowWaterParameters",
     "EmbeddedDGOptions", "RecoveryOptions", "SUPGOptions",
-    "SpongeLayerParameters", "DiffusionParameters"
+    "SpongeLayerParameters", "DiffusionParameters", "BoundaryLayerParameters"
 ]
 
 
@@ -185,3 +185,18 @@ class DiffusionParameters(Configuration):
 
     kappa = None
     mu = None
+
+
+class BoundaryLayerParameters(Configuration):
+    """
+    Parameters for the idealised wind drag, surface flux and boundary layer
+    mixing schemes.
+    """
+
+    coeff_drag_0 = 7e-4         # Zeroth drag coefficient (dimensionless)
+    coeff_drag_1 = 6.5e-5       # First drag coefficient (s/m)
+    coeff_drag_2 = 2e-3         # Second drag coefficient (dimensionless)
+    coeff_heat = 1.1e-3         # Dimensionless surface sensible heat coefficient
+    coeff_evap = 1.1e-3         # Dimensionless surface evaporation coefficient
+    height_surface_layer = 75.  # Height (m) of surface level (usually lowest level)
+    mu = 100.                   # Interior penalty coefficient for vertical diffusion

gusto/coordinates.py

@@ -7,6 +7,7 @@
 from gusto.logging import logger
 from firedrake import SpatialCoordinate, Function
 import numpy as np
+import pandas as pd
 
 
 class Coordinates(object):
@@ -154,3 +155,124 @@ def register_space(self, domain, space_name):
                     new_coords = comm.recv(source=procid, tag=my_tag)
                     (low_lim, up_lim) = self.parallel_array_lims[space_name][procid][:]
                     self.global_chi_coords[space_name][i, low_lim:up_lim+1] = new_coords
+
+    def get_column_data(self, field, domain):
+        """
+        Reshapes a field's data into columns.
+
+        Args:
+            field (:class:`Function`): the field whose data needs sorting.
+            domain (:class:`Domain`): the domain used to register coordinates
+                if this hasn't already been done.
+
+        Returns:
+            tuple of :class:`numpy.ndarray`: a 2D array of data, arranged in
+                columns, and the data pairing the indices of the data with the
+                ordered column data.
+        """
+
+        space_name = field.function_space().name
+        if space_name not in self.chi_coords.keys():
+            self.register_space(domain, space_name)
+        coords = self.chi_coords[space_name]
+
+        data_is_3d = (len(coords) == 3)
+        coords_X = coords[0].dat.data
+        coords_Y = coords[1].dat.data if data_is_3d else None
+        coords_Z = coords[-1].dat.data
+
+        # ------------------------------------------------------------------------ #
+        # Round data to ensure sorting in dataframe is OK
+        # ------------------------------------------------------------------------ #
+
+        # Work out digits to round to, based on number of points and range of coords
+        num_points = np.size(coords_X)
+        data_range = np.max(coords_X) - np.min(coords_X)
+        if data_range > np.finfo(type(data_range)).tiny:
+            digits = int(np.floor(-np.log10(data_range / num_points)) + 3)
+            coords_X = coords_X.round(digits)
+
+        if data_is_3d:
+            data_range = np.max(coords_Y) - np.min(coords_Y)
+            if data_range > np.finfo(type(data_range)).tiny:
+                # Only round if there is already some range
+                digits = int(np.floor(-np.log10(data_range / num_points)) + 3)
+                coords_Y = coords_Y.round(digits)
+
+        # -------------------------------------------------------------------- #
+        # Make data frame
+        # -------------------------------------------------------------------- #
+
+        data_dict = {'field': field.dat.data, 'X': coords_X, 'Z': coords_Z,
+                     'index': range(len(field.dat.data))}
+        if coords_Y is not None:
+            data_dict['Y'] = coords_Y
+
+        # Put everything into a pandas dataframe
+        data = pd.DataFrame(data_dict)
+
+        # Sort array by X and Y coordinates
+        if data_is_3d:
+            data = data.sort_values(by=['X', 'Y', 'Z'])
+            first_X, first_Y = data['X'].values[0], data['Y'].values[0]
+            first_point = data[(np.isclose(data['X'], first_X))
+                               & (np.isclose(data['Y'], first_Y))]
+
+        else:
+            data = data.sort_values(by=['X', 'Z'])
+            first_X = data['X'].values[0]
+            first_point = data[np.isclose(data['X'], first_X)]
+
+        # Number of levels should correspond to the number of points with the first
+        # coordinate values
+        num_levels = len(first_point)
+        assert len(data) % num_levels == 0, 'Unable to nicely divide data into levels'
+
+        # -------------------------------------------------------------------- #
+        # Create new arrays to store structured data
+        # -------------------------------------------------------------------- #
+
+        num_hori_points = int(len(data) / num_levels)
+        field_data = np.zeros((num_hori_points, num_levels))
+        coords_X = np.zeros((num_hori_points, num_levels))
+        if data_is_3d:
+            coords_Y = np.zeros((num_hori_points, num_levels))
+        coords_Z = np.zeros((num_hori_points, num_levels))
+        index_data = np.zeros((num_hori_points, num_levels), dtype=int)
+
+        # -------------------------------------------------------------------- #
+        # Fill arrays, on the basis of the dataframe already being sorted
+        # -------------------------------------------------------------------- #
+
+        for lev_idx in range(num_levels):
+            data_slice = slice(lev_idx, num_hori_points*num_levels+lev_idx, num_levels)
+            field_data[:, lev_idx] = data['field'].values[data_slice]
+            coords_X[:, lev_idx] = data['X'].values[data_slice]
+            if data_is_3d:
+                coords_Y[:, lev_idx] = data['Y'].values[data_slice]
+            coords_Z[:, lev_idx] = data['Z'].values[data_slice]
+            index_data[:, lev_idx] = data['index'].values[data_slice]
+
+        return field_data, index_data
+
+    def set_field_from_column_data(self, field, columnwise_data, index_data):
+        """
+        Fills in field data from some columnwise data.
+
+        Args:
+            field (:class:`Function`): the field whose data shall be filled.
+            columnwise_data (:class:`numpy.ndarray`): the field data arranged
+                into columns, to be written into the field.
+            index_data (:class:`numpy.ndarray`): the indices of the original
+                field data, arranged like the columnwise data.
+
+        Returns:
+            :class:`Function`: the updated field.
+        """
+
+        _, num_levels = np.shape(columnwise_data)
+
+        for lev_idx in range(num_levels):
+            field.dat.data[index_data[:, lev_idx]] = columnwise_data[:, lev_idx]
+
+        return field

gusto/diagnostics.py

@@ -14,6 +14,7 @@
 from gusto.recovery import Recoverer, BoundaryMethod
 from gusto.equations import CompressibleEulerEquations
 from gusto.active_tracers import TracerVariableType, Phases
+from gusto.logging import logger
 import numpy as np
 
 __all__ = ["Diagnostics", "CourantNumber", "Gradient", "XComponent", "YComponent",
@@ -25,7 +26,7 @@
            "ThermodynamicKineticEnergy", "Dewpoint", "Temperature", "Theta_d",
            "RelativeHumidity", "Pressure", "Exner_Vt", "HydrostaticImbalance", "Precipitation",
            "PotentialVorticity", "RelativeVorticity", "AbsoluteVorticity", "Divergence",
-           "TracerDensity"]
+           "BruntVaisalaFrequencySquared", "TracerDensity"]
 
 
 class Diagnostics(object):
@@ -221,6 +222,8 @@ def setup(self, domain, state_fields, space=None):
     def compute(self):
         """Compute the diagnostic field from the current state."""
 
+        logger.debug(f'Computing diagnostic {self.name} with {self.method} method')
+
         if self.method == 'interpolate':
             self.evaluator.interpolate()
         elif self.method == 'assign':
@@ -944,6 +947,51 @@ def setup(self, domain, state_fields):
         super().setup(domain, state_fields)
 
 
+class BruntVaisalaFrequencySquared(DiagnosticField):
+    """The diagnostic for the Brunt-Väisälä frequency."""
+    name = "Brunt-Vaisala_squared"
+
+    def __init__(self, equations, space=None, method='interpolate'):
+        """
+        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'.
+        """
+        self.parameters = equations.parameters
+        # Work out required fields
+        if isinstance(equations, CompressibleEulerEquations):
+            required_fields = ['theta']
+            if equations.active_tracers is not None and len(equations.active_tracers) > 1:
+                # TODO: I think theta here should be theta_e, which would be
+                # easiest if this is a ThermodynamicDiagnostic. But in the dry
+                # case, our numerical theta_e does not reduce to the numerical
+                # dry theta
+                raise NotImplementedError(
+                    'Brunt-Vaisala diagnostic not implemented for moist equations')
+        else:
+            raise NotImplementedError(
+                f'Brunt-Vaisala diagnostic not implemented for {type(equations)}')
+        super().__init__(space=space, method=method, required_fields=tuple(required_fields))
+
+    def setup(self, domain, state_fields):
+        """
+        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.
+        """
+        theta = state_fields('theta')
+        self.expr = self.parameters.g/theta * dot(domain.k, grad(theta))
+        super().setup(domain, state_fields)
+
+
 class Sum(DiagnosticField):
     """Base diagnostic for computing the sum of two fields."""
     def __init__(self, field_name1, field_name2):

gusto/domain.py

@@ -132,7 +132,7 @@ def __init__(self, mesh, dt, family, degree=None,
             radius_field = Function(CG1)
             radius_field.interpolate(r_shell)
             # TODO: this should use global min kernel
-            radius = np.min(radius_field.dat.data_ro)
+            radius = Constant(np.min(radius_field.dat.data_ro))
         else:
             radius = None
 
@@ -146,12 +146,47 @@ def __init__(self, mesh, dt, family, degree=None,
         _ = self.spaces('DG1_equispaced')
         self.coords.register_space(self, 'DG1_equispaced')
 
+        # Set height above surface (requires coordinates)
+        if hasattr(mesh, "_base_mesh"):
+            self.set_height_above_surface()
+
         # -------------------------------------------------------------------- #
         # Construct metadata about domain
         # -------------------------------------------------------------------- #
 
         self.metadata = construct_domain_metadata(mesh, self.coords, self.on_sphere)
 
+    def set_height_above_surface(self):
+        """
+        Sets a coordinate field which corresponds to height above the domain's
+        surface.
+        """
+
+        from firedrake import dot
+
+        x = SpatialCoordinate(self.mesh)
+
+        # Make a height field in CG1
+        CG1 = FunctionSpace(self.mesh, "CG", 1, name='CG1')
+        self.spaces.add_space('CG1', CG1)
+        self.coords.register_space(self, 'CG1')
+        CG1_height = Function(CG1)
+        CG1_height.interpolate(dot(self.k, x))
+        height_above_surface = Function(CG1)
+
+        # Turn height into columnwise data
+        columnwise_height, index_data = self.coords.get_column_data(CG1_height, self)
+
+        # Find minimum height in each column
+        surface_height_1d = np.min(columnwise_height, axis=1)
+        height_above_surface_data = columnwise_height - surface_height_1d[:, None]
+
+        self.coords.set_field_from_column_data(height_above_surface,
+                                               height_above_surface_data,
+                                               index_data)
+
+        self.height_above_surface = height_above_surface
+
 
 def construct_domain_metadata(mesh, coords, on_sphere):
     """

gusto/forcing.py

@@ -6,7 +6,7 @@
 )
 from gusto.fml import drop, replace_subject, name
 from gusto.labels import (
-    transport, diffusion, time_derivative, hydrostatic
+    transport, diffusion, time_derivative, hydrostatic, physics_label
 )
 from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual
 
@@ -47,7 +47,7 @@ def __init__(self, equation, alpha):
 
         # drop terms relating to transport and diffusion
         residual = equation.residual.label_map(
-            lambda t: any(t.has_label(transport, diffusion, return_tuple=True)), drop)
+            lambda t: any(t.has_label(transport, diffusion, physics_label, return_tuple=True)), drop)
 
         # the lhs of both of the explicit and implicit solvers is just
         # the time derivative form