diff --git a/LICENSE b/LICENSE index fc95efea0..e15fad85d 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ The MIT License (MIT) -Copyright (c) 2016 Firedrake +Copyright (c) 2016 Gusto Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/examples/shallow_water/moist_thermal_williamson5.py b/examples/shallow_water/moist_thermal_williamson5.py index 314f7a75e..f022e1058 100644 --- a/examples/shallow_water/moist_thermal_williamson5.py +++ b/examples/shallow_water/moist_thermal_williamson5.py @@ -4,7 +4,7 @@ """ from gusto import * from firedrake import (IcosahedralSphereMesh, SpatialCoordinate, - as_vector, pi, sqrt, min_value, exp, cos) + as_vector, pi, sqrt, min_value, exp, cos, sin) import sys # ----------------------------------------------------------------- # diff --git a/gusto/__init__.py b/gusto/__init__.py index 37f870157..f39e35d64 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -1,29 +1,15 @@ # Start logging first, incase anything goes wrong -from gusto.logging import * # noqa +from gusto.core.logging import * # noqa set_log_handler() -from gusto.active_tracers import * # noqa -from gusto.common_forms import * # noqa -from gusto.configuration import * # noqa -from gusto.coord_transforms import * # noqa -from gusto.domain import * # noqa +from gusto.core import * # noqa from gusto.diagnostics import * # noqa -from gusto.diffusion_methods import * # noqa from gusto.equations import * # noqa -from gusto.forcing import * # noqa -from gusto.initialisation_tools import * # noqa -from gusto.io import * # noqa -from gusto.labels import * # noqa -from gusto.limiters import * # noqa -from gusto.linear_solvers import * # noqa -from gusto.meshes import * # noqa -from gusto.numerical_integrator import * # noqa +from gusto.initialisation import * # noqa from gusto.physics import * # noqa -from gusto.preconditioners import * # noqa from gusto.recovery import * # noqa +from gusto.solvers import * # noqa from gusto.spatial_methods import * # noqa from gusto.time_discretisation import * # noqa -from gusto.timeloop import * # noqa -from gusto.transport_methods import * # noqa -from gusto.wrappers import * # noqa +from gusto.timestepping import * # noqa diff --git a/gusto/core/__init__.py b/gusto/core/__init__.py new file mode 100644 index 000000000..163f9ddff --- /dev/null +++ b/gusto/core/__init__.py @@ -0,0 +1,11 @@ +from gusto.core.configuration import * # noqa +from gusto.core.coordinates import * # noqa +from gusto.core.coord_transforms import * # noqa +from gusto.core.domain import * # noqa +from gusto.core.fields import * # noqa +from gusto.core.function_spaces import * # noqa +from gusto.core.io import * # noqa +from gusto.core.kernels import * # noqa +from gusto.core.labels import * # noqa +from gusto.core.logging import * # noqa +from gusto.core.meshes import * # noqa \ No newline at end of file diff --git a/gusto/configuration.py b/gusto/core/configuration.py similarity index 100% rename from gusto/configuration.py rename to gusto/core/configuration.py diff --git a/gusto/coord_transforms.py b/gusto/core/coord_transforms.py similarity index 100% rename from gusto/coord_transforms.py rename to gusto/core/coord_transforms.py diff --git a/gusto/coordinates.py b/gusto/core/coordinates.py similarity index 99% rename from gusto/coordinates.py rename to gusto/core/coordinates.py index c86c85d3d..69aebbf13 100644 --- a/gusto/coordinates.py +++ b/gusto/core/coordinates.py @@ -3,8 +3,8 @@ Coordinate fields are stored in specified VectorFunctionSpaces. """ -from gusto.coord_transforms import lonlatr_from_xyz, rotated_lonlatr_coords -from gusto.logging import logger +from gusto.core.coord_transforms import lonlatr_from_xyz, rotated_lonlatr_coords +from gusto.core.logging import logger from firedrake import SpatialCoordinate, Function import numpy as np import pandas as pd diff --git a/gusto/domain.py b/gusto/core/domain.py similarity index 98% rename from gusto/domain.py rename to gusto/core/domain.py index 8f22083c6..0cf23fb7f 100644 --- a/gusto/domain.py +++ b/gusto/core/domain.py @@ -4,8 +4,8 @@ model's time interval. """ -from gusto.coordinates import Coordinates -from gusto.function_spaces import Spaces, check_degree_args +from gusto.core.coordinates import Coordinates +from gusto.core.function_spaces import Spaces, check_degree_args from firedrake import (Constant, SpatialCoordinate, sqrt, CellNormal, cross, inner, grad, VectorFunctionSpace, Function, FunctionSpace, perp) diff --git a/gusto/fields.py b/gusto/core/fields.py similarity index 100% rename from gusto/fields.py rename to gusto/core/fields.py diff --git a/gusto/function_spaces.py b/gusto/core/function_spaces.py similarity index 99% rename from gusto/function_spaces.py rename to gusto/core/function_spaces.py index 185273420..97820d454 100644 --- a/gusto/function_spaces.py +++ b/gusto/core/function_spaces.py @@ -3,7 +3,7 @@ used by the model. """ -from gusto import logger +from gusto.core.logging import logger from firedrake import (HCurl, HDiv, FunctionSpace, FiniteElement, TensorProductElement, interval) diff --git a/gusto/io.py b/gusto/core/io.py similarity index 99% rename from gusto/io.py rename to gusto/core/io.py index 40352085f..a9a09f1fe 100644 --- a/gusto/io.py +++ b/gusto/core/io.py @@ -6,13 +6,13 @@ import sys import time from gusto.diagnostics import Diagnostics, CourantNumber -from gusto.meshes import get_flat_latlon_mesh +from gusto.core.meshes import get_flat_latlon_mesh from firedrake import (Function, functionspaceimpl, Constant, DumbCheckpoint, FILE_CREATE, FILE_READ, CheckpointFile) from firedrake.output import VTKFile from pyop2.mpi import MPI import numpy as np -from gusto.logging import logger, update_logfile_location +from gusto.core.logging import logger, update_logfile_location __all__ = ["pick_up_mesh", "IO"] diff --git a/gusto/kernels.py b/gusto/core/kernels.py similarity index 100% rename from gusto/kernels.py rename to gusto/core/kernels.py diff --git a/gusto/labels.py b/gusto/core/labels.py similarity index 98% rename from gusto/labels.py rename to gusto/core/labels.py index 61d132bea..06ce4fa8b 100644 --- a/gusto/labels.py +++ b/gusto/core/labels.py @@ -3,7 +3,7 @@ import ufl from firedrake import Function from firedrake.fml import Term, Label, LabelledForm -from gusto.configuration import IntegrateByParts, TransportEquationType +from gusto.core.configuration import IntegrateByParts, TransportEquationType from types import MethodType dynamics_label = Label("dynamics", validator=lambda value: type(value) is str) diff --git a/gusto/logging.py b/gusto/core/logging.py similarity index 99% rename from gusto/logging.py rename to gusto/core/logging.py index dc116275b..c6df8eab4 100644 --- a/gusto/logging.py +++ b/gusto/core/logging.py @@ -1,7 +1,7 @@ """Gusto Logging Module All logging functionality for Gusto is controlled from -``gusto.logging``. A logger object ``logging.getLogger("gusto")`` is +``gusto.core.logging``. A logger object ``logging.getLogger("gusto")`` is created internally. The primary means of configuration is via environment variables, the diff --git a/gusto/meshes.py b/gusto/core/meshes.py similarity index 100% rename from gusto/meshes.py rename to gusto/core/meshes.py diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py deleted file mode 100644 index fd6d7d4ae..000000000 --- a/gusto/diagnostics.py +++ /dev/null @@ -1,1830 +0,0 @@ -"""Common diagnostic fields.""" - - -from firedrake import assemble, dot, dx, Function, sqrt, ln, \ - TestFunction, TrialFunction, Constant, grad, inner, curl, \ - LinearVariationalProblem, LinearVariationalSolver, FacetNormal, \ - ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, avg, jump, pi, \ - TensorFunctionSpace, SpatialCoordinate, as_vector, \ - Projector, Interpolator, FunctionSpace, FiniteElement, \ - TensorProductElement -from firedrake.assign import Assigner -from ufl.domain import extract_unique_domain - -from abc import ABCMeta, abstractmethod, abstractproperty -from gusto.equations import CompressibleEulerEquations -import gusto.thermodynamics as tde -from gusto.coord_transforms import rotated_lonlatr_vectors -from gusto.recovery import Recoverer, BoundaryMethod -from gusto.active_tracers import TracerVariableType, Phases -from gusto.logging import logger -from gusto.kernels import MinKernel, MaxKernel -import numpy as np - -__all__ = ["Diagnostics", "CourantNumber", "Gradient", "XComponent", "YComponent", - "ZComponent", "MeridionalComponent", "ZonalComponent", "RadialComponent", - "RichardsonNumber", "Entropy", "PhysicalEntropy", "DynamicEntropy", "Energy", "KineticEnergy", - "ShallowWaterKineticEnergy", "ShallowWaterPotentialEnergy", "ShallowWaterPotentialEnstrophy", - "CompressibleKineticEnergy", "Exner", "Sum", "Difference", "SteadyStateError", - "Perturbation", "Theta_e", "InternalEnergy", "PotentialEnergy", - "ThermodynamicKineticEnergy", "Dewpoint", "Temperature", "Theta_d", - "RelativeHumidity", "Pressure", "Exner_Vt", "HydrostaticImbalance", "Precipitation", - "PotentialVorticity", "RelativeVorticity", "AbsoluteVorticity", "Divergence", - "BruntVaisalaFrequencySquared", "TracerDensity"] - - -class Diagnostics(object): - """ - Stores all diagnostic fields, and controls global diagnostics computation. - - This object stores the diagnostic fields to be output, and the computation - of global values from them (such as global maxima or norms). - """ - - available_diagnostics = ["min", "max", "rms", "l2", "total"] - - def __init__(self, *fields): - """ - Args: - *fields: list of :class:`Function` objects of fields to be output. - """ - - self.fields = list(fields) - - def register(self, *fields): - """ - Registers diagnostic fields for outputting. - - Args: - *fields: list of :class:`Function` objects of fields to be output. - """ - - fset = set(self.fields) - for f in fields: - if f not in fset: - self.fields.append(f) - - @staticmethod - def min(f): - """ - Finds the global minimum DoF value of a field. - - Args: - f (:class:`Function`): field to compute diagnostic for. - """ - min_kernel = MinKernel() - return min_kernel.apply(f) - - @staticmethod - def max(f): - """ - Finds the global maximum DoF value of a field. - - Args: - f (:class:`Function`): field to compute diagnostic for. - """ - max_kernel = MaxKernel() - return max_kernel.apply(f) - - @staticmethod - def rms(f): - """ - Calculates the root-mean-square of a field. - - Args: - f (:class:`Function`): field to compute diagnostic for. - """ - - area = assemble(1*dx(domain=extract_unique_domain(f))) - return sqrt(assemble(inner(f, f)*dx)/area) - - @staticmethod - def l2(f): - """ - Calculates the L2 norm of a field. - - Args: - f (:class:`Function`): field to compute diagnostic for. - """ - - return sqrt(assemble(inner(f, f)*dx)) - - @staticmethod - def total(f): - """ - Calculates the total of a field. Only applicable for fields with - scalar-values. - - Args: - f (:class:`Function`): field to compute diagnostic for. - """ - - if len(f.ufl_shape) == 0: - return assemble(f * dx) - else: - pass - - -class DiagnosticField(object, metaclass=ABCMeta): - """Base object to represent diagnostic fields for outputting.""" - def __init__(self, space=None, method='interpolate', required_fields=()): - """ - 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 (). - """ - - assert method in ['interpolate', 'project', 'solve', 'assign'], \ - f'Invalid evaluation method {self.method} for diagnostic {self.name}' - - self._initialised = False - self.required_fields = required_fields - self.space = space - self.method = method - self.expr = None - self.to_dump = True - - # Property to allow graceful failures if solve method not valid - if not hasattr(self, "solve_implemented"): - self.solve_implemented = False - - if method == 'solve' and not self.solve_implemented: - raise NotImplementedError(f'Solve method has not been implemented for diagnostic {self.name}') - - @abstractproperty - def name(self): - """The name of this diagnostic field""" - pass - - @abstractmethod - 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. - """ - - 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) - - 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 - - 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': - self.evaluator.assign() - elif self.method == 'project': - self.evaluator.project() - elif self.method == 'solve': - self.evaluator.solve() - - def __call__(self): - """Return the diagnostic field computed from the current state.""" - self.compute() - return self.field - - -class CourantNumber(DiagnosticField): - """Dimensionless Courant number diagnostic field.""" - name = "CourantNumber" - - def __init__(self, velocity='u', component='whole', name=None, to_dump=True, - space=None, method='interpolate', required_fields=()): - """ - Args: - velocity (str or :class:`ufl.Expr`, optional): the velocity field to - take the Courant number of. Can be a string referring to an - existing field, or an expression. If it is an expression, the - name argument is required. Defaults to 'u'. - component (str, optional): the component of the velocity to use for - calculating the Courant number. Valid values are "whole", - "horizontal" or "vertical". Defaults to "whole". - name (str, optional): the name to append to "CourantNumber" to form - the name of this diagnostic. This argument must be provided if - the velocity is an expression (rather than a string). Defaults - to None. - to_dump (bool, optional): whether this diagnostic should be dumped. - Defaults to True. - 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 (). - """ - if component not in ["whole", "horizontal", "vertical"]: - raise ValueError(f'component arg {component} not valid. Allowed ' - + 'values are "whole", "horizontal" and "vertical"') - self.component = component - - # Work out whether to take Courant number from field or expression - if type(velocity) is str: - # Default name should just be CourantNumber - if velocity == 'u': - self.name = 'CourantNumber' - elif name is None: - self.name = 'CourantNumber_'+velocity - else: - self.name = 'CourantNumber_'+name - if component != 'whole': - self.name += '_'+component - else: - if name is None: - raise ValueError('CourantNumber diagnostic: if provided ' - + 'velocity is an expression then the name ' - + 'argument must be provided') - self.name = 'CourantNumber_'+name - - self.velocity = velocity - super().__init__(space=space, method=method, required_fields=required_fields) - - # Done after super init to ensure that it is not always set to True - self.to_dump = to_dump - - 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. - """ - - V = FunctionSpace(domain.mesh, "DG", 0) - test = TestFunction(V) - cell_volume = Function(V) - self.cell_flux = Function(V) - - # Calculate cell volumes - One = Function(V).assign(1) - assemble(One*test*dx, tensor=cell_volume) - - # Get the velocity that is being used - if type(self.velocity) is str: - u = state_fields(self.velocity) - else: - u = self.velocity - - # Determine the component of the velocity - if self.component == "whole": - u_expr = u - elif self.component == "vertical": - u_expr = dot(u, domain.k)*domain.k - elif self.component == "horizontal": - u_expr = u - dot(u, domain.k)*domain.k - - # Work out which facet integrals to use - if domain.mesh.extruded: - dS_calc = dS_v + dS_h - ds_calc = ds_v + ds_t + ds_b - else: - dS_calc = dS - ds_calc = ds - - # Set up form for DG flux - n = FacetNormal(domain.mesh) - un = 0.5*(inner(-u_expr, n) + abs(inner(-u_expr, n))) - self.cell_flux_form = 2*avg(un*test)*dS_calc + un*test*ds_calc - - # Final Courant number expression - self.expr = self.cell_flux * domain.dt / cell_volume - - super().setup(domain, state_fields) - - def compute(self): - """Compute the diagnostic field from the current state.""" - - assemble(self.cell_flux_form, tensor=self.cell_flux) - super().compute() - - -class Gradient(DiagnosticField): - """Diagnostic for computing the gradient of fields.""" - def __init__(self, name, space=None, method='solve'): - """ - Args: - name (str): name of the field to compute the gradient of. - 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 'solve'. - """ - self.fname = name - self.solve_implemented = True - super().__init__(space=space, method=method, required_fields=(name,)) - - @property - def name(self): - """Gives the name of this diagnostic field.""" - return self.fname+"_gradient" - - 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. - """ - f = state_fields(self.fname) - - mesh_dim = domain.mesh.geometric_dimension() - try: - field_dim = state_fields(self.fname).ufl_shape[0] - except IndexError: - field_dim = 1 - shape = (mesh_dim, ) * field_dim - space = TensorFunctionSpace(domain.mesh, "CG", 1, shape=shape, name=f'Tensor{field_dim}_CG1') - - if self.method != 'solve': - self.expr = grad(f) - - super().setup(domain, state_fields, space=space) - - # Set up problem now that self.field has been set up - if self.method == 'solve': - test = TestFunction(space) - trial = TrialFunction(space) - n = FacetNormal(domain.mesh) - a = inner(test, trial)*dx - L = -inner(div(test), f)*dx - if space.extruded: - L += dot(dot(test, n), f)*(ds_t + ds_b) - prob = LinearVariationalProblem(a, L, self.field) - self.evaluator = LinearVariationalSolver(prob) - - -class Divergence(DiagnosticField): - """Diagnostic for computing the divergence of vector-valued fields.""" - def __init__(self, name='u', space=None, method='interpolate'): - """ - Args: - name (str, optional): name of the field to compute the gradient of. - Defaults to 'u', in which case this takes the divergence of the - wind field. - space (:class:`FunctionSpace`, optional): the function space to - evaluate the diagnostic field in. Defaults to None, in which - case the default space is the domain's DG space. - 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.fname = name - super().__init__(space=space, method=method, required_fields=(self.fname,)) - - @property - def name(self): - """Gives the name of this diagnostic field.""" - return self.fname+"_divergence" - - 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. - """ - f = state_fields(self.fname) - self.expr = div(f) - space = domain.spaces("DG") - super().setup(domain, state_fields, space=space) - - -class VectorComponent(DiagnosticField): - """Base diagnostic for orthogonal components of vector-valued fields.""" - def __init__(self, name, space=None, method='interpolate'): - """ - Args: - name (str): name of the field to compute the component of. - space (:class:`FunctionSpace`, optional): the function space to - evaluate the diagnostic field in. Defaults to None, in which - case the default space is the domain's DG space. - 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.fname = name - super().__init__(space=space, method=method, required_fields=(name,)) - - def setup(self, domain, state_fields, unit_vector): - """ - 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. - unit_vector (:class:`ufl.Expr`): the unit vector to extract the - component for. This assumes an orthogonal coordinate system. - """ - f = state_fields(self.fname) - self.expr = dot(f, unit_vector) - super().setup(domain, state_fields) - - -class XComponent(VectorComponent): - """The geocentric Cartesian x-component of a vector-valued field.""" - @property - def name(self): - """Gives the name of this diagnostic field.""" - return self.fname+"_x" - - 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. - """ - dim = domain.mesh.topological_dimension() - e_x = as_vector([Constant(1.0)]+[Constant(0.0)]*(dim-1)) - super().setup(domain, state_fields, e_x) - - -class YComponent(VectorComponent): - """The geocentric Cartesian y-component of a vector-valued field.""" - @property - def name(self): - """Gives the name of this diagnostic field.""" - return self.fname+"_y" - - 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. - """ - assert domain.metadata['domain_type'] not in ['interval', 'vertical_slice'], \ - f'Y-component diagnostic cannot be used with domain {domain.metadata["domain_type"]}' - dim = domain.mesh.topological_dimension() - e_y = as_vector([Constant(0.0), Constant(1.0)]+[Constant(0.0)]*(dim-2)) - super().setup(domain, state_fields, e_y) - - -class ZComponent(VectorComponent): - """The geocentric Cartesian z-component of a vector-valued field.""" - @property - def name(self): - """Gives the name of this diagnostic field.""" - return self.fname+"_z" - - 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. - """ - assert domain.metadata['domain_type'] not in ['interval', 'plane'], \ - f'Z-component diagnostic cannot be used with domain {domain.metadata["domain_type"]}' - dim = domain.mesh.topological_dimension() - e_x = as_vector([Constant(0.0)]*(dim-1)+[Constant(1.0)]) - super().setup(domain, state_fields, e_x) - - -class SphericalComponent(VectorComponent): - """Base diagnostic for computing spherical-polar components of fields.""" - def __init__(self, name, rotated_pole=None, space=None, method='interpolate'): - """ - Args: - name (str): name of the field to compute the component of. - rotated_pole (tuple of floats, optional): a tuple of floats - (lon, lat) of the new pole, in the original coordinate system. - The longitude and latitude must be expressed in radians. - Defaults to None, corresponding to a pole of (0, pi/2). - space (:class:`FunctionSpace`, optional): the function space to - evaluate the diagnostic field in. Defaults to None, in which - case the default space is the domain's DG space. - 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.rotated_pole = (0.0, pi/2) if rotated_pole is None else rotated_pole - super().__init__(name=name, space=space, method=method) - - def _check_args(self, domain, field): - """ - Checks the validity of the domain and field for taking the spherical - component diagnostic. - - Args: - domain (:class:`Domain`): the model's domain object. - field (:class:`Function`): the field to take the component of. - """ - - # check geometric dimension is 3D - if domain.mesh.geometric_dimension() != 3: - raise ValueError('Spherical components only work when the geometric dimension is 3!') - - if np.prod(field.ufl_shape) != 3: - raise ValueError('Components can only be found of a vector function space in 3D.') - - -class MeridionalComponent(SphericalComponent): - """The meridional component of a vector-valued field.""" - @property - def name(self): - """Gives the name of this diagnostic field.""" - return self.fname+"_meridional" - - 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. - """ - f = state_fields(self.fname) - self._check_args(domain, f) - xyz = SpatialCoordinate(domain.mesh) - _, e_lat, _ = rotated_lonlatr_vectors(xyz, self.rotated_pole) - super().setup(domain, state_fields, e_lat) - - -class ZonalComponent(SphericalComponent): - """The zonal component of a vector-valued field.""" - @property - def name(self): - """Gives the name of this diagnostic field.""" - return self.fname+"_zonal" - - 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. - """ - f = state_fields(self.fname) - self._check_args(domain, f) - xyz = SpatialCoordinate(domain.mesh) - e_lon, _, _ = rotated_lonlatr_vectors(xyz, self.rotated_pole) - super().setup(domain, state_fields, e_lon) - - -class RadialComponent(SphericalComponent): - """The radial component of a vector-valued field.""" - @property - def name(self): - """Gives the name of this diagnostic field.""" - return self.fname+"_radial" - - 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. - """ - f = state_fields(self.fname) - self._check_args(domain, f) - xyz = SpatialCoordinate(domain.mesh) - _, _, e_r = rotated_lonlatr_vectors(xyz, self.rotated_pole) - super().setup(domain, state_fields, e_r) - - -class RichardsonNumber(DiagnosticField): - """Dimensionless Richardson number diagnostic field.""" - name = "RichardsonNumber" - - def __init__(self, density_field, factor=1., space=None, method='interpolate'): - u""" - Args: - density_field (str): the name of the density field. - factor (float, optional): a factor to multiply the Brunt-Väisälä - frequency by. Defaults to 1. - 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'. - """ - super().__init__(space=space, method=method, required_fields=(density_field, "u_gradient")) - self.density_field = density_field - self.factor = Constant(factor) - - 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. - """ - rho_grad = self.density_field+"_gradient" - grad_density = state_fields(rho_grad) - gradu = state_fields("u_gradient") - - denom = 0. - z_dim = domain.mesh.geometric_dimension() - 1 - u_dim = state_fields("u").ufl_shape[0] - for i in range(u_dim-1): - denom += gradu[i, z_dim]**2 - Nsq = self.factor*grad_density[z_dim] - self.expr = Nsq/denom - super().setup(domain, state_fields) - - -class Entropy(DiagnosticField): - """Base diagnostic field for entropy diagnostic """ - - 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.equations = equations - if isinstance(equations, CompressibleEulerEquations): - required_fields = ['rho', 'theta'] - else: - raise NotImplementedError(f'entropy not yet implemented for {type(equations)}') - - super().__init__(space=space, method=method, required_fields=tuple(required_fields)) - - -class PhysicalEntropy(Entropy): - u"""Physical entropy ρ*ln(θ) for Compressible Euler equations""" - name = "PhysicalEntropy" - - 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. - """ - rho = state_fields('rho') - theta = state_fields('theta') - self.expr = rho * ln(theta) - super().setup(domain, state_fields) - - -class DynamicEntropy(Entropy): - u"""Dynamic entropy 0.5*ρ*θ^2 for Compressible Euler equations""" - name = "DynamicEntropy" - - 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. - """ - rho = state_fields('rho') - theta = state_fields('theta') - self.expr = 0.5 * rho * theta**2 - super().setup(domain, state_fields) - - -# TODO: unify all energy diagnostics -- should be based on equation -class Energy(DiagnosticField): - """Base diagnostic field for computing energy density fields.""" - def kinetic(self, u, factor=None): - """ - Computes a kinetic energy term. - - Args: - u (:class:`ufl.Expr`): the velocity variable. - factor (:class:`ufl.Expr`, optional): factor to multiply the term by - (e.g. a density variable). Defaults to None. - - Returns: - :class:`ufl.Expr`: the kinetic energy term - """ - if factor is not None: - energy = 0.5*factor*dot(u, u) - else: - energy = 0.5*dot(u, u) - return energy - - -class KineticEnergy(Energy): - """Diagnostic kinetic energy density.""" - name = "KineticEnergy" - - def __init__(self, space=None, method='interpolate'): - """ - 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'. - """ - super().__init__(space=space, method=method, required_fields=("u")) - - 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. - """ - u = state_fields("u") - self.expr = self.kinetic(u) - super().setup(domain, state_fields) - - -class ShallowWaterKineticEnergy(Energy): - """Diagnostic shallow-water kinetic energy density.""" - name = "ShallowWaterKineticEnergy" - - def __init__(self, space=None, method='interpolate'): - """ - 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'. - """ - super().__init__(space=space, method=method, required_fields=("D", "u")) - - 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. - """ - u = state_fields("u") - D = state_fields("D") - self.expr = self.kinetic(u, D) - super().setup(domain, state_fields) - - -class ShallowWaterPotentialEnergy(Energy): - """Diagnostic shallow-water potential energy density.""" - name = "ShallowWaterPotentialEnergy" - - def __init__(self, parameters, space=None, method='interpolate'): - """ - Args: - parameters (:class:`ShallowWaterParameters`): the configuration - object containing the physical parameters for this equation. - 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 = parameters - super().__init__(space=space, method=method, required_fields=("D")) - - 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. - """ - g = self.parameters.g - D = state_fields("D") - self.expr = 0.5*g*D**2 - super().setup(domain, state_fields) - - -class ShallowWaterPotentialEnstrophy(DiagnosticField): - """Diagnostic (dry) compressible kinetic energy density.""" - def __init__(self, base_field_name="PotentialVorticity", space=None, - method='interpolate'): - """ - Args: - base_field_name (str, optional): the base potential vorticity field - to compute the enstrophy from. Defaults to "PotentialVorticity". - 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'. - """ - base_enstrophy_names = ["PotentialVorticity", "RelativeVorticity", "AbsoluteVorticity"] - if base_field_name not in base_enstrophy_names: - raise ValueError( - f"Don't know how to compute enstrophy with base_field_name={base_field_name};" - + f"base_field_name should be one of {base_enstrophy_names}") - # Work out required fields - if base_field_name in ["PotentialVorticity", "AbsoluteVorticity"]: - required_fields = (base_field_name, "D") - elif base_field_name == "RelativeVorticity": - required_fields = (base_field_name, "D", "coriolis") - else: - raise NotImplementedError(f'Enstrophy with vorticity {base_field_name} not implemented') - - super().__init__(space=space, method=method, required_fields=required_fields) - self.base_field_name = base_field_name - - @property - def name(self): - """Gives the name of this diagnostic field.""" - base_name = "SWPotentialEnstrophy" - return "_from_".join((base_name, self.base_field_name)) - - 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. - """ - if self.base_field_name == "PotentialVorticity": - pv = state_fields("PotentialVorticity") - D = state_fields("D") - self.expr = 0.5*pv**2*D - elif self.base_field_name == "RelativeVorticity": - zeta = state_fields("RelativeVorticity") - D = state_fields("D") - f = state_fields("coriolis") - self.expr = 0.5*(zeta + f)**2/D - elif self.base_field_name == "AbsoluteVorticity": - zeta_abs = state_fields("AbsoluteVorticity") - D = state_fields("D") - self.expr = 0.5*(zeta_abs)**2/D - else: - raise NotImplementedError(f'Enstrophy with {self.base_field_name} not implemented') - super().setup(domain, state_fields) - - -class CompressibleKineticEnergy(Energy): - """Diagnostic (dry) compressible kinetic energy density.""" - name = "CompressibleKineticEnergy" - - def __init__(self, space=None, method='interpolate'): - """ - 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'. - """ - super().__init__(space=space, method=method, required_fields=("rho", "u")) - - 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. - """ - u = state_fields("u") - rho = state_fields("rho") - self.expr = self.kinetic(u, rho) - super().setup(domain, state_fields) - - -class Exner(DiagnosticField): - """The diagnostic Exner pressure field.""" - def __init__(self, parameters, reference=False, space=None, method='interpolate'): - """ - Args: - parameters (:class:`CompressibleParameters`): the configuration - object containing the physical parameters for this equation. - reference (bool, optional): whether to compute the reference Exner - pressure field or not. Defaults to False. - 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 = parameters - self.reference = reference - if reference: - self.rho_name = "rho_bar" - self.theta_name = "theta_bar" - else: - self.rho_name = "rho" - self.theta_name = "theta" - super().__init__(space=space, method=method, required_fields=(self.rho_name, self.theta_name)) - - @property - def name(self): - """Gives the name of this diagnostic field.""" - if self.reference: - return "Exner_bar" - else: - return "Exner" - - 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. - """ - rho = state_fields(self.rho_name) - theta = state_fields(self.theta_name) - self.expr = tde.exner_pressure(self.parameters, rho, theta) - 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): - """ - Args: - field_name1 (str): the name of one field to be added. - field_name2 (str): the name of the other field to be added. - """ - super().__init__(method='assign', required_fields=(field_name1, field_name2)) - self.field_name1 = field_name1 - self.field_name2 = field_name2 - - @property - def name(self): - """Gives the name of this diagnostic field.""" - return self.field_name1+"_plus_"+self.field_name2 - - 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. - """ - field1 = state_fields(self.field_name1) - field2 = state_fields(self.field_name2) - space = field1.function_space() - self.expr = field1 + field2 - super().setup(domain, state_fields, space=space) - - -class Difference(DiagnosticField): - """Base diagnostic for calculating the difference between two fields.""" - def __init__(self, field_name1, field_name2): - """ - Args: - field_name1 (str): the name of the field to be subtracted from. - field_name2 (str): the name of the field to be subtracted. - """ - super().__init__(method='assign', required_fields=(field_name1, field_name2)) - self.field_name1 = field_name1 - self.field_name2 = field_name2 - - @property - def name(self): - """Gives the name of this diagnostic field.""" - return self.field_name1+"_minus_"+self.field_name2 - - 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. - """ - - field1 = state_fields(self.field_name1) - field2 = state_fields(self.field_name2) - self.expr = field1 - field2 - space = field1.function_space() - super().setup(domain, state_fields, space=space) - - -class SteadyStateError(Difference): - """Base diagnostic for computing the steady-state error in a field.""" - def __init__(self, name): - """ - Args: - name (str): name of the field to take the steady-state error of. - """ - self.field_name1 = name - self.field_name2 = name+'_init' - DiagnosticField.__init__(self, method='assign', required_fields=(name, self.field_name2)) - - 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. - """ - # Check if initial field already exists -- otherwise needs creating - if not hasattr(state_fields, self.field_name2): - field1 = state_fields(self.field_name1) - field2 = state_fields(self.field_name2, space=field1.function_space(), - pick_up=True, dump=False) - # Attach state fields to self so that we can pick it up in compute - self.state_fields = state_fields - # The initial value for fields may not have already been set yet so we - # postpone setting it until the compute method is called - self.init_field_set = False - else: - field1 = state_fields(self.field_name1) - field2 = state_fields(self.field_name2, space=field1.function_space(), - pick_up=True, dump=False) - # By default set this new field to the current value - # This may be overwritten if picking up from a checkpoint - field2.assign(field1) - self.state_fields = state_fields - self.init_field_set = True - - super().setup(domain, state_fields) - - def compute(self): - # The first time the compute method is called we set the initial field. - # We do not want to do this if picking up from a checkpoint - if not self.init_field_set: - # Set initial field - full_field = self.state_fields(self.field_name1) - init_field = self.state_fields(self.field_name2) - init_field.assign(full_field) - - self.init_field_set = True - - super().compute() - - @property - def name(self): - """Gives the name of this diagnostic field.""" - return self.field_name1+"_error" - - -class Perturbation(Difference): - """Base diagnostic for computing perturbations from a reference profile.""" - def __init__(self, name): - """ - Args: - name (str): name of the field to take the perturbation of. - """ - self.field_name1 = name - self.field_name2 = name+'_bar' - DiagnosticField.__init__(self, method='assign', required_fields=(name, self.field_name2)) - - @property - def name(self): - """Gives the name of this diagnostic field.""" - return self.field_name1+"_perturbation" - - 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. - """ - # Check if initial field already exists -- otherwise needs creating - if not hasattr(state_fields, self.field_name2): - field1 = state_fields(self.field_name1) - _ = state_fields(self.field_name2, space=field1.function_space(), - pick_up=True, dump=False) - - super().setup(domain, state_fields) - - -# TODO: unify thermodynamic diagnostics -class ThermodynamicDiagnostic(DiagnosticField): - """Base thermodynamic diagnostic field, computing many common fields.""" - - 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.equations = equations - self.parameters = equations.parameters - # Work out required fields - if isinstance(equations, CompressibleEulerEquations): - required_fields = ['rho', 'theta'] - if equations.active_tracers is not None: - for active_tracer in equations.active_tracers: - if active_tracer.chemical == 'H2O': - required_fields.append(active_tracer.name) - else: - raise NotImplementedError(f'Thermodynamic diagnostics not implemented for {type(equations)}') - super().__init__(space=space, method=method, required_fields=tuple(required_fields)) - - def _setup_thermodynamics(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. - """ - - self.Vtheta = domain.spaces('theta') - h_deg = self.Vtheta.ufl_element().degree()[0] - v_deg = self.Vtheta.ufl_element().degree()[1]-1 - boundary_method = BoundaryMethod.extruded if (v_deg == 0 and h_deg == 0) else None - - # Extract all fields - self.rho = state_fields("rho") - self.theta = state_fields("theta") - # Rho must be averaged to Vtheta - self.rho_averaged = Function(self.Vtheta) - self.recoverer = Recoverer(self.rho, self.rho_averaged, boundary_method=boundary_method) - - zero_expr = Constant(0.0)*self.theta - self.r_v = zero_expr # Water vapour - self.r_l = zero_expr # Liquid water - self.r_t = zero_expr # All water mixing ratios - for active_tracer in self.equations.active_tracers: - if active_tracer.chemical == "H2O": - if active_tracer.variable_type != TracerVariableType.mixing_ratio: - raise NotImplementedError('Only mixing ratio tracers are implemented') - if active_tracer.phase == Phases.gas: - self.r_v += state_fields(active_tracer.name) - elif active_tracer.phase == Phases.liquid: - self.r_l += state_fields(active_tracer.name) - self.r_t += state_fields(active_tracer.name) - - # Store the most common expressions - self.exner = tde.exner_pressure(self.parameters, self.rho_averaged, self.theta) - self.T = tde.T(self.parameters, self.theta, self.exner, r_v=self.r_v) - self.p = tde.p(self.parameters, self.exner) - - def compute(self): - """Compute the thermodynamic diagnostic.""" - self.recoverer.project() - super().compute() - - -class Theta_e(ThermodynamicDiagnostic): - """The moist equivalent potential temperature diagnostic field.""" - name = "Theta_e" - - 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. - """ - self._setup_thermodynamics(domain, state_fields) - self.expr = tde.theta_e(self.parameters, self.T, self.p, self.r_v, self.r_t) - super().setup(domain, state_fields, space=self.Vtheta) - - -class InternalEnergy(ThermodynamicDiagnostic): - """The moist compressible internal energy density.""" - name = "InternalEnergy" - - 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. - """ - self._setup_thermodynamics(domain, state_fields) - self.expr = tde.internal_energy(self.parameters, self.rho_averaged, self.T, r_v=self.r_v, r_l=self.r_l) - super().setup(domain, state_fields, space=self.Vtheta) - - -class PotentialEnergy(ThermodynamicDiagnostic): - """The moist compressible potential energy density.""" - name = "PotentialEnergy" - - 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. - """ - x = SpatialCoordinate(domain.mesh) - self._setup_thermodynamics(domain, state_fields) - z = Function(self.rho_averaged.function_space()) - z.interpolate(dot(x, domain.k)) - self.expr = self.rho_averaged * (1 + self.r_t) * self.parameters.g * z - super().setup(domain, state_fields, space=domain.spaces("DG")) - - -# TODO: this needs consolidating with energy diagnostics -class ThermodynamicKineticEnergy(ThermodynamicDiagnostic): - """The moist compressible kinetic energy density.""" - name = "ThermodynamicKineticEnergy" - - 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.equations = equations - self.parameters = equations.parameters - # Work out required fields - if isinstance(equations, CompressibleEulerEquations): - required_fields = ['rho', 'u'] - if equations.active_tracers is not None: - for active_tracer in equations.active_tracers: - if active_tracer.chemical == 'H2O': - required_fields.append(active_tracer.name) - else: - raise NotImplementedError(f'Thermodynamic K.E. 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. - """ - u = state_fields('u') - self._setup_thermodynamics(domain, state_fields) - self.expr = 0.5 * self.rho_averaged * (1 + self.r_t) * dot(u, u) - super().setup(domain, state_fields, space=domain.spaces("DG")) - - -class Dewpoint(ThermodynamicDiagnostic): - """The dewpoint temperature diagnostic field.""" - name = "Dewpoint" - - 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. - """ - 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) - - -class Temperature(ThermodynamicDiagnostic): - """The absolute temperature diagnostic field.""" - name = "Temperature" - - 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. - """ - self._setup_thermodynamics(domain, state_fields) - self.expr = self.T - super().setup(domain, state_fields, space=self.Vtheta) - - -class Theta_d(ThermodynamicDiagnostic): - """The dry potential temperature diagnostic field.""" - name = "Theta_d" - - 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. - """ - self._setup_thermodynamics(domain, state_fields) - self.expr = self.theta / (1 + self.r_v * self.parameters.R_v / self.parameters.R_d) - super().setup(domain, state_fields, space=self.Vtheta) - - -class RelativeHumidity(ThermodynamicDiagnostic): - """The relative humidity diagnostic field.""" - name = "RelativeHumidity" - - 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. - """ - self._setup_thermodynamics(domain, state_fields) - self.expr = tde.RH(self.parameters, self.r_v, self.T, self.p) - super().setup(domain, state_fields, space=self.Vtheta) - - -class Pressure(ThermodynamicDiagnostic): - """The pressure field computed in the 'theta' space.""" - name = "Pressure_Vt" - - 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. - """ - self._setup_thermodynamics(domain, state_fields) - self.expr = self.p - super().setup(domain, state_fields, space=self.Vtheta) - - -class Exner_Vt(ThermodynamicDiagnostic): - """The Exner pressure field computed in the 'theta' space.""" - name = "Exner_Vt" - - 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. - """ - self._setup_thermodynamics(domain, state_fields) - self.expr = self.exner - super().setup(domain, state_fields, space=self.Vtheta) - - -# TODO: this doesn't contain the effects of moisture -# TODO: this has not been implemented for other equation sets -class HydrostaticImbalance(DiagnosticField): - """Hydrostatic imbalance diagnostic field.""" - name = "HydrostaticImbalance" - - 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'. - """ - # Work out required fields - if isinstance(equations, CompressibleEulerEquations): - required_fields = ['rho', 'theta', 'rho_bar', 'theta_bar'] - if equations.active_tracers is not None: - for active_tracer in equations.active_tracers: - if active_tracer.chemical == 'H2O': - required_fields.append(active_tracer.name) - self.equations = equations - self.parameters = equations.parameters - else: - raise NotImplementedError(f'Hydrostatic Imbalance not implemented for {type(equations)}') - super().__init__(space=space, method=method, required_fields=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. - """ - Vu = domain.spaces("HDiv") - rho = state_fields("rho") - rhobar = state_fields("rho_bar") - theta = state_fields("theta") - thetabar = state_fields("theta_bar") - exner = tde.exner_pressure(self.parameters, rho, theta) - exnerbar = tde.exner_pressure(self.parameters, rhobar, thetabar) - - cp = Constant(self.parameters.cp) - n = FacetNormal(domain.mesh) - - # TODO: not sure about this expression! - # Gravity does not appear, and why are there reference profiles? - F = TrialFunction(Vu) - w = TestFunction(Vu) - imbalance = Function(Vu) - a = inner(w, F)*dx - L = (- cp*div((theta-thetabar)*w)*exnerbar*dx - + cp*jump((theta-thetabar)*w, n)*avg(exnerbar)*dS_v - - cp*div(thetabar*w)*(exner-exnerbar)*dx - + cp*jump(thetabar*w, n)*avg(exner-exnerbar)*dS_v) - - bcs = self.equations.bcs['u'] - - imbalanceproblem = LinearVariationalProblem(a, L, imbalance, bcs=bcs) - self.imbalance_solver = LinearVariationalSolver(imbalanceproblem) - self.expr = dot(imbalance, domain.k) - super().setup(domain, state_fields) - - def compute(self): - """Compute and return the diagnostic field from the current state. - """ - self.imbalance_solver.solve() - super().compute() - - -class Precipitation(DiagnosticField): - """ - The total precipitation falling through the domain's bottom surface. - - This is normalised by unit area, giving a result in kg / m^2. - """ - name = "Precipitation" - - def __init__(self): - self.solve_implemented = True - required_fields = ('rain', 'rainfall_velocity', 'rho') - super().__init__(method='solve', required_fields=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. - """ - if not hasattr(domain.spaces, "DG0"): - DG0 = domain.spaces.create_space("DG0", "DG", 0) - else: - DG0 = domain.spaces("DG0") - assert DG0.extruded, 'Cannot compute precipitation on a non-extruded mesh' - self.space = DG0 - - # Gather fields - rain = state_fields('rain') - rho = state_fields('rho') - v = state_fields('rainfall_velocity') - # Set up problem - self.phi = TestFunction(DG0) - flux = TrialFunction(DG0) - self.flux = Function(DG0) # Flux to solve for - area = Function(DG0) # Need to compute normalisation (area) - - eqn_lhs = self.phi * flux * dx - area_rhs = self.phi * ds_b - eqn_rhs = domain.dt * self.phi * (rain * dot(- v, domain.k) * rho / area) * ds_b - - # Compute area normalisation - area_prob = LinearVariationalProblem(eqn_lhs, area_rhs, area) - area_solver = LinearVariationalSolver(area_prob) - area_solver.solve() - - # setup solver - rain_prob = LinearVariationalProblem(eqn_lhs, eqn_rhs, self.flux) - self.solver = LinearVariationalSolver(rain_prob) - self.field = state_fields(self.name, space=DG0, dump=True, pick_up=True) - # Initialise field to zero, if picking up this will be overridden - self.field.assign(0.0) - - def compute(self): - """Increment the precipitation diagnostic.""" - self.solver.solve() - self.field.assign(self.field + self.flux) - - -class Vorticity(DiagnosticField): - """Base diagnostic field class for shallow-water vorticity variables.""" - - def setup(self, domain, state_fields, vorticity_type=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. - vorticity_type (str, optional): denotes which type of vorticity to - be computed ('relative', 'absolute' or 'potential'). Defaults to - None. - """ - - vorticity_types = ["relative", "absolute", "potential"] - if vorticity_type not in vorticity_types: - raise ValueError(f"vorticity type must be one of {vorticity_types}, not {vorticity_type}") - space = domain.spaces("H1") - - u = state_fields("u") - if vorticity_type in ["absolute", "potential"]: - f = state_fields("coriolis") - if vorticity_type == "potential": - D = state_fields("D") - - if self.method != 'solve': - if vorticity_type == "potential": - self.expr = (curl(u) + f) / D - elif vorticity_type == "absolute": - self.expr = curl(u) + f - elif vorticity_type == "relative": - self.expr = curl(u) - - super().setup(domain, state_fields, space=space) - - # Set up problem now that self.field has been set up - if self.method == 'solve': - gamma = TestFunction(space) - q = TrialFunction(space) - - if vorticity_type == "potential": - a = q*gamma*D*dx - else: - a = q*gamma*dx - - L = (- inner(domain.perp(grad(gamma)), u))*dx - if vorticity_type != "relative": - f = state_fields("coriolis") - L += gamma*f*dx - - problem = LinearVariationalProblem(a, L, self.field) - self.evaluator = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"}) - - -class PotentialVorticity(Vorticity): - u"""Diagnostic field for shallow-water potential vorticity, q=(∇×(u+f))/D""" - name = "PotentialVorticity" - - def __init__(self, space=None, method='solve'): - """ - 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 'solve'. - """ - self.solve_implemented = True - super().__init__(space=space, method=method, - required_fields=('u', 'D', 'coriolis')) - - 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. - """ - super().setup(domain, state_fields, vorticity_type="potential") - - -class AbsoluteVorticity(Vorticity): - u"""Diagnostic field for absolute vorticity, ζ=∇×(u+f)""" - name = "AbsoluteVorticity" - - def __init__(self, space=None, method='solve'): - """ - 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 'solve'. - """ - self.solve_implemented = True - super().__init__(space=space, method=method, required_fields=('u', 'coriolis')) - - 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. - """ - super().setup(domain, state_fields, vorticity_type="absolute") - - -class RelativeVorticity(Vorticity): - u"""Diagnostic field for relative vorticity, ζ=∇×u""" - name = "RelativeVorticity" - - def __init__(self, space=None, method='solve'): - """ - 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 'solve'. - """ - self.solve_implemented = True - super().__init__(space=space, method=method, required_fields=('u',)) - - 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. - """ - super().setup(domain, state_fields, vorticity_type="relative") - - -class TracerDensity(DiagnosticField): - """Diagnostic for computing the density of a tracer. This is - computed as the product of a mixing ratio and dry density""" - - @property - def name(self): - """Gives the name of this diagnostic field. This records - the mixing ratio and density names, in case multiple tracer - densities are used.""" - return "TracerDensity_"+self.mixing_ratio_name+'_'+self.density_name - - def __init__(self, mixing_ratio_name, density_name, space=None, method='interpolate'): - """ - Args: - mixing_ratio_name (str): the name of the tracer mixing ratio variable - density_name (str): the name of the tracer density variable - space (:class:`FunctionSpace`, optional): the function space to - evaluate the diagnostic field in. Defaults to None, in which - case a new space will be constructed for this diagnostic. This - space will have enough a high enough degree to accurately compute - the product of the mixing ratio and density. - method (str, optional): a string specifying the method of evaluation - for this diagnostic. Valid options are 'interpolate', 'project' and - 'assign'. Defaults to 'interpolate'. - """ - super().__init__(space=space, method=method, required_fields=(mixing_ratio_name, density_name)) - - self.mixing_ratio_name = mixing_ratio_name - self.density_name = density_name - - 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. - """ - m_X = state_fields(self.mixing_ratio_name) - rho_d = state_fields(self.density_name) - self.expr = m_X*rho_d - - if self.space is None: - # Construct a space for the diagnostic that has enough - # degrees to accurately capture the tracer density. This - # will be the sum of the degrees of the individual mixing ratio - # and density function spaces. - m_X_space = m_X.function_space() - rho_d_space = rho_d.function_space() - - if domain.spaces.extruded_mesh: - # Extract the base horizontal and vertical elements - # for the mixing ratio and density. - m_X_horiz = m_X_space.ufl_element().sub_elements[0] - m_X_vert = m_X_space.ufl_element().sub_elements[1] - rho_d_horiz = rho_d_space.ufl_element().sub_elements[0] - rho_d_vert = rho_d_space.ufl_element().sub_elements[1] - - horiz_degree = m_X_horiz.degree() + rho_d_horiz.degree() - vert_degree = m_X_vert.degree() + rho_d_vert.degree() - - cell = domain.mesh._base_mesh.ufl_cell().cellname() - horiz_elt = FiniteElement('DG', cell, horiz_degree) - vert_elt = FiniteElement('DG', cell, vert_degree) - elt = TensorProductElement(horiz_elt, vert_elt) - else: - m_X_degree = m_X_space.ufl_element().degree() - rho_d_degree = rho_d_space.ufl_element().degree() - degree = m_X_degree + rho_d_degree - - cell = domain.mesh.ufl_cell().cellname() - elt = FiniteElement('DG', cell, degree) - - tracer_density_space = FunctionSpace(domain.mesh, elt, name='tracer_density_space') - super().setup(domain, state_fields, space=tracer_density_space) - - else: - super().setup(domain, state_fields) diff --git a/gusto/diagnostics/__init__.py b/gusto/diagnostics/__init__.py new file mode 100644 index 000000000..2d37128b9 --- /dev/null +++ b/gusto/diagnostics/__init__.py @@ -0,0 +1,3 @@ +from gusto.diagnostics.diagnostics import * # noqa +from gusto.diagnostics.shallow_water_diagnostics import * # noqa +from gusto.diagnostics.compressible_euler_diagnostics import * # noqa \ No newline at end of file diff --git a/gusto/diagnostics/compressible_euler_diagnostics.py b/gusto/diagnostics/compressible_euler_diagnostics.py new file mode 100644 index 000000000..ce03468d6 --- /dev/null +++ b/gusto/diagnostics/compressible_euler_diagnostics.py @@ -0,0 +1,653 @@ +"""Common diagnostic fields for the compressible Euler equations.""" + +from firedrake import (dot, dx, Function, ln, TestFunction, TrialFunction, + Constant, grad, inner, LinearVariationalProblem, + LinearVariationalSolver, FacetNormal, ds_b, dS_v, div, + avg, jump, SpatialCoordinate) + +from gusto.diagnostics.diagnostics import DiagnosticField, Energy +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", + "Temperature", "Theta_d", "RelativeHumidity", "Pressure", "Exner_Vt", + "HydrostaticImbalance", "Precipitation", "BruntVaisalaFrequencySquared"] + + +class RichardsonNumber(DiagnosticField): + """Dimensionless Richardson number diagnostic field.""" + name = "RichardsonNumber" + + def __init__(self, density_field, factor=1., space=None, method='interpolate'): + u""" + Args: + density_field (str): the name of the density field. + factor (float, optional): a factor to multiply the Brunt-Väisälä + frequency by. Defaults to 1. + 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'. + """ + super().__init__(space=space, method=method, required_fields=(density_field, "u_gradient")) + self.density_field = density_field + self.factor = Constant(factor) + + 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. + """ + rho_grad = self.density_field+"_gradient" + grad_density = state_fields(rho_grad) + gradu = state_fields("u_gradient") + + denom = 0. + z_dim = domain.mesh.geometric_dimension() - 1 + u_dim = state_fields("u").ufl_shape[0] + for i in range(u_dim-1): + denom += gradu[i, z_dim]**2 + Nsq = self.factor*grad_density[z_dim] + self.expr = Nsq/denom + super().setup(domain, state_fields) + + +class Entropy(DiagnosticField): + """Base diagnostic field for entropy diagnostic """ + + 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.equations = equations + if isinstance(equations, CompressibleEulerEquations): + required_fields = ['rho', 'theta'] + else: + raise NotImplementedError(f'entropy not yet implemented for {type(equations)}') + + super().__init__(space=space, method=method, required_fields=tuple(required_fields)) + + +class PhysicalEntropy(Entropy): + u"""Physical entropy ρ*ln(θ) for Compressible Euler equations""" + name = "PhysicalEntropy" + + 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. + """ + rho = state_fields('rho') + theta = state_fields('theta') + self.expr = rho * ln(theta) + super().setup(domain, state_fields) + + +class DynamicEntropy(Entropy): + u"""Dynamic entropy 0.5*ρ*θ^2 for Compressible Euler equations""" + name = "DynamicEntropy" + + 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. + """ + rho = state_fields('rho') + theta = state_fields('theta') + self.expr = 0.5 * rho * theta**2 + super().setup(domain, state_fields) + + +class CompressibleKineticEnergy(Energy): + """Diagnostic (dry) compressible kinetic energy density.""" + name = "CompressibleKineticEnergy" + + def __init__(self, space=None, method='interpolate'): + """ + 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'. + """ + super().__init__(space=space, method=method, required_fields=("rho", "u")) + + 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. + """ + u = state_fields("u") + rho = state_fields("rho") + self.expr = self.kinetic(u, rho) + super().setup(domain, state_fields) + + +class Exner(DiagnosticField): + """The diagnostic Exner pressure field.""" + def __init__(self, parameters, reference=False, space=None, method='interpolate'): + """ + Args: + parameters (:class:`CompressibleParameters`): the configuration + object containing the physical parameters for this equation. + reference (bool, optional): whether to compute the reference Exner + pressure field or not. Defaults to False. + 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 = parameters + self.reference = reference + if reference: + self.rho_name = "rho_bar" + self.theta_name = "theta_bar" + else: + self.rho_name = "rho" + self.theta_name = "theta" + super().__init__(space=space, method=method, required_fields=(self.rho_name, self.theta_name)) + + @property + def name(self): + """Gives the name of this diagnostic field.""" + if self.reference: + return "Exner_bar" + else: + return "Exner" + + 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. + """ + rho = state_fields(self.rho_name) + theta = state_fields(self.theta_name) + self.expr = tde.exner_pressure(self.parameters, rho, theta) + 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) + + +# TODO: unify thermodynamic diagnostics +class ThermodynamicDiagnostic(DiagnosticField): + """Base thermodynamic diagnostic field, computing many common fields.""" + + 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.equations = equations + self.parameters = equations.parameters + # Work out required fields + if isinstance(equations, CompressibleEulerEquations): + required_fields = ['rho', 'theta'] + if equations.active_tracers is not None: + for active_tracer in equations.active_tracers: + if active_tracer.chemical == 'H2O': + required_fields.append(active_tracer.name) + else: + raise NotImplementedError(f'Thermodynamic diagnostics not implemented for {type(equations)}') + super().__init__(space=space, method=method, required_fields=tuple(required_fields)) + + def _setup_thermodynamics(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. + """ + + self.Vtheta = domain.spaces('theta') + h_deg = self.Vtheta.ufl_element().degree()[0] + v_deg = self.Vtheta.ufl_element().degree()[1]-1 + boundary_method = BoundaryMethod.extruded if (v_deg == 0 and h_deg == 0) else None + + # Extract all fields + self.rho = state_fields("rho") + self.theta = state_fields("theta") + # Rho must be averaged to Vtheta + self.rho_averaged = Function(self.Vtheta) + self.recoverer = Recoverer(self.rho, self.rho_averaged, boundary_method=boundary_method) + + zero_expr = Constant(0.0)*self.theta + self.r_v = zero_expr # Water vapour + self.r_l = zero_expr # Liquid water + self.r_t = zero_expr # All water mixing ratios + for active_tracer in self.equations.active_tracers: + if active_tracer.chemical == "H2O": + if active_tracer.variable_type != TracerVariableType.mixing_ratio: + raise NotImplementedError('Only mixing ratio tracers are implemented') + if active_tracer.phase == Phases.gas: + self.r_v += state_fields(active_tracer.name) + elif active_tracer.phase == Phases.liquid: + self.r_l += state_fields(active_tracer.name) + self.r_t += state_fields(active_tracer.name) + + # Store the most common expressions + self.exner = tde.exner_pressure(self.parameters, self.rho_averaged, self.theta) + self.T = tde.T(self.parameters, self.theta, self.exner, r_v=self.r_v) + self.p = tde.p(self.parameters, self.exner) + + def compute(self): + """Compute the thermodynamic diagnostic.""" + self.recoverer.project() + super().compute() + + +class Theta_e(ThermodynamicDiagnostic): + """The moist equivalent potential temperature diagnostic field.""" + name = "Theta_e" + + 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. + """ + self._setup_thermodynamics(domain, state_fields) + self.expr = tde.theta_e(self.parameters, self.T, self.p, self.r_v, self.r_t) + super().setup(domain, state_fields, space=self.Vtheta) + + +class InternalEnergy(ThermodynamicDiagnostic): + """The moist compressible internal energy density.""" + name = "InternalEnergy" + + 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. + """ + self._setup_thermodynamics(domain, state_fields) + self.expr = tde.internal_energy(self.parameters, self.rho_averaged, self.T, r_v=self.r_v, r_l=self.r_l) + super().setup(domain, state_fields, space=self.Vtheta) + + +class PotentialEnergy(ThermodynamicDiagnostic): + """The moist compressible potential energy density.""" + name = "PotentialEnergy" + + 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. + """ + x = SpatialCoordinate(domain.mesh) + self._setup_thermodynamics(domain, state_fields) + z = Function(self.rho_averaged.function_space()) + z.interpolate(dot(x, domain.k)) + self.expr = self.rho_averaged * (1 + self.r_t) * self.parameters.g * z + super().setup(domain, state_fields, space=domain.spaces("DG")) + + +# TODO: this needs consolidating with energy diagnostics +class ThermodynamicKineticEnergy(ThermodynamicDiagnostic): + """The moist compressible kinetic energy density.""" + name = "ThermodynamicKineticEnergy" + + 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.equations = equations + self.parameters = equations.parameters + # Work out required fields + if isinstance(equations, CompressibleEulerEquations): + required_fields = ['rho', 'u'] + if equations.active_tracers is not None: + for active_tracer in equations.active_tracers: + if active_tracer.chemical == 'H2O': + required_fields.append(active_tracer.name) + else: + raise NotImplementedError(f'Thermodynamic K.E. 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. + """ + u = state_fields('u') + self._setup_thermodynamics(domain, state_fields) + self.expr = 0.5 * self.rho_averaged * (1 + self.r_t) * dot(u, u) + super().setup(domain, state_fields, space=domain.spaces("DG")) + + +class Dewpoint(ThermodynamicDiagnostic): + """The dewpoint temperature diagnostic field.""" + name = "Dewpoint" + + 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. + """ + 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) + + +class Temperature(ThermodynamicDiagnostic): + """The absolute temperature diagnostic field.""" + name = "Temperature" + + 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. + """ + self._setup_thermodynamics(domain, state_fields) + self.expr = self.T + super().setup(domain, state_fields, space=self.Vtheta) + + +class Theta_d(ThermodynamicDiagnostic): + """The dry potential temperature diagnostic field.""" + name = "Theta_d" + + 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. + """ + self._setup_thermodynamics(domain, state_fields) + self.expr = self.theta / (1 + self.r_v * self.parameters.R_v / self.parameters.R_d) + super().setup(domain, state_fields, space=self.Vtheta) + + +class RelativeHumidity(ThermodynamicDiagnostic): + """The relative humidity diagnostic field.""" + name = "RelativeHumidity" + + 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. + """ + self._setup_thermodynamics(domain, state_fields) + self.expr = tde.RH(self.parameters, self.r_v, self.T, self.p) + super().setup(domain, state_fields, space=self.Vtheta) + + +class Pressure(ThermodynamicDiagnostic): + """The pressure field computed in the 'theta' space.""" + name = "Pressure_Vt" + + 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. + """ + self._setup_thermodynamics(domain, state_fields) + self.expr = self.p + super().setup(domain, state_fields, space=self.Vtheta) + + +class Exner_Vt(ThermodynamicDiagnostic): + """The Exner pressure field computed in the 'theta' space.""" + name = "Exner_Vt" + + 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. + """ + self._setup_thermodynamics(domain, state_fields) + self.expr = self.exner + super().setup(domain, state_fields, space=self.Vtheta) + + +# TODO: this doesn't contain the effects of moisture +# TODO: this has not been implemented for other equation sets +class HydrostaticImbalance(DiagnosticField): + """Hydrostatic imbalance diagnostic field.""" + name = "HydrostaticImbalance" + + 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'. + """ + # Work out required fields + if isinstance(equations, CompressibleEulerEquations): + required_fields = ['rho', 'theta', 'rho_bar', 'theta_bar'] + if equations.active_tracers is not None: + for active_tracer in equations.active_tracers: + if active_tracer.chemical == 'H2O': + required_fields.append(active_tracer.name) + self.equations = equations + self.parameters = equations.parameters + else: + raise NotImplementedError(f'Hydrostatic Imbalance not implemented for {type(equations)}') + super().__init__(space=space, method=method, required_fields=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. + """ + Vu = domain.spaces("HDiv") + rho = state_fields("rho") + rhobar = state_fields("rho_bar") + theta = state_fields("theta") + thetabar = state_fields("theta_bar") + exner = tde.exner_pressure(self.parameters, rho, theta) + exnerbar = tde.exner_pressure(self.parameters, rhobar, thetabar) + + cp = Constant(self.parameters.cp) + n = FacetNormal(domain.mesh) + + # TODO: not sure about this expression! + # Gravity does not appear, and why are there reference profiles? + F = TrialFunction(Vu) + w = TestFunction(Vu) + imbalance = Function(Vu) + a = inner(w, F)*dx + L = (- cp*div((theta-thetabar)*w)*exnerbar*dx + + cp*jump((theta-thetabar)*w, n)*avg(exnerbar)*dS_v + - cp*div(thetabar*w)*(exner-exnerbar)*dx + + cp*jump(thetabar*w, n)*avg(exner-exnerbar)*dS_v) + + bcs = self.equations.bcs['u'] + + imbalanceproblem = LinearVariationalProblem(a, L, imbalance, bcs=bcs) + self.imbalance_solver = LinearVariationalSolver(imbalanceproblem) + self.expr = dot(imbalance, domain.k) + super().setup(domain, state_fields) + + def compute(self): + """Compute and return the diagnostic field from the current state. + """ + self.imbalance_solver.solve() + super().compute() + + +class Precipitation(DiagnosticField): + """ + The total precipitation falling through the domain's bottom surface. + + This is normalised by unit area, giving a result in kg / m^2. + """ + name = "Precipitation" + + def __init__(self): + self.solve_implemented = True + required_fields = ('rain', 'rainfall_velocity', 'rho') + super().__init__(method='solve', required_fields=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. + """ + if not hasattr(domain.spaces, "DG0"): + DG0 = domain.spaces.create_space("DG0", "DG", 0) + else: + DG0 = domain.spaces("DG0") + assert DG0.extruded, 'Cannot compute precipitation on a non-extruded mesh' + self.space = DG0 + + # Gather fields + rain = state_fields('rain') + rho = state_fields('rho') + v = state_fields('rainfall_velocity') + # Set up problem + self.phi = TestFunction(DG0) + flux = TrialFunction(DG0) + self.flux = Function(DG0) # Flux to solve for + area = Function(DG0) # Need to compute normalisation (area) + + eqn_lhs = self.phi * flux * dx + area_rhs = self.phi * ds_b + eqn_rhs = domain.dt * self.phi * (rain * dot(- v, domain.k) * rho / area) * ds_b + + # Compute area normalisation + area_prob = LinearVariationalProblem(eqn_lhs, area_rhs, area) + area_solver = LinearVariationalSolver(area_prob) + area_solver.solve() + + # setup solver + rain_prob = LinearVariationalProblem(eqn_lhs, eqn_rhs, self.flux) + self.solver = LinearVariationalSolver(rain_prob) + self.field = state_fields(self.name, space=DG0, dump=True, pick_up=True) + # Initialise field to zero, if picking up this will be overridden + self.field.assign(0.0) + + def compute(self): + """Increment the precipitation diagnostic.""" + self.solver.solve() + self.field.assign(self.field + self.flux) diff --git a/gusto/diagnostics/diagnostics.py b/gusto/diagnostics/diagnostics.py new file mode 100644 index 000000000..5fabe2eb5 --- /dev/null +++ b/gusto/diagnostics/diagnostics.py @@ -0,0 +1,920 @@ +"""Common diagnostic fields.""" + + +from firedrake import (assemble, 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, + TensorProductElement) +from firedrake.assign import Assigner +from ufl.domain import extract_unique_domain + +from abc import ABCMeta, abstractmethod, abstractproperty +from gusto.core.coord_transforms import rotated_lonlatr_vectors +from gusto.core.logging import logger +from gusto.core.kernels import MinKernel, MaxKernel +import numpy as np + +__all__ = ["Diagnostics", "DiagnosticField", "CourantNumber", "Gradient", + "XComponent", "YComponent", "ZComponent", "MeridionalComponent", + "ZonalComponent", "RadialComponent", "Energy", "KineticEnergy", + "Sum", "Difference", "SteadyStateError", "Perturbation", + "Divergence", "TracerDensity"] + + +class Diagnostics(object): + """ + Stores all diagnostic fields, and controls global diagnostics computation. + + This object stores the diagnostic fields to be output, and the computation + of global values from them (such as global maxima or norms). + """ + + available_diagnostics = ["min", "max", "rms", "l2", "total"] + + def __init__(self, *fields): + """ + Args: + *fields: list of :class:`Function` objects of fields to be output. + """ + + self.fields = list(fields) + + def register(self, *fields): + """ + Registers diagnostic fields for outputting. + + Args: + *fields: list of :class:`Function` objects of fields to be output. + """ + + fset = set(self.fields) + for f in fields: + if f not in fset: + self.fields.append(f) + + @staticmethod + def min(f): + """ + Finds the global minimum DoF value of a field. + + Args: + f (:class:`Function`): field to compute diagnostic for. + """ + min_kernel = MinKernel() + return min_kernel.apply(f) + + @staticmethod + def max(f): + """ + Finds the global maximum DoF value of a field. + + Args: + f (:class:`Function`): field to compute diagnostic for. + """ + max_kernel = MaxKernel() + return max_kernel.apply(f) + + @staticmethod + def rms(f): + """ + Calculates the root-mean-square of a field. + + Args: + f (:class:`Function`): field to compute diagnostic for. + """ + + area = assemble(1*dx(domain=extract_unique_domain(f))) + return sqrt(assemble(inner(f, f)*dx)/area) + + @staticmethod + def l2(f): + """ + Calculates the L2 norm of a field. + + Args: + f (:class:`Function`): field to compute diagnostic for. + """ + + return sqrt(assemble(inner(f, f)*dx)) + + @staticmethod + def total(f): + """ + Calculates the total of a field. Only applicable for fields with + scalar-values. + + Args: + f (:class:`Function`): field to compute diagnostic for. + """ + + if len(f.ufl_shape) == 0: + return assemble(f * dx) + else: + pass + + +class DiagnosticField(object, metaclass=ABCMeta): + """Base object to represent diagnostic fields for outputting.""" + def __init__(self, space=None, method='interpolate', required_fields=()): + """ + 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 (). + """ + + assert method in ['interpolate', 'project', 'solve', 'assign'], \ + f'Invalid evaluation method {self.method} for diagnostic {self.name}' + + self._initialised = False + self.required_fields = required_fields + self.space = space + self.method = method + self.expr = None + self.to_dump = True + + # Property to allow graceful failures if solve method not valid + if not hasattr(self, "solve_implemented"): + self.solve_implemented = False + + if method == 'solve' and not self.solve_implemented: + raise NotImplementedError(f'Solve method has not been implemented for diagnostic {self.name}') + + @abstractproperty + def name(self): + """The name of this diagnostic field""" + pass + + @abstractmethod + 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. + """ + + 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) + + 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 + + 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': + self.evaluator.assign() + elif self.method == 'project': + self.evaluator.project() + elif self.method == 'solve': + self.evaluator.solve() + + def __call__(self): + """Return the diagnostic field computed from the current state.""" + self.compute() + return self.field + + +class CourantNumber(DiagnosticField): + """Dimensionless Courant number diagnostic field.""" + name = "CourantNumber" + + def __init__(self, velocity='u', component='whole', name=None, to_dump=True, + space=None, method='interpolate', required_fields=()): + """ + Args: + velocity (str or :class:`ufl.Expr`, optional): the velocity field to + take the Courant number of. Can be a string referring to an + existing field, or an expression. If it is an expression, the + name argument is required. Defaults to 'u'. + component (str, optional): the component of the velocity to use for + calculating the Courant number. Valid values are "whole", + "horizontal" or "vertical". Defaults to "whole". + name (str, optional): the name to append to "CourantNumber" to form + the name of this diagnostic. This argument must be provided if + the velocity is an expression (rather than a string). Defaults + to None. + to_dump (bool, optional): whether this diagnostic should be dumped. + Defaults to True. + 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 (). + """ + if component not in ["whole", "horizontal", "vertical"]: + raise ValueError(f'component arg {component} not valid. Allowed ' + + 'values are "whole", "horizontal" and "vertical"') + self.component = component + + # Work out whether to take Courant number from field or expression + if type(velocity) is str: + # Default name should just be CourantNumber + if velocity == 'u': + self.name = 'CourantNumber' + elif name is None: + self.name = 'CourantNumber_'+velocity + else: + self.name = 'CourantNumber_'+name + if component != 'whole': + self.name += '_'+component + else: + if name is None: + raise ValueError('CourantNumber diagnostic: if provided ' + + 'velocity is an expression then the name ' + + 'argument must be provided') + self.name = 'CourantNumber_'+name + + self.velocity = velocity + super().__init__(space=space, method=method, required_fields=required_fields) + + # Done after super init to ensure that it is not always set to True + self.to_dump = to_dump + + 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. + """ + + V = FunctionSpace(domain.mesh, "DG", 0) + test = TestFunction(V) + cell_volume = Function(V) + self.cell_flux = Function(V) + + # Calculate cell volumes + One = Function(V).assign(1) + assemble(One*test*dx, tensor=cell_volume) + + # Get the velocity that is being used + if type(self.velocity) is str: + u = state_fields(self.velocity) + else: + u = self.velocity + + # Determine the component of the velocity + if self.component == "whole": + u_expr = u + elif self.component == "vertical": + u_expr = dot(u, domain.k)*domain.k + elif self.component == "horizontal": + u_expr = u - dot(u, domain.k)*domain.k + + # Work out which facet integrals to use + if domain.mesh.extruded: + dS_calc = dS_v + dS_h + ds_calc = ds_v + ds_t + ds_b + else: + dS_calc = dS + ds_calc = ds + + # Set up form for DG flux + n = FacetNormal(domain.mesh) + un = 0.5*(inner(-u_expr, n) + abs(inner(-u_expr, n))) + self.cell_flux_form = 2*avg(un*test)*dS_calc + un*test*ds_calc + + # Final Courant number expression + self.expr = self.cell_flux * domain.dt / cell_volume + + super().setup(domain, state_fields) + + def compute(self): + """Compute the diagnostic field from the current state.""" + + assemble(self.cell_flux_form, tensor=self.cell_flux) + super().compute() + + +class Gradient(DiagnosticField): + """Diagnostic for computing the gradient of fields.""" + def __init__(self, name, space=None, method='solve'): + """ + Args: + name (str): name of the field to compute the gradient of. + 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 'solve'. + """ + self.fname = name + self.solve_implemented = True + super().__init__(space=space, method=method, required_fields=(name,)) + + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.fname+"_gradient" + + 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. + """ + f = state_fields(self.fname) + + mesh_dim = domain.mesh.geometric_dimension() + try: + field_dim = state_fields(self.fname).ufl_shape[0] + except IndexError: + field_dim = 1 + shape = (mesh_dim, ) * field_dim + space = TensorFunctionSpace(domain.mesh, "CG", 1, shape=shape, name=f'Tensor{field_dim}_CG1') + + if self.method != 'solve': + self.expr = grad(f) + + super().setup(domain, state_fields, space=space) + + # Set up problem now that self.field has been set up + if self.method == 'solve': + test = TestFunction(space) + trial = TrialFunction(space) + n = FacetNormal(domain.mesh) + a = inner(test, trial)*dx + L = -inner(div(test), f)*dx + if space.extruded: + L += dot(dot(test, n), f)*(ds_t + ds_b) + prob = LinearVariationalProblem(a, L, self.field) + self.evaluator = LinearVariationalSolver(prob) + + +class Divergence(DiagnosticField): + """Diagnostic for computing the divergence of vector-valued fields.""" + def __init__(self, name='u', space=None, method='interpolate'): + """ + Args: + name (str, optional): name of the field to compute the gradient of. + Defaults to 'u', in which case this takes the divergence of the + wind field. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case the default space is the domain's DG space. + 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.fname = name + super().__init__(space=space, method=method, required_fields=(self.fname,)) + + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.fname+"_divergence" + + 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. + """ + f = state_fields(self.fname) + self.expr = div(f) + space = domain.spaces("DG") + super().setup(domain, state_fields, space=space) + + +class VectorComponent(DiagnosticField): + """Base diagnostic for orthogonal components of vector-valued fields.""" + def __init__(self, name, space=None, method='interpolate'): + """ + Args: + name (str): name of the field to compute the component of. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case the default space is the domain's DG space. + 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.fname = name + super().__init__(space=space, method=method, required_fields=(name,)) + + def setup(self, domain, state_fields, unit_vector): + """ + 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. + unit_vector (:class:`ufl.Expr`): the unit vector to extract the + component for. This assumes an orthogonal coordinate system. + """ + f = state_fields(self.fname) + self.expr = dot(f, unit_vector) + super().setup(domain, state_fields) + + +class XComponent(VectorComponent): + """The geocentric Cartesian x-component of a vector-valued field.""" + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.fname+"_x" + + 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. + """ + dim = domain.mesh.topological_dimension() + e_x = as_vector([Constant(1.0)]+[Constant(0.0)]*(dim-1)) + super().setup(domain, state_fields, e_x) + + +class YComponent(VectorComponent): + """The geocentric Cartesian y-component of a vector-valued field.""" + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.fname+"_y" + + 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. + """ + assert domain.metadata['domain_type'] not in ['interval', 'vertical_slice'], \ + f'Y-component diagnostic cannot be used with domain {domain.metadata["domain_type"]}' + dim = domain.mesh.topological_dimension() + e_y = as_vector([Constant(0.0), Constant(1.0)]+[Constant(0.0)]*(dim-2)) + super().setup(domain, state_fields, e_y) + + +class ZComponent(VectorComponent): + """The geocentric Cartesian z-component of a vector-valued field.""" + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.fname+"_z" + + 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. + """ + assert domain.metadata['domain_type'] not in ['interval', 'plane'], \ + f'Z-component diagnostic cannot be used with domain {domain.metadata["domain_type"]}' + dim = domain.mesh.topological_dimension() + e_x = as_vector([Constant(0.0)]*(dim-1)+[Constant(1.0)]) + super().setup(domain, state_fields, e_x) + + +class SphericalComponent(VectorComponent): + """Base diagnostic for computing spherical-polar components of fields.""" + def __init__(self, name, rotated_pole=None, space=None, method='interpolate'): + """ + Args: + name (str): name of the field to compute the component of. + rotated_pole (tuple of floats, optional): a tuple of floats + (lon, lat) of the new pole, in the original coordinate system. + The longitude and latitude must be expressed in radians. + Defaults to None, corresponding to a pole of (0, pi/2). + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case the default space is the domain's DG space. + 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.rotated_pole = (0.0, pi/2) if rotated_pole is None else rotated_pole + super().__init__(name=name, space=space, method=method) + + def _check_args(self, domain, field): + """ + Checks the validity of the domain and field for taking the spherical + component diagnostic. + + Args: + domain (:class:`Domain`): the model's domain object. + field (:class:`Function`): the field to take the component of. + """ + + # check geometric dimension is 3D + if domain.mesh.geometric_dimension() != 3: + raise ValueError('Spherical components only work when the geometric dimension is 3!') + + if np.prod(field.ufl_shape) != 3: + raise ValueError('Components can only be found of a vector function space in 3D.') + + +class MeridionalComponent(SphericalComponent): + """The meridional component of a vector-valued field.""" + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.fname+"_meridional" + + 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. + """ + f = state_fields(self.fname) + self._check_args(domain, f) + xyz = SpatialCoordinate(domain.mesh) + _, e_lat, _ = rotated_lonlatr_vectors(xyz, self.rotated_pole) + super().setup(domain, state_fields, e_lat) + + +class ZonalComponent(SphericalComponent): + """The zonal component of a vector-valued field.""" + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.fname+"_zonal" + + 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. + """ + f = state_fields(self.fname) + self._check_args(domain, f) + xyz = SpatialCoordinate(domain.mesh) + e_lon, _, _ = rotated_lonlatr_vectors(xyz, self.rotated_pole) + super().setup(domain, state_fields, e_lon) + + +class RadialComponent(SphericalComponent): + """The radial component of a vector-valued field.""" + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.fname+"_radial" + + 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. + """ + f = state_fields(self.fname) + self._check_args(domain, f) + xyz = SpatialCoordinate(domain.mesh) + _, _, e_r = rotated_lonlatr_vectors(xyz, self.rotated_pole) + super().setup(domain, state_fields, e_r) + + +# TODO: unify all energy diagnostics -- should be based on equation +class Energy(DiagnosticField): + """Base diagnostic field for computing energy density fields.""" + def kinetic(self, u, factor=None): + """ + Computes a kinetic energy term. + + Args: + u (:class:`ufl.Expr`): the velocity variable. + factor (:class:`ufl.Expr`, optional): factor to multiply the term by + (e.g. a density variable). Defaults to None. + + Returns: + :class:`ufl.Expr`: the kinetic energy term + """ + if factor is not None: + energy = 0.5*factor*dot(u, u) + else: + energy = 0.5*dot(u, u) + return energy + + +class KineticEnergy(Energy): + """Diagnostic kinetic energy density.""" + name = "KineticEnergy" + + def __init__(self, space=None, method='interpolate'): + """ + 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'. + """ + super().__init__(space=space, method=method, required_fields=("u")) + + 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. + """ + u = state_fields("u") + self.expr = self.kinetic(u) + super().setup(domain, state_fields) + + +class Sum(DiagnosticField): + """Base diagnostic for computing the sum of two fields.""" + def __init__(self, field_name1, field_name2): + """ + Args: + field_name1 (str): the name of one field to be added. + field_name2 (str): the name of the other field to be added. + """ + super().__init__(method='assign', required_fields=(field_name1, field_name2)) + self.field_name1 = field_name1 + self.field_name2 = field_name2 + + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.field_name1+"_plus_"+self.field_name2 + + 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. + """ + field1 = state_fields(self.field_name1) + field2 = state_fields(self.field_name2) + space = field1.function_space() + self.expr = field1 + field2 + super().setup(domain, state_fields, space=space) + + +class Difference(DiagnosticField): + """Base diagnostic for calculating the difference between two fields.""" + def __init__(self, field_name1, field_name2): + """ + Args: + field_name1 (str): the name of the field to be subtracted from. + field_name2 (str): the name of the field to be subtracted. + """ + super().__init__(method='assign', required_fields=(field_name1, field_name2)) + self.field_name1 = field_name1 + self.field_name2 = field_name2 + + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.field_name1+"_minus_"+self.field_name2 + + 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. + """ + + field1 = state_fields(self.field_name1) + field2 = state_fields(self.field_name2) + self.expr = field1 - field2 + space = field1.function_space() + super().setup(domain, state_fields, space=space) + + +class SteadyStateError(Difference): + """Base diagnostic for computing the steady-state error in a field.""" + def __init__(self, name): + """ + Args: + name (str): name of the field to take the steady-state error of. + """ + self.field_name1 = name + self.field_name2 = name+'_init' + DiagnosticField.__init__(self, method='assign', required_fields=(name, self.field_name2)) + + 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. + """ + # Check if initial field already exists -- otherwise needs creating + if not hasattr(state_fields, self.field_name2): + field1 = state_fields(self.field_name1) + field2 = state_fields(self.field_name2, space=field1.function_space(), + pick_up=True, dump=False) + # Attach state fields to self so that we can pick it up in compute + self.state_fields = state_fields + # The initial value for fields may not have already been set yet so we + # postpone setting it until the compute method is called + self.init_field_set = False + else: + field1 = state_fields(self.field_name1) + field2 = state_fields(self.field_name2, space=field1.function_space(), + pick_up=True, dump=False) + # By default set this new field to the current value + # This may be overwritten if picking up from a checkpoint + field2.assign(field1) + self.state_fields = state_fields + self.init_field_set = True + + super().setup(domain, state_fields) + + def compute(self): + # The first time the compute method is called we set the initial field. + # We do not want to do this if picking up from a checkpoint + if not self.init_field_set: + # Set initial field + full_field = self.state_fields(self.field_name1) + init_field = self.state_fields(self.field_name2) + init_field.assign(full_field) + + self.init_field_set = True + + super().compute() + + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.field_name1+"_error" + + +class Perturbation(Difference): + """Base diagnostic for computing perturbations from a reference profile.""" + def __init__(self, name): + """ + Args: + name (str): name of the field to take the perturbation of. + """ + self.field_name1 = name + self.field_name2 = name+'_bar' + DiagnosticField.__init__(self, method='assign', required_fields=(name, self.field_name2)) + + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.field_name1+"_perturbation" + + 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. + """ + # Check if initial field already exists -- otherwise needs creating + if not hasattr(state_fields, self.field_name2): + field1 = state_fields(self.field_name1) + _ = state_fields(self.field_name2, space=field1.function_space(), + pick_up=True, dump=False) + + super().setup(domain, state_fields) + + +class TracerDensity(DiagnosticField): + """Diagnostic for computing the density of a tracer. This is + computed as the product of a mixing ratio and dry density""" + + @property + def name(self): + """Gives the name of this diagnostic field. This records + the mixing ratio and density names, in case multiple tracer + densities are used.""" + return "TracerDensity_"+self.mixing_ratio_name+'_'+self.density_name + + def __init__(self, mixing_ratio_name, density_name, space=None, method='interpolate'): + """ + Args: + mixing_ratio_name (str): the name of the tracer mixing ratio variable + density_name (str): the name of the tracer density variable + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a new space will be constructed for this diagnostic. This + space will have enough a high enough degree to accurately compute + the product of the mixing ratio and density. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project' and + 'assign'. Defaults to 'interpolate'. + """ + super().__init__(space=space, method=method, required_fields=(mixing_ratio_name, density_name)) + + self.mixing_ratio_name = mixing_ratio_name + self.density_name = density_name + + 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. + """ + m_X = state_fields(self.mixing_ratio_name) + rho_d = state_fields(self.density_name) + self.expr = m_X*rho_d + + if self.space is None: + # Construct a space for the diagnostic that has enough + # degrees to accurately capture the tracer density. This + # will be the sum of the degrees of the individual mixing ratio + # and density function spaces. + m_X_space = m_X.function_space() + rho_d_space = rho_d.function_space() + + if domain.spaces.extruded_mesh: + # Extract the base horizontal and vertical elements + # for the mixing ratio and density. + m_X_horiz = m_X_space.ufl_element().sub_elements[0] + m_X_vert = m_X_space.ufl_element().sub_elements[1] + rho_d_horiz = rho_d_space.ufl_element().sub_elements[0] + rho_d_vert = rho_d_space.ufl_element().sub_elements[1] + + horiz_degree = m_X_horiz.degree() + rho_d_horiz.degree() + vert_degree = m_X_vert.degree() + rho_d_vert.degree() + + cell = domain.mesh._base_mesh.ufl_cell().cellname() + horiz_elt = FiniteElement('DG', cell, horiz_degree) + vert_elt = FiniteElement('DG', cell, vert_degree) + elt = TensorProductElement(horiz_elt, vert_elt) + else: + m_X_degree = m_X_space.ufl_element().degree() + rho_d_degree = rho_d_space.ufl_element().degree() + degree = m_X_degree + rho_d_degree + + cell = domain.mesh.ufl_cell().cellname() + elt = FiniteElement('DG', cell, degree) + + tracer_density_space = FunctionSpace(domain.mesh, elt, name='tracer_density_space') + super().setup(domain, state_fields, space=tracer_density_space) + + else: + super().setup(domain, state_fields) diff --git a/gusto/diagnostics/shallow_water_diagnostics.py b/gusto/diagnostics/shallow_water_diagnostics.py new file mode 100644 index 000000000..dce6fac69 --- /dev/null +++ b/gusto/diagnostics/shallow_water_diagnostics.py @@ -0,0 +1,276 @@ +"""Common diagnostic fields for the Shallow Water equations.""" + + +from firedrake import (dx, TestFunction, TrialFunction, grad, inner, curl, + LinearVariationalProblem, LinearVariationalSolver) +from gusto.diagnostics.diagnostics import DiagnosticField, Energy + +__all__ = ["ShallowWaterKineticEnergy", "ShallowWaterPotentialEnergy", + "ShallowWaterPotentialEnstrophy", "PotentialVorticity", + "RelativeVorticity", "AbsoluteVorticity"] + + +class ShallowWaterKineticEnergy(Energy): + """Diagnostic shallow-water kinetic energy density.""" + name = "ShallowWaterKineticEnergy" + + def __init__(self, space=None, method='interpolate'): + """ + 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'. + """ + super().__init__(space=space, method=method, required_fields=("D", "u")) + + 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. + """ + u = state_fields("u") + D = state_fields("D") + self.expr = self.kinetic(u, D) + super().setup(domain, state_fields) + + +class ShallowWaterPotentialEnergy(Energy): + """Diagnostic shallow-water potential energy density.""" + name = "ShallowWaterPotentialEnergy" + + def __init__(self, parameters, space=None, method='interpolate'): + """ + Args: + parameters (:class:`ShallowWaterParameters`): the configuration + object containing the physical parameters for this equation. + 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 = parameters + super().__init__(space=space, method=method, required_fields=("D")) + + 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. + """ + g = self.parameters.g + D = state_fields("D") + self.expr = 0.5*g*D**2 + super().setup(domain, state_fields) + + +class ShallowWaterPotentialEnstrophy(DiagnosticField): + """Diagnostic (dry) compressible kinetic energy density.""" + def __init__(self, base_field_name="PotentialVorticity", space=None, + method='interpolate'): + """ + Args: + base_field_name (str, optional): the base potential vorticity field + to compute the enstrophy from. Defaults to "PotentialVorticity". + 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'. + """ + base_enstrophy_names = ["PotentialVorticity", "RelativeVorticity", "AbsoluteVorticity"] + if base_field_name not in base_enstrophy_names: + raise ValueError( + f"Don't know how to compute enstrophy with base_field_name={base_field_name};" + + f"base_field_name should be one of {base_enstrophy_names}") + # Work out required fields + if base_field_name in ["PotentialVorticity", "AbsoluteVorticity"]: + required_fields = (base_field_name, "D") + elif base_field_name == "RelativeVorticity": + required_fields = (base_field_name, "D", "coriolis") + else: + raise NotImplementedError(f'Enstrophy with vorticity {base_field_name} not implemented') + + super().__init__(space=space, method=method, required_fields=required_fields) + self.base_field_name = base_field_name + + @property + def name(self): + """Gives the name of this diagnostic field.""" + base_name = "SWPotentialEnstrophy" + return "_from_".join((base_name, self.base_field_name)) + + 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. + """ + if self.base_field_name == "PotentialVorticity": + pv = state_fields("PotentialVorticity") + D = state_fields("D") + self.expr = 0.5*pv**2*D + elif self.base_field_name == "RelativeVorticity": + zeta = state_fields("RelativeVorticity") + D = state_fields("D") + f = state_fields("coriolis") + self.expr = 0.5*(zeta + f)**2/D + elif self.base_field_name == "AbsoluteVorticity": + zeta_abs = state_fields("AbsoluteVorticity") + D = state_fields("D") + self.expr = 0.5*(zeta_abs)**2/D + else: + raise NotImplementedError(f'Enstrophy with {self.base_field_name} not implemented') + super().setup(domain, state_fields) + + +class Vorticity(DiagnosticField): + """Base diagnostic field class for shallow-water vorticity variables.""" + + def setup(self, domain, state_fields, vorticity_type=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. + vorticity_type (str, optional): denotes which type of vorticity to + be computed ('relative', 'absolute' or 'potential'). Defaults to + None. + """ + + vorticity_types = ["relative", "absolute", "potential"] + if vorticity_type not in vorticity_types: + raise ValueError(f"vorticity type must be one of {vorticity_types}, not {vorticity_type}") + space = domain.spaces("H1") + + u = state_fields("u") + if vorticity_type in ["absolute", "potential"]: + f = state_fields("coriolis") + if vorticity_type == "potential": + D = state_fields("D") + + if self.method != 'solve': + if vorticity_type == "potential": + self.expr = (curl(u) + f) / D + elif vorticity_type == "absolute": + self.expr = curl(u) + f + elif vorticity_type == "relative": + self.expr = curl(u) + + super().setup(domain, state_fields, space=space) + + # Set up problem now that self.field has been set up + if self.method == 'solve': + gamma = TestFunction(space) + q = TrialFunction(space) + + if vorticity_type == "potential": + a = q*gamma*D*dx + else: + a = q*gamma*dx + + L = (- inner(domain.perp(grad(gamma)), u))*dx + if vorticity_type != "relative": + f = state_fields("coriolis") + L += gamma*f*dx + + problem = LinearVariationalProblem(a, L, self.field) + self.evaluator = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"}) + + +class PotentialVorticity(Vorticity): + u"""Diagnostic field for shallow-water potential vorticity, q=(∇×(u+f))/D""" + name = "PotentialVorticity" + + def __init__(self, space=None, method='solve'): + """ + 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 'solve'. + """ + self.solve_implemented = True + super().__init__(space=space, method=method, + required_fields=('u', 'D', 'coriolis')) + + 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. + """ + super().setup(domain, state_fields, vorticity_type="potential") + + +class AbsoluteVorticity(Vorticity): + u"""Diagnostic field for absolute vorticity, ζ=∇×(u+f)""" + name = "AbsoluteVorticity" + + def __init__(self, space=None, method='solve'): + """ + 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 'solve'. + """ + self.solve_implemented = True + super().__init__(space=space, method=method, required_fields=('u', 'coriolis')) + + 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. + """ + super().setup(domain, state_fields, vorticity_type="absolute") + + +class RelativeVorticity(Vorticity): + u"""Diagnostic field for relative vorticity, ζ=∇×u""" + name = "RelativeVorticity" + + def __init__(self, space=None, method='solve'): + """ + 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 'solve'. + """ + self.solve_implemented = True + super().__init__(space=space, method=method, required_fields=('u',)) + + 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. + """ + super().setup(domain, state_fields, vorticity_type="relative") diff --git a/gusto/equations/__init__.py b/gusto/equations/__init__.py new file mode 100644 index 000000000..d4cd2c3b4 --- /dev/null +++ b/gusto/equations/__init__.py @@ -0,0 +1,9 @@ +from gusto.equations.active_tracers import * # noqa +from gusto.equations.common_forms import * # noqa +from gusto.equations.prognostic_equations import * # noqa +from gusto.equations.transport_equations import * # noqa +from gusto.equations.diffusion_equations import * # noqa +from gusto.equations.advection_diffusion_equations import * # noqa +from gusto.equations.shallow_water_equations import * # noqa +from gusto.equations.boussinesq_equations import * # noqa +from gusto.equations.compressible_euler_equations import * # noqa diff --git a/gusto/active_tracers.py b/gusto/equations/active_tracers.py similarity index 98% rename from gusto/active_tracers.py rename to gusto/equations/active_tracers.py index cbc6a442d..d7bf2ffe6 100644 --- a/gusto/active_tracers.py +++ b/gusto/equations/active_tracers.py @@ -8,8 +8,8 @@ """ from enum import Enum -from gusto.configuration import TransportEquationType -from gusto.logging import logger +from gusto.core.configuration import TransportEquationType +from gusto.core.logging import logger __all__ = ["TracerVariableType", "Phases", "ActiveTracer", "WaterVapour", "CloudWater", "Rain"] diff --git a/gusto/equations/advection_diffusion_equations.py b/gusto/equations/advection_diffusion_equations.py new file mode 100644 index 000000000..b54172019 --- /dev/null +++ b/gusto/equations/advection_diffusion_equations.py @@ -0,0 +1,44 @@ +"""Defines the advection-diffusion equation in weak form.""" + +from firedrake import inner, dx +from firedrake.fml import subject +from gusto.core.labels import time_derivative, prognostic +from gusto.equations.common_forms import advection_form, diffusion_form +from gusto.equations.prognostic_equations import PrognosticEquation + +__all__ = ["AdvectionDiffusionEquation"] + + +class AdvectionDiffusionEquation(PrognosticEquation): + u"""The advection-diffusion equation, ∂q/∂t + (u.∇)q = ∇.(κ∇q)""" + + def __init__(self, domain, function_space, field_name, Vu=None, + diffusion_parameters=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + function_space (:class:`FunctionSpace`): the function space that the + equation's prognostic is defined on. + field_name (str): name of the prognostic field. + Vu (:class:`FunctionSpace`, optional): the function space for the + velocity field. If this is Defaults to None. + diffusion_parameters (:class:`DiffusionParameters`, optional): + parameters describing the diffusion to be applied. + """ + + super().__init__(domain, function_space, field_name) + + if Vu is not None: + domain.spaces.add_space("HDiv", Vu, overwrite_space=True) + V = domain.spaces("HDiv") + u = self.prescribed_fields("u", V) + + test = self.test + q = self.X + mass_form = time_derivative(inner(q, test)*dx) + transport_form = advection_form(test, q, u) + diffusive_form = diffusion_form(test, q, diffusion_parameters.kappa) + + self.residual = prognostic(subject( + mass_form + transport_form + diffusive_form, q), field_name) diff --git a/gusto/equations/boussinesq_equations.py b/gusto/equations/boussinesq_equations.py new file mode 100644 index 000000000..eae09369e --- /dev/null +++ b/gusto/equations/boussinesq_equations.py @@ -0,0 +1,290 @@ +"""Defines the Boussinesq equations.""" + +from firedrake import inner, dx, div, cross, split +from firedrake.fml import subject +from gusto.core.labels import ( + time_derivative, transport, 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 +) +from gusto.equations.prognostic_equations import PrognosticEquationSet + +__all__ = ["BoussinesqEquations", "LinearBoussinesqEquations"] + + +class BoussinesqEquations(PrognosticEquationSet): + """ + Class for the Boussinesq equations, which evolve the velocity + 'u', the pressure 'p' and the buoyancy 'b'. Can be compressible or + incompressible, depending on the value of the input flag, which defaults + to compressible. + + The compressible form of the equations is + ∂u/∂t + (u.∇)u + 2Ω×u + ∇p + b*k = 0, \n + ∂p/∂t + cs**2 ∇.u = p, \n + ∂b/∂t + (u.∇)b = 0, \n + where k is the vertical unit vector, Ω is the planet's rotation vector + and cs is the sound speed. + + For the incompressible form of the equations, the pressure features as + a Lagrange multiplier to enforce incompressibility. The equations are \n + ∂u/∂t + (u.∇)u + 2Ω×u + ∇p + b*k = 0, \n + ∇.u = p, \n + ∂b/∂t + (u.∇)b = 0, \n + where k is the vertical unit vector and Ω is the planet's rotation vector. + """ + + def __init__(self, domain, parameters, + compressible=True, + Omega=None, + space_names=None, + linearisation_map='default', + u_transport_option="vector_invariant_form", + no_normal_flow_bc_ids=None, + active_tracers=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. + compressible (bool, optional): flag to indicate whether the + equations are compressible. Defaults to True + Omega (:class:`ufl.Expr`, optional): an expression for the planet's + rotation vector. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + 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. + u_transport_option (str, optional): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form' and + 'circulation_form'. + Defaults to 'vector_invariant_form'. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations.. Defaults to None. + + Raises: + NotImplementedError: active tracers are not implemented. + """ + + field_names = ['u', 'p', 'b'] + + if space_names is None: + space_names = {'u': 'HDiv', 'p': 'L2', 'b': 'theta'} + + if active_tracers is not None: + raise NotImplementedError('Tracers not implemented for Boussinesq equations') + + if active_tracers is None: + active_tracers = [] + + if linearisation_map == 'default': + # Default linearisation is time derivatives and scalar transport terms + # Don't include active tracers + linearisation_map = lambda t: \ + t.get(prognostic) in ['u', 'p', 'b'] \ + and (t.has_label(time_derivative) + 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, + active_tracers=active_tracers) + + self.parameters = parameters + self.compressible = compressible + + 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) + + # -------------------------------------------------------------------- # + # Time Derivative Terms + # -------------------------------------------------------------------- # + mass_form = self.generate_mass_terms() + + # -------------------------------------------------------------------- # + # Transport Terms + # -------------------------------------------------------------------- # + # 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') + elif u_transport_option == "circulation_form": + ke_form = prognostic(kinetic_energy_form(w, u, u), 'u') + u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), 'u') + ke_form + else: + raise ValueError("Invalid u_transport_option: %s" % u_transport_option) + + # Buoyancy transport + b_adv = prognostic(advection_form(gamma, b, u), 'b') + if self.linearisation_map(b_adv.terms[0]): + linear_b_adv = linear_advection_form(gamma, b_bar, u_trial) + b_adv = linearisation(b_adv, linear_b_adv) + + if compressible: + # Pressure transport + p_adv = prognostic(advection_form(phi, p, u), 'p') + if self.linearisation_map(p_adv.terms[0]): + linear_p_adv = linear_advection_form(phi, p_bar, u_trial) + 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) + + # Add transport of tracers + if len(active_tracers) > 0: + adv_form += self.generate_tracer_transport_terms(active_tracers) + + # -------------------------------------------------------------------- # + # Pressure Gradient Term + # -------------------------------------------------------------------- # + pressure_gradient_form = pressure_gradient(subject(prognostic( + -div(w)*p*dx, 'u'), self.X)) + + # -------------------------------------------------------------------- # + # Gravitational Term + # -------------------------------------------------------------------- # + gravity_form = gravity(subject(prognostic( + -b*inner(w, domain.k)*dx, 'u'), self.X)) + + # -------------------------------------------------------------------- # + # Divergence Term + # -------------------------------------------------------------------- # + if compressible: + cs = parameters.cs + linear_div_form = divergence(subject( + prognostic(cs**2 * phi * div(u_trial) * dx, 'p'), self.X)) + divergence_form = divergence(linearisation( + subject(prognostic(cs**2 * phi * div(u) * dx, 'p'), self.X), + linear_div_form)) + else: + # This enforces that div(u) = 0 + # The p features here so that the div(u) evaluated in the + # "forcing" step replaces the whole pressure field, rather than + # merely providing an increment to it. + linear_div_form = incompressible( + subject(prognostic(phi*(p_trial-div(u_trial))*dx, 'p'), self.X)) + divergence_form = incompressible(linearisation( + subject(prognostic(phi*(p-div(u))*dx, 'p'), self.X), + linear_div_form)) + + residual = (mass_form + adv_form + divergence_form + + pressure_gradient_form + gravity_form) + + # -------------------------------------------------------------------- # + # Extra Terms (Coriolis) + # -------------------------------------------------------------------- # + if Omega is not None: + # TODO: add linearisation + residual += coriolis(subject(prognostic( + inner(w, cross(2*Omega, u))*dx, 'u'), self.X)) + + # -------------------------------------------------------------------- # + # Linearise equations + # -------------------------------------------------------------------- # + # Add linearisations to equations + self.residual = self.generate_linear_terms(residual, self.linearisation_map) + + +class LinearBoussinesqEquations(BoussinesqEquations): + """ + Class for the Boussinesq equations, which evolve the velocity + 'u', the pressure 'p' and the buoyancy 'b'. Can be compressible or + incompressible, depending on the value of the input flag, which defaults + to compressible. + + The compressible form of the equations is + ∂u/∂t + (u.∇)u + 2Ω×u + ∇p + b*k = 0, \n + ∂p/∂t + cs**2 ∇.u = p, \n + ∂b/∂t + (u.∇)b = 0, \n + where k is the vertical unit vector, Ω is the planet's rotation vector + and cs is the sound speed. + + For the incompressible form of the equations, the pressure features as + a Lagrange multiplier to enforce incompressibility. The equations are \n + ∂u/∂t + (u.∇)u + 2Ω×u + ∇p + b*k = 0, \n + ∇.u = p, \n + ∂b/∂t + (u.∇)b = 0, \n + where k is the vertical unit vector and Ω is the planet's rotation vector. + """ + + def __init__(self, domain, parameters, + compressible=True, + Omega=None, + space_names=None, + linearisation_map='default', + u_transport_option="vector_invariant_form", + no_normal_flow_bc_ids=None, + active_tracers=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. + compressible (bool, optional): flag to indicate whether the + equations are compressible. Defaults to True + Omega (:class:`ufl.Expr`, optional): an expression for the planet's + rotation vector. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + 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. + u_transport_option (str, optional): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form' and + 'circulation_form'. + Defaults to 'vector_invariant_form'. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations.. Defaults to None. + + Raises: + 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, + Omega=Omega, + space_names=space_names, + linearisation_map=linearisation_map, + u_transport_option=u_transport_option, + no_normal_flow_bc_ids=no_normal_flow_bc_ids, + active_tracers=active_tracers) + + # Use the underlying routine to do a first linearisation of + # the equations + self.linearise_equation_set() diff --git a/gusto/common_forms.py b/gusto/equations/common_forms.py similarity index 98% rename from gusto/common_forms.py rename to gusto/equations/common_forms.py index 42b1bc815..91ce01368 100644 --- a/gusto/common_forms.py +++ b/gusto/equations/common_forms.py @@ -5,9 +5,9 @@ from firedrake import (dx, dot, grad, div, inner, outer, cross, curl, split, TestFunction, TestFunctions, TrialFunctions) from firedrake.fml import subject, drop -from gusto.configuration import TransportEquationType -from gusto.labels import (transport, transporting_velocity, diffusion, - prognostic, linearisation) +from gusto.core.configuration import TransportEquationType +from gusto.core.labels import (transport, transporting_velocity, diffusion, + prognostic, linearisation) __all__ = ["advection_form", "advection_form_1d", "continuity_form", "continuity_form_1d", "vector_invariant_form", diff --git a/gusto/equations/compressible_euler_equations.py b/gusto/equations/compressible_euler_equations.py new file mode 100644 index 000000000..b6b89cd8c --- /dev/null +++ b/gusto/equations/compressible_euler_equations.py @@ -0,0 +1,374 @@ +"""Defines variants of the compressible Euler equations.""" + +from firedrake import ( + sin, pi, inner, dx, div, cross, FunctionSpace, FacetNormal, jump, avg, dS_v, + conditional, SpatialCoordinate, split, Constant +) +from firedrake.fml import subject, replace_subject +from gusto.core.labels import ( + time_derivative, transport, prognostic, hydrostatic, linearisation, + pressure_gradient, coriolis, gravity, sponge +) +from gusto.equations.thermodynamics import exner_pressure +from gusto.equations.common_forms import ( + advection_form, continuity_form, vector_invariant_form, + kinetic_energy_form, advection_equation_circulation_form, + diffusion_form, linear_continuity_form, linear_advection_form +) +from gusto.equations.active_tracers import Phases, TracerVariableType +from gusto.equations.prognostic_equations import PrognosticEquationSet + +__all__ = ["CompressibleEulerEquations", "HydrostaticCompressibleEulerEquations"] + + +class CompressibleEulerEquations(PrognosticEquationSet): + """ + Class for the compressible Euler equations, which evolve the velocity 'u', + the dry density 'rho' and the (virtual dry) potential temperature 'theta', + solving: \n + ∂u/∂t + (u.∇)u + 2Ω×u + c_p*θ*∇Π + g = 0, \n + ∂ρ/∂t + ∇.(ρ*u) = 0, \n + ∂θ/∂t + (u.∇)θ = 0, \n + where Π is the Exner pressure, g is the gravitational vector, Ω is the + planet's rotation vector and c_p is the heat capacity of dry air at constant + pressure. + """ + + def __init__(self, domain, parameters, Omega=None, sponge_options=None, + extra_terms=None, space_names=None, + linearisation_map='default', + u_transport_option="vector_invariant_form", + diffusion_options=None, + no_normal_flow_bc_ids=None, + active_tracers=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. + Omega (:class:`ufl.Expr`, optional): an expression for the planet's + rotation vector. Defaults to None. + sponge_options (:class:`SpongeLayerParameters`, optional): any + parameters for applying a sponge layer to the upper boundary. + Defaults to None. + extra_terms (:class:`ufl.Expr`, optional): any extra terms to be + included in the equation set. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + 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. + u_transport_option (str, optional): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form' and + 'circulation_form'. + Defaults to 'vector_invariant_form'. + diffusion_options (:class:`DiffusionParameters`, optional): any + options to specify for applying diffusion terms to variables. + Defaults to None. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations.. Defaults to None. + + Raises: + NotImplementedError: only mixing ratio tracers are implemented. + """ + + field_names = ['u', 'rho', 'theta'] + + if space_names is None: + space_names = {'u': 'HDiv', 'rho': 'L2', 'theta': 'theta'} + + if active_tracers is None: + active_tracers = [] + + if linearisation_map == 'default': + # Default linearisation is time derivatives and scalar transport terms + # Don't include active tracers + linearisation_map = lambda t: \ + t.get(prognostic) in ['u', 'rho', 'theta'] \ + and (t.has_label(time_derivative) + or (t.get(prognostic) != 'u' 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, + active_tracers=active_tracers) + + self.parameters = parameters + g = parameters.g + cp = parameters.cp + + w, phi, gamma = self.tests[0:3] + u, rho, theta = split(self.X)[0:3] + u_trial = split(self.trials)[0] + _, rho_bar, theta_bar = split(self.X_ref)[0:3] + zero_expr = Constant(0.0)*theta + exner = exner_pressure(parameters, rho, theta) + n = FacetNormal(domain.mesh) + + # -------------------------------------------------------------------- # + # Time Derivative Terms + # -------------------------------------------------------------------- # + mass_form = self.generate_mass_terms() + + # -------------------------------------------------------------------- # + # Transport Terms + # -------------------------------------------------------------------- # + # 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') + elif u_transport_option == "circulation_form": + ke_form = prognostic(kinetic_energy_form(w, u, u), 'u') + u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), 'u') + ke_form + else: + raise ValueError("Invalid u_transport_option: %s" % u_transport_option) + + # Density transport (conservative form) + rho_adv = prognostic(continuity_form(phi, rho, u), 'rho') + # Transport term needs special linearisation + if self.linearisation_map(rho_adv.terms[0]): + linear_rho_adv = linear_continuity_form(phi, rho_bar, u_trial) + rho_adv = linearisation(rho_adv, linear_rho_adv) + + # Potential temperature transport (advective form) + theta_adv = prognostic(advection_form(gamma, theta, u), 'theta') + # Transport term needs special linearisation + if self.linearisation_map(theta_adv.terms[0]): + linear_theta_adv = linear_advection_form(gamma, theta_bar, u_trial) + theta_adv = linearisation(theta_adv, linear_theta_adv) + + adv_form = subject(u_adv + rho_adv + theta_adv, self.X) + + # Add transport of tracers + if len(active_tracers) > 0: + adv_form += self.generate_tracer_transport_terms(active_tracers) + + # -------------------------------------------------------------------- # + # Pressure Gradient Term + # -------------------------------------------------------------------- # + # First get total mass of tracers + tracer_mr_total = zero_expr + for tracer in active_tracers: + if tracer.variable_type == TracerVariableType.mixing_ratio: + idx = self.field_names.index(tracer.name) + tracer_mr_total += split(self.X)[idx] + else: + raise NotImplementedError('Only mixing ratio tracers are implemented') + theta_v = theta / (Constant(1.0) + tracer_mr_total) + + pressure_gradient_form = pressure_gradient(subject(prognostic( + cp*(-div(theta_v*w)*exner*dx + + jump(theta_v*w, n)*avg(exner)*dS_v), 'u'), self.X)) + + # -------------------------------------------------------------------- # + # Gravitational Term + # -------------------------------------------------------------------- # + gravity_form = gravity(subject(prognostic(g*inner(domain.k, w)*dx, + 'u'), self.X)) + + residual = (mass_form + adv_form + pressure_gradient_form + + gravity_form) + + # -------------------------------------------------------------------- # + # Moist Thermodynamic Divergence Term + # -------------------------------------------------------------------- # + if len(active_tracers) > 0: + cv = parameters.cv + c_vv = parameters.c_vv + c_pv = parameters.c_pv + c_pl = parameters.c_pl + R_d = parameters.R_d + R_v = parameters.R_v + + # Get gas and liquid moisture mixing ratios + mr_l = zero_expr + mr_v = zero_expr + + for tracer in active_tracers: + if tracer.chemical == 'H2O': + if tracer.variable_type == TracerVariableType.mixing_ratio: + idx = self.field_names.index(tracer.name) + if tracer.phase == Phases.gas: + mr_v += split(self.X)[idx] + elif tracer.phase == Phases.liquid: + mr_l += split(self.X)[idx] + else: + raise NotImplementedError('Only mixing ratio tracers are implemented') + + c_vml = cv + mr_v * c_vv + mr_l * c_pl + c_pml = cp + mr_v * c_pv + mr_l * c_pl + R_m = R_d + mr_v * R_v + + residual += subject(prognostic( + gamma * theta * div(u) + * (R_m / c_vml - (R_d * c_pml) / (cp * c_vml))*dx, 'theta'), self.X) + + # -------------------------------------------------------------------- # + # Extra Terms (Coriolis, Sponge, Diffusion and others) + # -------------------------------------------------------------------- # + if Omega is not None: + # TODO: add linearisation + residual += coriolis(subject(prognostic( + inner(w, cross(2*Omega, u))*dx, "u"), self.X)) + + if sponge_options is not None: + W_DG = FunctionSpace(domain.mesh, "DG", 2) + x = SpatialCoordinate(domain.mesh) + z = x[len(x)-1] + H = sponge_options.H + zc = sponge_options.z_level + assert float(zc) < float(H), \ + "The sponge level is set above the height the your domain" + mubar = sponge_options.mubar + muexpr = conditional(z <= zc, + 0.0, + mubar*sin((pi/2.)*(z-zc)/(H-zc))**2) + self.mu = self.prescribed_fields("sponge", W_DG).interpolate(muexpr) + + residual += sponge(subject(prognostic( + self.mu*inner(w, domain.k)*inner(u, domain.k)*dx, 'u'), self.X)) + + if diffusion_options is not None: + for field, diffusion in diffusion_options: + idx = self.field_names.index(field) + test = self.tests[idx] + fn = split(self.X)[idx] + residual += subject( + prognostic(diffusion_form(test, fn, diffusion.kappa), field), + self.X) + + if extra_terms is not None: + for field, term in extra_terms: + idx = self.field_names.index(field) + test = self.tests[idx] + residual += subject(prognostic( + inner(test, term)*dx, field), self.X) + + # -------------------------------------------------------------------- # + # Linearise equations + # -------------------------------------------------------------------- # + # Add linearisations to equations + self.residual = self.generate_linear_terms(residual, self.linearisation_map) + + +class HydrostaticCompressibleEulerEquations(CompressibleEulerEquations): + """ + The hydrostatic form of the compressible Euler equations. In this case the + vertical velocity derivative is zero in the equations, so only 'u_h', the + horizontal component of the velocity is allowed to vary in time. The + equations, for velocity 'u', dry density 'rho' and (dry) potential + temperature 'theta' are: \n + ∂u_h/∂t + (u.∇)u_h + 2Ω×u + c_p*θ*∇Π + g = 0, \n + ∂ρ/∂t + ∇.(ρ*u) = 0, \n + ∂θ/∂t + (u.∇)θ = 0, \n + where Π is the Exner pressure, g is the gravitational vector, Ω is the + planet's rotation vector and c_p is the heat capacity of dry air at constant + pressure. + + This is implemented through a hydrostatic switch to the compressible Euler + equations. + """ + + def __init__(self, domain, parameters, Omega=None, sponge=None, + extra_terms=None, space_names=None, linearisation_map='default', + u_transport_option="vector_invariant_form", + diffusion_options=None, + no_normal_flow_bc_ids=None, + active_tracers=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. + Omega (:class:`ufl.Expr`, optional): an expression for the planet's + rotation vector. Defaults to None. + sponge (:class:`ufl.Expr`, optional): an expression for a sponge + layer. Defaults to None. + extra_terms (:class:`ufl.Expr`, optional): any extra terms to be + included in the equation set. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + 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. + u_transport_option (str, optional): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form' and + 'circulation_form'. + Defaults to 'vector_invariant_form'. + diffusion_options (:class:`DiffusionOptions`, optional): any options + to specify for applying diffusion terms to variables. Defaults + to None. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations.. Defaults to None. + + Raises: + NotImplementedError: only mixing ratio tracers are implemented. + """ + + super().__init__(domain, parameters, Omega=Omega, sponge=sponge, + extra_terms=extra_terms, space_names=space_names, + linearisation_map=linearisation_map, + u_transport_option=u_transport_option, + diffusion_options=diffusion_options, + no_normal_flow_bc_ids=no_normal_flow_bc_ids, + active_tracers=active_tracers) + + self.residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=lambda t: hydrostatic(t, self.hydrostatic_projection(t)) + ) + + k = self.domain.k + u = split(self.X)[0] + self.residual += hydrostatic( + subject( + prognostic( + -inner(k, self.tests[0]) * inner(k, u) * dx, "u"), + self.X)) + + def hydrostatic_projection(self, t): + """ + Performs the 'hydrostatic' projection. + + Takes a term involving a vector prognostic variable and replaces the + prognostic with only its horizontal components. + + Args: + t (:class:`Term`): the term to perform the projection upon. + + Returns: + :class:`LabelledForm`: the labelled form containing the new term. + + Raises: + AssertionError: spherical geometry is not yet implemented. + """ + + # TODO: make this more general, i.e. should work on the sphere + if self.domain.on_sphere: + raise NotImplementedError("The hydrostatic projection is not yet " + + "implemented for spherical geometry") + k = Constant((*self.domain.k, 0, 0)) + X = t.get(subject) + + new_subj = X - k * inner(X, k) + return replace_subject(new_subj)(t) diff --git a/gusto/equations/diffusion_equations.py b/gusto/equations/diffusion_equations.py new file mode 100644 index 000000000..0e4ec07da --- /dev/null +++ b/gusto/equations/diffusion_equations.py @@ -0,0 +1,34 @@ +"""Defines the diffusion equation in weak form.""" + +from firedrake import inner, dx +from firedrake.fml import subject +from gusto.core.labels import time_derivative, prognostic +from gusto.equations.common_forms import diffusion_form +from gusto.equations.prognostic_equations import PrognosticEquation + +__all__ = ["DiffusionEquation"] + + +class DiffusionEquation(PrognosticEquation): + u"""Discretises the diffusion equation, ∂q/∂t = ∇.(κ∇q)""" + + def __init__(self, domain, function_space, field_name, + diffusion_parameters): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + function_space (:class:`FunctionSpace`): the function space that the + equation's prognostic is defined on. + field_name (str): name of the prognostic field. + diffusion_parameters (:class:`DiffusionParameters`): parameters + describing the diffusion to be applied. + """ + super().__init__(domain, function_space, field_name) + + test = self.test + q = self.X + mass_form = time_derivative(inner(q, test)*dx) + diffusive_form = diffusion_form(test, q, diffusion_parameters.kappa) + + self.residual = prognostic(subject(mass_form + diffusive_form, q), field_name) diff --git a/gusto/equations.py b/gusto/equations/equation.py similarity index 99% rename from gusto/equations.py rename to gusto/equations/equation.py index c07aed867..1410b0802 100644 --- a/gusto/equations.py +++ b/gusto/equations/equation.py @@ -11,13 +11,13 @@ Term, all_terms, keep, drop, Label, subject, replace_subject, replace_trial_function ) -from gusto.fields import PrescribedFields -from gusto.labels import ( +from gusto.core import PrescribedFields +from gusto.core.labels import ( time_derivative, transport, prognostic, hydrostatic, linearisation, pressure_gradient, coriolis, divergence, gravity, incompressible, sponge ) -from gusto.thermodynamics import exner_pressure -from gusto.common_forms import ( +from gusto.equations.thermodynamics import exner_pressure +from gusto.equations.common_forms import ( advection_form, advection_form_1d, continuity_form, continuity_form_1d, vector_invariant_form, kinetic_energy_form, advection_equation_circulation_form, @@ -25,8 +25,8 @@ linear_continuity_form, linear_continuity_form_1d, linear_advection_form, tracer_conservative_form ) -from gusto.active_tracers import ActiveTracer, Phases, TracerVariableType -from gusto.configuration import TransportEquationType +from gusto.equations.active_tracers import ActiveTracer, Phases, TracerVariableType +from gusto.core.configuration import TransportEquationType import ufl diff --git a/gusto/equations/prognostic_equations.py b/gusto/equations/prognostic_equations.py new file mode 100644 index 000000000..9f88bf31c --- /dev/null +++ b/gusto/equations/prognostic_equations.py @@ -0,0 +1,428 @@ +"""Objects describing geophysical fluid equations to be solved in weak form.""" + +from abc import ABCMeta +from firedrake import ( + TestFunction, Function, inner, dx, MixedFunctionSpace, TestFunctions, + TrialFunction, DirichletBC, split, action +) +from firedrake.fml import ( + Term, all_terms, keep, drop, Label, subject, + replace_subject, replace_trial_function +) +from gusto.core import PrescribedFields +from gusto.core.labels import time_derivative, prognostic, linearisation +from gusto.equations.common_forms import ( + advection_form, continuity_form, tracer_conservative_form +) +from gusto.equations.active_tracers import ActiveTracer +from gusto.core.configuration import TransportEquationType +import ufl + +__all__ = ["PrognosticEquation", "PrognosticEquationSet"] + + +class PrognosticEquation(object, metaclass=ABCMeta): + """Base class for prognostic equations.""" + + def __init__(self, domain, function_space, field_name): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + function_space (:class:`FunctionSpace`): the function space that the + equation's prognostic is defined on. + field_name (str): name of the prognostic field. + """ + + self.domain = domain + self.function_space = function_space + self.X = Function(function_space) + self.field_name = field_name + self.bcs = {} + self.prescribed_fields = PrescribedFields() + + if len(function_space) > 1: + assert hasattr(self, "field_names") + for fname in self.field_names: + self.bcs[fname] = [] + else: + # To avoid confusion, only add "self.test" when not mixed FS + self.test = TestFunction(function_space) + + self.bcs[field_name] = [] + + def label_terms(self, term_filter, label): + """ + Labels terms in the equation, subject to the term filter. + + + Args: + term_filter (func): a function, taking terms as an argument, that + is used to filter terms. + label (:class:`Label`): the label to be applied to the terms. + """ + assert type(label) == Label + self.residual = self.residual.label_map(term_filter, map_if_true=label) + + +class PrognosticEquationSet(PrognosticEquation, metaclass=ABCMeta): + """ + Base class for solving a set of prognostic equations. + + A prognostic equation set contains multiple prognostic variables, which are + solved for simultaneously in a :class:`MixedFunctionSpace`. This base class + contains common routines for these equation sets. + """ + + def __init__(self, field_names, domain, space_names, + linearisation_map=None, no_normal_flow_bc_ids=None, + active_tracers=None): + """ + Args: + field_names (list): a list of strings for names of the prognostic + variables for the equation set. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + space_names (dict): a dictionary of strings for names of the + function spaces to use for the spatial discretisation. The keys + are the names of the prognostic variables. + linearisation_map (func, optional): a function specifying which + terms in the equation set to linearise. Defaults to None. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations.. Defaults to None. + """ + + self.field_names = field_names + self.space_names = space_names + self.active_tracers = active_tracers + self.linearisation_map = lambda t: False if linearisation_map is None else linearisation_map(t) + + # Build finite element spaces + self.spaces = [domain.spaces(space_name) for space_name in + [self.space_names[field_name] for field_name in self.field_names]] + + # Add active tracers to the list of prognostics + if active_tracers is None: + active_tracers = [] + self.add_tracers_to_prognostics(domain, active_tracers) + + # Make the full mixed function space + W = MixedFunctionSpace(self.spaces) + + # Can now call the underlying PrognosticEquation + full_field_name = "_".join(self.field_names) + super().__init__(domain, W, full_field_name) + + # Set up test functions, trials and prognostics + self.tests = TestFunctions(W) + self.trials = TrialFunction(W) + self.X_ref = Function(W) + + # Set up no-normal-flow boundary conditions + if no_normal_flow_bc_ids is None: + no_normal_flow_bc_ids = [] + self.set_no_normal_flow_bcs(domain, no_normal_flow_bc_ids) + + # ======================================================================== # + # Set up time derivative / mass terms + # ======================================================================== # + + def generate_mass_terms(self): + """ + Builds the weak time derivative terms for the equation set. + + Generates the weak time derivative terms ("mass terms") for all the + prognostic variables of the equation set. + + Returns: + :class:`LabelledForm`: a labelled form containing the mass terms. + """ + + for i, (test, field_name) in enumerate(zip(self.tests, self.field_names)): + prog = split(self.X)[i] + mass = subject(prognostic(inner(prog, test)*dx, field_name), self.X) + if i == 0: + mass_form = time_derivative(mass) + else: + mass_form += time_derivative(mass) + + return mass_form + + # ======================================================================== # + # Linearisation Routines + # ======================================================================== # + + def generate_linear_terms(self, residual, linearisation_map): + """ + Generate the linearised forms for the equation set. + + Generates linear forms for each of the terms in the equation set + (unless specified otherwise). The linear forms are then added to the + terms through a `linearisation` :class:`Label`. + + Linear forms are generated by replacing the `subject` using the + `ufl.derivative` to obtain the forms linearised around reference states. + + Terms that already have a `linearisation` label are left. + + Args: + residual (:class:`LabelledForm`): the residual of the equation set. + A labelled form containing all the terms of the equation set. + linearisation_map (func): a function describing the terms to be + linearised. + + Returns: + :class:`LabelledForm`: the residual with linear terms attached to + each term as labels. + """ + + from functools import partial + + # Function to check if term should be linearised + def should_linearise(term): + return (not term.has_label(linearisation) and linearisation_map(term)) + + # Linearise a term, and add the linearisation as a label + def linearise(term, X, X_ref, du): + linear_term = Term(action(ufl.derivative(term.form, X), du), term.labels) + return linearisation(term, replace_subject(X_ref)(linear_term)) + + # Add linearisations to all terms that need linearising + residual = residual.label_map( + should_linearise, + map_if_true=partial(linearise, X=self.X, X_ref=self.X_ref, du=self.trials), + map_if_false=keep, + ) + + return residual + + def linearise_equation_set(self): + """ + Linearises the equation set. + + Linearises the whole equation set, replacing all the equations with + the complete linearisation. Terms without linearisations are dropped. + All labels are carried over, and the original linearisations containing + the trial function are kept as labels to the new terms. + """ + + # Replace all terms with their linearisations, drop terms without + self.residual = self.residual.label_map( + lambda t: t.has_label(linearisation), + map_if_true=lambda t: Term(t.get(linearisation).form, t.labels), + map_if_false=drop) + + # Replace trial functions with the prognostics + self.residual = self.residual.label_map( + all_terms, replace_trial_function(self.X)) + + # ======================================================================== # + # Boundary Condition Routines + # ======================================================================== # + + def set_no_normal_flow_bcs(self, domain, no_normal_flow_bc_ids): + """ + Sets up the boundary conditions for no-normal flow at domain boundaries. + + Sets up the no-normal-flow boundary conditions, storing the + :class:`DirichletBC` object at each specified boundary. There must be + a velocity variable named 'u' to apply the boundary conditions to. + + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + no_normal_flow_bc_ids (list): A list of IDs of the domain boundaries + at which no normal flow will be enforced. + + Raises: + NotImplementedError: if there is no velocity field (with name 'u') + in the equation set. + """ + + if 'u' not in self.field_names: + raise NotImplementedError( + 'No-normal-flow boundary conditions can only be applied ' + + 'when there is a variable called "u" and none was found') + + Vu = domain.spaces("HDiv") + # we only apply no normal-flow BCs when extruded mesh is non periodic + if Vu.extruded and not Vu.ufl_domain().topology.extruded_periodic: + self.bcs['u'].append(DirichletBC(Vu, 0.0, "bottom")) + self.bcs['u'].append(DirichletBC(Vu, 0.0, "top")) + for id in no_normal_flow_bc_ids: + self.bcs['u'].append(DirichletBC(Vu, 0.0, id)) + + # Add all boundary conditions to mixed function space + W = self.X.function_space() + self.bcs[self.field_name] = [] + for idx, field_name in enumerate(self.field_names): + for bc in self.bcs[field_name]: + self.bcs[self.field_name].append(DirichletBC(W.sub(idx), bc.function_arg, bc.sub_domain)) + + # ======================================================================== # + # Active Tracer Routines + # ======================================================================== # + + def add_tracers_to_prognostics(self, domain, active_tracers): + """ + Augments the equation set with specified active tracer variables. + + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + active_tracers (list): A list of :class:`ActiveTracer` objects that + encode the metadata for the active tracers. + + Raises: + ValueError: the equation set already contains a variable with the + name of the active tracer. + """ + + # Loop through tracer fields and add field names and spaces + for tracer in active_tracers: + if isinstance(tracer, ActiveTracer): + if tracer.name not in self.field_names: + self.field_names.append(tracer.name) + else: + raise ValueError(f'There is already a field named {tracer.name}') + + # Add name of space to self.space_names, but check for conflict + # with the tracer's name + if tracer.name in self.space_names: + assert self.space_names[tracer.name] == tracer.space, \ + 'space_name dict provided to equation has space ' \ + + f'{self.space_names[tracer.name]} for tracer ' \ + + f'{tracer.name} which conflicts with the space ' \ + + f'{tracer.space} specified in the ActiveTracer object' + else: + self.space_names[tracer.name] = tracer.space + self.spaces.append(domain.spaces(tracer.space)) + else: + raise TypeError(f'Tracers must be ActiveTracer objects, not {type(tracer)}') + + def generate_tracer_mass_terms(self, active_tracers): + """ + Adds the mass forms for the active tracers to the equation set. + + Args: + active_tracers (list): A list of :class:`ActiveTracer` objects that + encode the metadata for the active tracers. + + Returns: + :class:`LabelledForm`: a labelled form containing the mass + terms for the active tracers. This is the usual mass form + unless using tracer_conservative, where it is multiplied + by the reference density. + """ + + for i, tracer in enumerate(active_tracers): + idx = self.field_names.index(tracer.name) + tracer_prog = split(self.X)[idx] + tracer_test = self.tests[idx] + + if tracer.transport_eqn == TransportEquationType.tracer_conservative: + ref_density_idx = self.field_names.index(tracer.density_name) + ref_density = split(self.X)[ref_density_idx] + q = tracer_prog*ref_density + mass = subject(prognostic(inner(q, tracer_test)*dx, + self.field_names[idx]), self.X) + else: + mass = subject(prognostic(inner(tracer_prog, tracer_test)*dx, + self.field_names[idx]), self.X) + + if i == 0: + mass_form = time_derivative(mass) + else: + mass_form += time_derivative(mass) + + return mass_form + + def generate_tracer_transport_terms(self, active_tracers): + """ + Adds the transport forms for the active tracers to the equation set. + + Args: + active_tracers (list): A list of :class:`ActiveTracer` objects that + encode the metadata for the active tracers. + + Raises: + ValueError: if the transport equation encoded in the active tracer + metadata is not valid. + + Returns: + :class:`LabelledForm`: a labelled form containing the transport + terms for the active tracers. + """ + + # By default return None if no tracers are to be transported + adv_form = None + no_tracer_transported = True + + if 'u' in self.field_names: + u_idx = self.field_names.index('u') + u = split(self.X)[u_idx] + elif 'u' in self.prescribed_fields._field_names: + u = self.prescribed_fields('u') + else: + raise ValueError('Cannot generate tracer transport terms ' + + 'as there is no velocity field') + + for _, tracer in enumerate(active_tracers): + if tracer.transport_eqn != TransportEquationType.no_transport: + idx = self.field_names.index(tracer.name) + tracer_prog = split(self.X)[idx] + tracer_test = self.tests[idx] + if tracer.transport_eqn == TransportEquationType.advective: + tracer_adv = prognostic( + advection_form(tracer_test, tracer_prog, u), + tracer.name) + elif tracer.transport_eqn == TransportEquationType.conservative: + tracer_adv = prognostic( + continuity_form(tracer_test, tracer_prog, u), + tracer.name) + elif tracer.transport_eqn == TransportEquationType.tracer_conservative: + ref_density_idx = self.field_names.index(tracer.density_name) + ref_density = split(self.X)[ref_density_idx] + tracer_adv = prognostic( + tracer_conservative_form(tracer_test, tracer_prog, + ref_density, u), tracer.name) + + else: + raise ValueError(f'Transport eqn {tracer.transport_eqn} not recognised') + + if no_tracer_transported: + # We arrive here for the first tracer to be transported + adv_form = subject(tracer_adv, self.X) + no_tracer_transported = False + else: + adv_form += subject(tracer_adv, self.X) + + return adv_form + + def get_active_tracer(self, field_name): + """ + Returns the active tracer metadata object for a particular field. + + Args: + field_name (str): the name of the field to return the metadata for. + + Returns: + :class:`ActiveTracer`: the object storing the metadata describing + the tracer. + """ + + active_tracer_to_return = None + + for active_tracer in self.active_tracers: + if active_tracer.name == field_name: + active_tracer_to_return = active_tracer + break + + if active_tracer_to_return is None: + raise RuntimeError(f'Unable to find active tracer {field_name}') + + return active_tracer_to_return diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py new file mode 100644 index 000000000..7e7453656 --- /dev/null +++ b/gusto/equations/shallow_water_equations.py @@ -0,0 +1,457 @@ +"""Classes for defining variants of the shallow-water equations.""" + +from firedrake import (inner, dx, div, FunctionSpace, FacetNormal, jump, avg, + dS, split) +from firedrake.fml import subject +from gusto.core.labels import (time_derivative, transport, prognostic, + linearisation, pressure_gradient, coriolis) +from gusto.equations.common_forms import ( + advection_form, advection_form_1d, continuity_form, + continuity_form_1d, vector_invariant_form, + kinetic_energy_form, advection_equation_circulation_form, diffusion_form_1d, + linear_continuity_form, linear_continuity_form_1d +) +from gusto.equations.prognostic_equations import PrognosticEquationSet + +__all__ = ["ShallowWaterEquations", "LinearShallowWaterEquations", + "ShallowWaterEquations_1d", "LinearShallowWaterEquations_1d"] + + +class ShallowWaterEquations(PrognosticEquationSet): + u""" + Class for the (rotating) shallow-water equations, which evolve the velocity + 'u' and the depth field 'D', via some variant of: \n + ∂u/∂t + (u.∇)u + f×u + g*∇(D+b) = 0, \n + ∂D/∂t + ∇.(D*u) = 0, \n + for Coriolis parameter 'f' and bottom surface 'b'. + """ + + def __init__(self, domain, parameters, fexpr=None, bexpr=None, + space_names=None, linearisation_map='default', + u_transport_option='vector_invariant_form', + no_normal_flow_bc_ids=None, active_tracers=None, + thermal=False): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. + fexpr (:class:`ufl.Expr`, optional): an expression for the Coroilis + parameter. Defaults to None. + bexpr (:class:`ufl.Expr`, optional): an expression for the bottom + surface of the fluid. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + in which case the spaces are taken from the de Rham complex. Any + buoyancy variable is taken by default to lie in the L2 space. + 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 both time derivatives, + the 'D' transport term and the pressure gradient term. + u_transport_option (str, optional): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form', and + 'circulation_form'. + Defaults to 'vector_invariant_form'. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations. Defaults to None. + thermal (flag, optional): specifies whether the equations have a + thermal or buoyancy variable that feeds back on the momentum. + Defaults to False. + + Raises: + NotImplementedError: active tracers are not yet implemented. + """ + + self.thermal = thermal + field_names = ['u', 'D'] + + if space_names is None: + space_names = {'u': 'HDiv', 'D': 'L2'} + + if active_tracers is None: + active_tracers = [] + + if self.thermal: + field_names.append('b') + if 'b' not in space_names.keys(): + space_names['b'] = 'L2' + + if linearisation_map == 'default': + # Default linearisation is time derivatives, pressure gradient and + # transport term from depth equation. Don't include active tracers + linearisation_map = lambda t: \ + t.get(prognostic) in ['u', 'D', 'b'] \ + and (any(t.has_label(time_derivative, pressure_gradient)) + or (t.get(prognostic) in ['D', 'b'] 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, + active_tracers=active_tracers) + + self.parameters = parameters + g = parameters.g + H = parameters.H + + w, phi = self.tests[0:2] + u, D = split(self.X)[0:2] + u_trial = split(self.trials)[0] + + # -------------------------------------------------------------------- # + # Time Derivative Terms + # -------------------------------------------------------------------- # + mass_form = self.generate_mass_terms() + + # -------------------------------------------------------------------- # + # Transport Terms + # -------------------------------------------------------------------- # + # 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') + elif u_transport_option == "circulation_form": + ke_form = prognostic(kinetic_energy_form(w, u, u), 'u') + u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), 'u') + ke_form + else: + raise ValueError("Invalid u_transport_option: %s" % u_transport_option) + + # Depth transport term + D_adv = prognostic(continuity_form(phi, D, u), 'D') + + # Transport term needs special linearisation + if self.linearisation_map(D_adv.terms[0]): + linear_D_adv = linear_continuity_form(phi, H, u_trial) + # Add linearisation to D_adv + D_adv = linearisation(D_adv, linear_D_adv) + + adv_form = subject(u_adv + D_adv, self.X) + + # Add transport of tracers + if len(active_tracers) > 0: + adv_form += self.generate_tracer_transport_terms(active_tracers) + # Add transport of buoyancy, if thermal shallow water equations + if self.thermal: + gamma = self.tests[2] + b = split(self.X)[2] + b_adv = prognostic(advection_form(gamma, b, u), 'b') + adv_form += subject(b_adv, self.X) + + # -------------------------------------------------------------------- # + # Pressure Gradient Term + # -------------------------------------------------------------------- # + # Add pressure gradient only if not doing thermal + if self.thermal: + residual = (mass_form + adv_form) + else: + pressure_gradient_form = pressure_gradient( + subject(prognostic(-g*div(w)*D*dx, 'u'), self.X)) + + residual = (mass_form + adv_form + pressure_gradient_form) + + # -------------------------------------------------------------------- # + # Extra Terms (Coriolis, Topography and Thermal) + # -------------------------------------------------------------------- # + # TODO: Is there a better way to store the Coriolis / topography fields? + # The current approach is that these are prescribed fields, stored in + # the equation, and initialised when the equation is + + if fexpr is not None: + V = FunctionSpace(domain.mesh, 'CG', 1) + f = self.prescribed_fields('coriolis', V).interpolate(fexpr) + coriolis_form = coriolis(subject( + prognostic(f*inner(domain.perp(u), w)*dx, "u"), self.X)) + # Add linearisation + if self.linearisation_map(coriolis_form.terms[0]): + linear_coriolis = coriolis( + subject(prognostic(f*inner(domain.perp(u_trial), w)*dx, 'u'), self.X)) + coriolis_form = linearisation(coriolis_form, linear_coriolis) + residual += coriolis_form + + if bexpr is not None: + topography = self.prescribed_fields('topography', domain.spaces('DG')).interpolate(bexpr) + if self.thermal: + n = FacetNormal(domain.mesh) + topography_form = subject(prognostic + (-topography*div(b*w)*dx + + jump(b*w, n)*avg(topography)*dS, + 'u'), self.X) + else: + topography_form = subject(prognostic + (-g*div(w)*topography*dx, 'u'), + self.X) + residual += topography_form + + # thermal source terms not involving topography + if self.thermal: + n = FacetNormal(domain.mesh) + source_form = subject(prognostic(-D*div(b*w)*dx + - 0.5*b*div(D*w)*dx + + jump(b*w, n)*avg(D)*dS + + 0.5*jump(D*w, n)*avg(b)*dS, + 'u'), self.X) + residual += source_form + + # -------------------------------------------------------------------- # + # Linearise equations + # -------------------------------------------------------------------- # + # Add linearisations to equations + self.residual = self.generate_linear_terms(residual, self.linearisation_map) + + +class LinearShallowWaterEquations(ShallowWaterEquations): + u""" + Class for the linear (rotating) shallow-water equations, which describe the + velocity 'u' and the depth field 'D', solving some variant of: \n + ∂u/∂t + f×u + g*∇(D+b) = 0, \n + ∂D/∂t + H*∇.(u) = 0, \n + for mean depth 'H', Coriolis parameter 'f' and bottom surface 'b'. + + This is set up the from the underlying :class:`ShallowWaterEquations`, + which is then linearised. + """ + + def __init__(self, domain, parameters, fexpr=None, bexpr=None, + space_names=None, linearisation_map='default', + u_transport_option="vector_invariant_form", + no_normal_flow_bc_ids=None, active_tracers=None, + thermal=False): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. + fexpr (:class:`ufl.Expr`, optional): an expression for the Coroilis + parameter. Defaults to None. + bexpr (:class:`ufl.Expr`, optional): an expression for the bottom + surface of the fluid. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + in which case the spaces are taken from the de Rham complex. Any + buoyancy variable is taken by default to lie in the L2 space. + 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 both time derivatives, + the 'D' transport term, pressure gradient and Coriolis terms. + u_transport_option (str, optional): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form' and + 'circulation_form'. + Defaults to 'vector_invariant_form'. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations. Defaults to None. + thermal (flag, optional): specifies whether the equations have a + thermal or buoyancy variable that feeds back on the momentum. + Defaults to False. + """ + + 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)) + or (t.get(prognostic) in ['D', 'b'] and t.has_label(transport))) + + super().__init__(domain, parameters, + fexpr=fexpr, bexpr=bexpr, space_names=space_names, + linearisation_map=linearisation_map, + u_transport_option=u_transport_option, + no_normal_flow_bc_ids=no_normal_flow_bc_ids, + active_tracers=active_tracers, thermal=thermal) + + # Use the underlying routine to do a first linearisation of the equations + self.linearise_equation_set() + + +class ShallowWaterEquations_1d(PrognosticEquationSet): + + u""" + Class for the (rotating) 1D shallow-water equations, which describe + the velocity 'u', 'v' and the depth field 'D', solving some variant of: \n + ∂u/∂t + u∂u/∂x - fv + g*∂D/∂x = 0, \n + ∂v/∂t + fu = 0, \n + ∂D/∂t + ∂(uD)/∂x = 0, \n + for mean depth 'H', Coriolis parameter 'f' and gravity 'g'. + + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. + fexpr (:class:`ufl.Expr`, optional): an expression for the Coroilis + parameter. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + 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 both time derivatives, + the 'D' transport term, pressure gradient and Coriolis terms. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations. Defaults to None. + """ + + def __init__(self, domain, parameters, + fexpr=None, + space_names=None, linearisation_map='default', + diffusion_options=None, + no_normal_flow_bc_ids=None, active_tracers=None): + + field_names = ['u', 'v', 'D'] + space_names = {'u': 'HDiv', 'v': 'L2', 'D': 'L2'} + + if active_tracers is not None: + raise NotImplementedError('Tracers not implemented for 1D shallow water equations') + + 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)) + or (t.get(prognostic) == 'D' 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, + active_tracers=active_tracers) + + self.parameters = parameters + g = parameters.g + H = parameters.H + + w1, w2, phi = self.tests + u, v, D = split(self.X) + u_trial = split(self.trials)[0] + + # -------------------------------------------------------------------- # + # Time Derivative Terms + # -------------------------------------------------------------------- # + mass_form = self.generate_mass_terms() + + # -------------------------------------------------------------------- # + # Transport Terms + # -------------------------------------------------------------------- # + # Velocity transport term + u_adv = prognostic(advection_form_1d(w1, u, u), 'u') + v_adv = prognostic(advection_form_1d(w2, v, u), 'v') + + # Depth transport term + D_adv = prognostic(continuity_form_1d(phi, D, u), 'D') + + # Transport term needs special linearisation + if self.linearisation_map(D_adv.terms[0]): + linear_D_adv = linear_continuity_form_1d(phi, H, u_trial) + # Add linearisation to D_adv + D_adv = linearisation(D_adv, linear_D_adv) + + adv_form = subject(u_adv + v_adv + D_adv, self.X) + + pressure_gradient_form = pressure_gradient(subject( + prognostic(-g * D * w1.dx(0) * dx, "u"), self.X)) + + self.residual = (mass_form + adv_form + + pressure_gradient_form) + + if fexpr is not None: + V = FunctionSpace(domain.mesh, 'CG', 1) + f = self.prescribed_fields('coriolis', V).interpolate(fexpr) + coriolis_form = coriolis(subject( + prognostic(-f * v * w1 * dx, "u") + + prognostic(f * u * w2 * dx, "v"), self.X)) + self.residual += coriolis_form + + if diffusion_options is not None: + for field, diffusion in diffusion_options: + idx = self.field_names.index(field) + test = self.tests[idx] + fn = split(self.X)[idx] + self.residual += subject( + prognostic(diffusion_form_1d(test, fn, diffusion.kappa), + field), + self.X) + + # -------------------------------------------------------------------- # + # Linearise equations + # -------------------------------------------------------------------- # + # Add linearisations to equations + self.residual = self.generate_linear_terms(self.residual, + self.linearisation_map) + + +class LinearShallowWaterEquations_1d(ShallowWaterEquations_1d): + u""" + Class for the linear (rotating) 1D shallow-water equations, which describe + the velocity 'u', 'v' and the depth field 'D', solving some variant of: \n + ∂u/∂t - fv + g*∂D/∂x = 0, \n + ∂v/∂t + fu = 0, \n + ∂D/∂t + H*∂u/∂x = 0, \n + for mean depth 'H', Coriolis parameter 'f' and gravity 'g'. + + This is set up the from the underlying :class:`ShallowWaterEquations_1d`, + which is then linearised. + """ + + def __init__(self, domain, parameters, fexpr=None, + space_names=None, linearisation_map='default', + no_normal_flow_bc_ids=None, active_tracers=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. + fexpr (:class:`ufl.Expr`, optional): an expression for the Coroilis + parameter. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + in which case the spaces are taken from the de Rham complex. Any + buoyancy variable is taken by default to lie in the L2 space. + 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 both time derivatives, + the 'D' transport term, pressure gradient and Coriolis terms. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations. Defaults to None. + """ + + 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)) + or (t.get(prognostic) == 'D' and t.has_label(transport))) + + super().__init__(domain, parameters, + fexpr=fexpr, space_names=space_names, + linearisation_map=linearisation_map, + no_normal_flow_bc_ids=no_normal_flow_bc_ids, + active_tracers=active_tracers) + + # Use the underlying routine to do a first linearisation of the equations + self.linearise_equation_set() diff --git a/gusto/thermodynamics.py b/gusto/equations/thermodynamics.py similarity index 100% rename from gusto/thermodynamics.py rename to gusto/equations/thermodynamics.py diff --git a/gusto/equations/transport_equations.py b/gusto/equations/transport_equations.py new file mode 100644 index 000000000..8be8b7f12 --- /dev/null +++ b/gusto/equations/transport_equations.py @@ -0,0 +1,128 @@ +"""Defines variants of the transport equation, in weak form.""" + +from firedrake import Function, inner, dx, MixedFunctionSpace, TestFunctions +from firedrake.fml import subject +from gusto.core.labels import time_derivative, prognostic +from gusto.equations.common_forms import advection_form, continuity_form +from gusto.equations.prognostic_equations import PrognosticEquation, PrognosticEquationSet + +__all__ = ["AdvectionEquation", "ContinuityEquation", "CoupledTransportEquation"] + + +class AdvectionEquation(PrognosticEquation): + u"""Discretises the advection equation, ∂q/∂t + (u.∇)q = 0""" + + def __init__(self, domain, function_space, field_name, Vu=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + function_space (:class:`FunctionSpace`): the function space that the + equation's prognostic is defined on. + field_name (str): name of the prognostic field. + Vu (:class:`FunctionSpace`, optional): the function space for the + velocity field. If this is not specified, uses the HDiv spaces + set up by the domain. Defaults to None. + """ + super().__init__(domain, function_space, field_name) + + if Vu is not None: + domain.spaces.add_space("HDiv", Vu, overwrite_space=True) + V = domain.spaces("HDiv") + u = self.prescribed_fields("u", V) + + test = self.test + q = self.X + mass_form = time_derivative(inner(q, test)*dx) + transport_form = advection_form(test, q, u) + + self.residual = prognostic(subject(mass_form + transport_form, q), field_name) + + +class ContinuityEquation(PrognosticEquation): + u"""Discretises the continuity equation, ∂q/∂t + ∇.(u*q) = 0""" + + def __init__(self, domain, function_space, field_name, Vu=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + function_space (:class:`FunctionSpace`): the function space that the + equation's prognostic is defined on. + field_name (str): name of the prognostic field. + Vu (:class:`FunctionSpace`, optional): the function space for the + velocity field. If this is not specified, uses the HDiv spaces + set up by the domain. Defaults to None. + """ + super().__init__(domain, function_space, field_name) + + if Vu is not None: + domain.spaces.add_space("HDiv", Vu, overwrite_space=True) + V = domain.spaces("HDiv") + u = self.prescribed_fields("u", V) + + test = self.test + q = self.X + mass_form = time_derivative(inner(q, test)*dx) + transport_form = continuity_form(test, q, u) + + self.residual = prognostic(subject(mass_form + transport_form, q), field_name) + + +class CoupledTransportEquation(PrognosticEquationSet): + u""" + Discretises the transport equation, \n + ∂q/∂t + (u.∇)q = F, \n + with the application of active tracers. + As there are multiple tracers or species that are interacting, q and F are + vectors. This equation can be enhanced through the addition of sources or + sinks (F) by applying it with physics schemes. + """ + def __init__(self, domain, active_tracers, Vu=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + active_tracers (list): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations. This is required for using this class; if there + is only a field to be advected, use the AdvectionEquation + instead. + Vu (:class:`FunctionSpace`, optional): the function space for the + velocity field. If this is not specified, uses the HDiv spaces + set up by the domain. Defaults to None. + """ + + self.active_tracers = active_tracers + self.terms_to_linearise = {} + self.field_names = [] + self.space_names = {} + + # Build finite element spaces + self.spaces = [] + + # Add active tracers to the list of prognostics + if active_tracers is None: + active_tracers = [] + self.add_tracers_to_prognostics(domain, active_tracers) + + # Make the full mixed function space + W = MixedFunctionSpace(self.spaces) + + full_field_name = "_".join(self.field_names) + PrognosticEquation.__init__(self, domain, W, full_field_name) + + if Vu is not None: + domain.spaces.add_space("HDiv", Vu, overwrite_space=True) + V = domain.spaces("HDiv") + _ = self.prescribed_fields("u", V) + + self.tests = TestFunctions(W) + self.X = Function(W) + + # Add mass forms for the tracers, which will use + # mass*density for any tracer_conservative terms + self.residual = self.generate_tracer_mass_terms(active_tracers) + + # Add transport of tracers + self.residual += self.generate_tracer_transport_terms(active_tracers) diff --git a/gusto/forcing.py b/gusto/forcing.py deleted file mode 100644 index b3cfa275f..000000000 --- a/gusto/forcing.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Discretisation of dynamic forcing terms, such as the pressure gradient.""" - -from firedrake import ( - Function, TrialFunctions, DirichletBC, LinearVariationalProblem, - LinearVariationalSolver -) -from firedrake.fml import drop, replace_subject -from gusto.labels import ( - transport, diffusion, time_derivative, hydrostatic, physics_label, - sponge, incompressible -) -from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual - - -__all__ = ["Forcing"] - - -class Forcing(object): - """ - Discretises forcing terms. - - This class describes the evaluation of forcing terms, for instance the - gravitational force, the Coriolis force or the pressure gradient force. - These are terms that can simply be evaluated, generally as part of some - semi-implicit time discretisation. - """ - - def __init__(self, equation, alpha): - """ - Args: - equation (:class:`PrognosticEquationSet`): the prognostic equations - containing the forcing terms. - alpha (:class:`Constant`): semi-implicit off-centering factor. An - alpha of 0 corresponds to fully explicit, while a factor of 1 - corresponds to fully implicit. - """ - - self.field_name = equation.field_name - implicit_terms = [incompressible, sponge] - dt = equation.domain.dt - - W = equation.function_space - self.x0 = Function(W) - self.xF = Function(W) - - # set up boundary conditions on the u subspace of W - bcs = [DirichletBC(W.sub(0), bc.function_arg, bc.sub_domain) for bc in equation.bcs['u']] - - # drop terms relating to transport, diffusion and physics - residual = equation.residual.label_map( - 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 - trials = TrialFunctions(W) - a = residual.label_map(lambda t: t.has_label(time_derivative), - replace_subject(trials), - map_if_false=drop) - - # the explicit forms are multiplied by (1-alpha) and moved to the rhs - L_explicit = -(1-alpha)*dt*residual.label_map( - lambda t: - any(t.has_label(time_derivative, hydrostatic, *implicit_terms, - return_tuple=True)), - drop, - replace_subject(self.x0)) - - # the implicit forms are multiplied by alpha and moved to the rhs - L_implicit = -alpha*dt*residual.label_map( - lambda t: - any(t.has_label( - time_derivative, hydrostatic, *implicit_terms, - return_tuple=True)), - drop, - replace_subject(self.x0)) - - # now add the terms that are always fully implicit - L_implicit -= dt*residual.label_map( - lambda t: any(t.has_label(*implicit_terms, return_tuple=True)), - replace_subject(self.x0), - drop) - - # the hydrostatic equations require some additional forms: - if any([t.has_label(hydrostatic) for t in residual]): - L_explicit += residual.label_map( - lambda t: t.has_label(hydrostatic), - replace_subject(self.x0), - drop) - - L_implicit -= residual.label_map( - lambda t: t.has_label(hydrostatic), - replace_subject(self.x0), - drop) - - # now we can set up the explicit and implicit problems - explicit_forcing_problem = LinearVariationalProblem( - a.form, L_explicit.form, self.xF, bcs=bcs - ) - - implicit_forcing_problem = LinearVariationalProblem( - a.form, L_implicit.form, self.xF, bcs=bcs - ) - - self.solvers = {} - self.solvers["explicit"] = LinearVariationalSolver( - explicit_forcing_problem, - options_prefix="ExplicitForcingSolver" - ) - self.solvers["implicit"] = LinearVariationalSolver( - implicit_forcing_problem, - options_prefix="ImplicitForcingSolver" - ) - - if logger.isEnabledFor(DEBUG): - self.solvers["explicit"].snes.ksp.setMonitor(logging_ksp_monitor_true_residual) - self.solvers["implicit"].snes.ksp.setMonitor(logging_ksp_monitor_true_residual) - - def apply(self, x_in, x_nl, x_out, label): - """ - Applies the discretisation for a forcing term F(x). - - This takes x_in and x_nl and computes F(x_nl), and updates x_out to \n - x_out = x_in + scale*F(x_nl) \n - where 'scale' is the appropriate semi-implicit factor. - - Args: - x_in (:class:`FieldCreator`): the field to be incremented. - x_nl (:class:`FieldCreator`): the field which the forcing term is - evaluated on. - x_out (:class:`FieldCreator`): the output field to be updated. - label (str): denotes which forcing to apply. Should be 'explicit' or - 'implicit'. TODO: there should be a check on this. Or this - should be an actual label. - """ - - self.x0.assign(x_nl(self.field_name)) - - self.solvers[label].solve() # places forcing in self.xF - - x_out.assign(x_in(self.field_name)) - x_out += self.xF diff --git a/gusto/initialisation/__init__.py b/gusto/initialisation/__init__.py new file mode 100644 index 000000000..6aa7f52ea --- /dev/null +++ b/gusto/initialisation/__init__.py @@ -0,0 +1,2 @@ +from gusto.initialisation.hydrostatic_initialisation import * # noqa +from gusto.initialisation.numerical_integrator import * # noqa \ No newline at end of file diff --git a/gusto/initialisation_tools.py b/gusto/initialisation/hydrostatic_initialisation.py similarity index 99% rename from gusto/initialisation_tools.py rename to gusto/initialisation/hydrostatic_initialisation.py index b0850812b..f4be78451 100644 --- a/gusto/initialisation_tools.py +++ b/gusto/initialisation/hydrostatic_initialisation.py @@ -1,4 +1,4 @@ -"""Tools for computing initial conditions, such as hydrostatic balance.""" +"""Tools for computing hydrostatically balanced initial conditions.""" from firedrake import MixedFunctionSpace, TrialFunctions, TestFunctions, \ TestFunction, TrialFunction, \ @@ -8,7 +8,7 @@ NonlinearVariationalProblem, NonlinearVariationalSolver, split, solve, \ FunctionSpace, errornorm, zero from gusto import thermodynamics -from gusto.logging import logger +from gusto.core import logger from gusto.recovery import Recoverer, BoundaryMethod diff --git a/gusto/numerical_integrator.py b/gusto/initialisation/numerical_integrator.py similarity index 100% rename from gusto/numerical_integrator.py rename to gusto/initialisation/numerical_integrator.py diff --git a/gusto/physics.py b/gusto/physics.py deleted file mode 100644 index a5a1ea525..000000000 --- a/gusto/physics.py +++ /dev/null @@ -1,1821 +0,0 @@ -""" -Objects to perform parametrisations of physical processes, or "physics". - -"PhysicsParametrisation" schemes are routines to compute updates to prognostic fields that -represent the action of non-fluid processes, or those fluid processes that are -unresolved. This module contains a set of these processes in the form of classes -with "evaluate" methods. -""" - -from abc import ABCMeta, abstractmethod -from firedrake import ( - Interpolator, conditional, Function, dx, sqrt, dot, min_value, - max_value, Constant, pi, Projector, grad, TestFunctions, split, - inner, TestFunction, exp, avg, outer, FacetNormal, - SpatialCoordinate, dS_v, NonlinearVariationalProblem, - NonlinearVariationalSolver -) -from firedrake.fml import identity, Term, subject -from gusto.active_tracers import Phases, TracerVariableType -from gusto.configuration import BoundaryLayerParameters -from gusto.recovery import Recoverer, BoundaryMethod -from gusto.equations import CompressibleEulerEquations -from gusto.labels import PhysicsLabel, transporting_velocity, transport, prognostic -from gusto.logging import logger -from gusto import thermodynamics -import ufl -import math -from enum import Enum -from types import FunctionType - - -__all__ = ["SaturationAdjustment", "Fallout", "Coalescence", "EvaporationOfRain", - "AdvectedMoments", "InstantRain", "SWSaturationAdjustment", - "SourceSink", "SurfaceFluxes", "WindDrag", "StaticAdjustment", - "SuppressVerticalWind", "BoundaryLayerMixing", "TerminatorToy"] - - -class PhysicsParametrisation(object, metaclass=ABCMeta): - """Base class for the parametrisation of physical processes for Gusto.""" - - def __init__(self, equation, label_name, parameters=None): - """ - Args: - equation (:class:`PrognosticEquationSet`): the model's equation. - label_name (str): name of physics scheme, to be passed to its label. - parameters (:class:`Configuration`, optional): parameters containing - the values of gas constants. Defaults to None, in which case the - parameters are obtained from the equation. - """ - - self.label = PhysicsLabel(label_name) - self.equation = equation - if parameters is None and hasattr(equation, 'parameters'): - self.parameters = equation.parameters - else: - self.parameters = parameters - - @abstractmethod - def evaluate(self): - """ - Computes the value of physics source and sink terms. - """ - pass - - -class SourceSink(PhysicsParametrisation): - """ - The source or sink of some variable, described through a UFL expression. - - A scheme representing the general source or sink of a variable, described - through a UFL expression. The expression should be for the rate of change - of the variable. It is intended that the source/sink is independent of the - prognostic variables. - - The expression can also be a time-varying expression. In which case a - function should be provided, taking a :class:`Constant` as an argument (to - represent the time.) - """ - - def __init__(self, equation, variable_name, rate_expression, - time_varying=False, method='interpolate'): - """ - Args: - equation (:class:`PrognosticEquationSet`): the model's equation. - variable_name (str): the name of the variable - rate_expression (:class:`ufl.Expr` or func): an expression giving - the rate of change of the variable. If a time-varying expression - is needed, this should be a function taking a single argument - representing the time. Then the `time_varying` argument must - be set to True. - time_varying (bool, optional): whether the source/sink expression - varies with time. Defaults to False. - method (str, optional): the method to use to evaluate the expression - for the source. Valid options are 'interpolate' or 'project'. - Default is 'interpolate'. - """ - - label_name = f'source_sink_{variable_name}' - super().__init__(equation, label_name, parameters=None) - - if method not in ['interpolate', 'project']: - raise ValueError(f'Method {method} for source/sink evaluation not valid') - self.method = method - self.time_varying = time_varying - self.variable_name = variable_name - - # Check the variable exists - if hasattr(equation, "field_names"): - assert variable_name in equation.field_names, \ - f'Field {variable_name} does not exist in the equation set' - else: - assert variable_name == equation.field_name, \ - f'Field {variable_name} does not exist in the equation' - - # Work out the appropriate function space - if hasattr(equation, "field_names"): - V_idx = equation.field_names.index(variable_name) - W = equation.function_space - V = W.sub(V_idx) - test = equation.tests[V_idx] - else: - V = equation.function_space - test = equation.test - - # Make source/sink term - self.source = Function(V) - equation.residual += self.label(subject(test * self.source * dx, equation.X), - self.evaluate) - - # Handle whether the expression is time-varying or not - if self.time_varying: - expression = rate_expression(equation.domain.t) - else: - expression = rate_expression - - # Handle method of evaluating source/sink - if self.method == 'interpolate': - self.source_interpolator = Interpolator(expression, V) - else: - self.source_projector = Projector(expression, V) - - # If not time-varying, evaluate for the first time here - if not self.time_varying: - if self.method == 'interpolate': - self.source.assign(self.source_interpolator.interpolate()) - else: - self.source.assign(self.source_projector.project()) - - 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. Unused. - dt: (:class:`Constant`): the timestep, which can be the time - interval for the scheme. Unused. - """ - if self.time_varying: - logger.info(f'Evaluating physics parametrisation {self.label.label}') - if self.method == 'interpolate': - self.source.assign(self.source_interpolator.interpolate()) - else: - self.source.assign(self.source_projector.project()) - else: - pass - - -class SaturationAdjustment(PhysicsParametrisation): - """ - Represents the phase change between water vapour and cloud liquid. - - This class computes updates to water vapour, cloud liquid and (virtual dry) - potential temperature, representing the action of condensation of water - vapour and/or evaporation of cloud liquid, with the associated latent heat - change. The parametrisation follows the saturation adjustment used in Bryan - and Fritsch (2002). - - Currently this is only implemented for use with mixing ratio variables, and - with "theta" assumed to be the virtual dry potential temperature. Latent - heating effects are always assumed, and the total mixing ratio is conserved. - A filter is applied to prevent generation of negative mixing ratios. - """ - - def __init__(self, equation, vapour_name='water_vapour', - cloud_name='cloud_water', latent_heat=True, parameters=None): - """ - Args: - equation (:class:`PrognosticEquationSet`): the model's equation. - vapour_name (str, optional): name of the water vapour variable. - Defaults to 'water_vapour'. - cloud_name (str, optional): name of the cloud water variable. - Defaults to 'cloud_water'. - latent_heat (bool, optional): whether to have latent heat exchange - feeding back from the phase change. Defaults to True. - parameters (:class:`Configuration`, optional): parameters containing - the values of gas constants. Defaults to None, in which case the - parameters are obtained from the equation. - - Raises: - NotImplementedError: currently this is only implemented for the - CompressibleEulerEquations. - """ - - label_name = 'saturation_adjustment' - self.explicit_only = True - super().__init__(equation, label_name, parameters=parameters) - - # TODO: make a check on the variable type of the active tracers - # if not a mixing ratio, we need to convert to mixing ratios - # this will be easier if we change equations to have dictionary of - # active tracer metadata corresponding to variable names - - # Check that fields exist - if vapour_name not in equation.field_names: - raise ValueError(f"Field {vapour_name} does not exist in the equation set") - if cloud_name not in equation.field_names: - raise ValueError(f"Field {cloud_name} does not exist in the equation set") - - # Make prognostic for physics scheme - parameters = self.parameters - self.X = Function(equation.X.function_space()) - self.latent_heat = latent_heat - - # Vapour and cloud variables are needed for every form of this scheme - cloud_idx = equation.field_names.index(cloud_name) - vap_idx = equation.field_names.index(vapour_name) - cloud_water = self.X.subfunctions[cloud_idx] - water_vapour = self.X.subfunctions[vap_idx] - - # Indices of variables in mixed function space - V_idxs = [vap_idx, cloud_idx] - V = equation.function_space.sub(vap_idx) # space in which to do the calculation - - # Get variables used to calculate saturation curve - if isinstance(equation, CompressibleEulerEquations): - rho_idx = equation.field_names.index('rho') - theta_idx = equation.field_names.index('theta') - rho = self.X.subfunctions[rho_idx] - theta = self.X.subfunctions[theta_idx] - if latent_heat: - V_idxs.append(theta_idx) - - # need to evaluate rho at theta-points, and do this via recovery - boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None - rho_averaged = Function(V) - self.rho_recoverer = Recoverer(rho, rho_averaged, boundary_method=boundary_method) - - exner = thermodynamics.exner_pressure(parameters, rho_averaged, theta) - T = thermodynamics.T(parameters, theta, exner, r_v=water_vapour) - p = thermodynamics.p(parameters, exner) - - else: - raise NotImplementedError( - 'Saturation adjustment only implemented for the Compressible Euler equations') - - # -------------------------------------------------------------------- # - # Compute heat capacities and calculate saturation curve - # -------------------------------------------------------------------- # - # Loop through variables to extract all liquid components - liquid_water = cloud_water - for active_tracer in equation.active_tracers: - if (active_tracer.phase == Phases.liquid - and active_tracer.chemical == 'H2O' and active_tracer.name != cloud_name): - liq_idx = equation.field_names.index(active_tracer.name) - liquid_water += self.X.subfunctions[liq_idx] - - # define some parameters as attributes - self.dt = Constant(0.0) - R_d = parameters.R_d - cp = parameters.cp - cv = parameters.cv - c_pv = parameters.c_pv - c_pl = parameters.c_pl - c_vv = parameters.c_vv - R_v = parameters.R_v - - # make useful fields - L_v = thermodynamics.Lv(parameters, T) - R_m = R_d + R_v * water_vapour - c_pml = cp + c_pv * water_vapour + c_pl * liquid_water - c_vml = cv + c_vv * water_vapour + c_pl * liquid_water - - # use Teten's formula to calculate the saturation curve - sat_expr = thermodynamics.r_sat(parameters, T, p) - - # -------------------------------------------------------------------- # - # Saturation adjustment expression - # -------------------------------------------------------------------- # - # make appropriate condensation rate - sat_adj_expr = (water_vapour - sat_expr) / self.dt - if latent_heat: - # As condensation/evaporation happens, the temperature changes - # so need to take this into account with an extra factor - sat_adj_expr = sat_adj_expr / (1.0 + ((L_v ** 2.0 * sat_expr) - / (cp * R_v * T ** 2.0))) - - # adjust the rate so that so negative values don't occur - sat_adj_expr = conditional(sat_adj_expr < 0, - max_value(sat_adj_expr, -cloud_water / self.dt), - min_value(sat_adj_expr, water_vapour / self.dt)) - - # -------------------------------------------------------------------- # - # Factors for multiplying source for different variables - # -------------------------------------------------------------------- # - # Factors need to have same shape as V_idxs - factors = [Constant(1.0), Constant(-1.0)] - if latent_heat and isinstance(equation, CompressibleEulerEquations): - factors.append(-theta * (cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml / (R_m * cp * c_vml))) - - # -------------------------------------------------------------------- # - # 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)] - - tests = [equation.tests[idx] for idx in V_idxs] - - # Add source terms to residual - for test, source in zip(tests, self.source): - equation.residual += self.label(subject(test * source * dx, - equation.X), self.evaluate) - - def evaluate(self, x_in, dt): - """ - Evaluates the source/sink for the saturation adjustment process. - - Args: - x_in (:class:`Function`): the (mixed) field to be evolved. - dt (:class:`Constant`): the time interval for the scheme. - """ - logger.info(f'Evaluating physics parametrisation {self.label.label}') - # Update the values of internal variables - self.dt.assign(dt) - self.X.assign(x_in) - if isinstance(self.equation, CompressibleEulerEquations): - self.rho_recoverer.project() - # Evaluate the source - for interpolator in self.source_interpolators: - interpolator.interpolate() - - -class AdvectedMoments(Enum): - """ - Enumerator describing the moments in the raindrop size distribution. - - This stores enumerators for the number of moments used to describe the - size distribution curve of raindrops. This can be used for deciding which - moments to advect in a precipitation scheme. - """ - - M0 = 0 # don't advect the distribution -- advect all rain at the same speed - M3 = 1 # advect the mass of the distribution - - -class Fallout(PhysicsParametrisation): - """ - Represents the fallout process of hydrometeors. - - Precipitation is described by downwards transport of tracer species. This - process determines the precipitation velocity from the `AdvectedMoments` - enumerator, which either: - (a) sets a terminal velocity of 5 m/s - (b) determines a rainfall size distribution based on a Gamma distribution, - as in Paluch (1979). The droplets are based on the mean mass of the rain - droplets (aka a single-moment scheme). - - Outflow boundary conditions are applied to the transport, so the rain will - flow out of the bottom of the domain. - - This is currently only implemented for "rain" in the compressible Euler - equation set. This variable must be a mixing ratio, It is only implemented - for Cartesian geometry. - """ - - def __init__(self, equation, rain_name, domain, transport_method, - moments=AdvectedMoments.M3): - """ - Args: - equation (:class:`PrognosticEquationSet`): the model's equation. - rain_name (str, optional): name of the rain variable. Defaults to - 'rain'. - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - transport_method (:class:`TransportMethod`): the spatial method - used for transporting the rain. - moments (int, optional): an :class:`AdvectedMoments` enumerator, - representing the number of moments of the size distribution of - raindrops to be transported. Defaults to `AdvectedMoments.M3`. - """ - - label_name = f'fallout_{rain_name}' - super().__init__(equation, label_name, parameters=None) - - # Check that fields exist - if rain_name not in equation.field_names: - raise ValueError(f"Field {rain_name} does not exist in the equation set") - - # Check if variable is a mixing ratio - rain_tracer = equation.get_active_tracer(rain_name) - if rain_tracer.variable_type != TracerVariableType.mixing_ratio: - raise NotImplementedError('Fallout only implemented when rain ' - + 'variable is a mixing ratio') - - # Set up rain and velocity - self.X = Function(equation.X.function_space()) - - rain_idx = equation.field_names.index(rain_name) - rain = self.X.subfunctions[rain_idx] - - Vu = domain.spaces("HDiv") - # TODO: there must be a better way than forcing this into the equation - v = equation.prescribed_fields(name='rainfall_velocity', space=Vu) - - # -------------------------------------------------------------------- # - # Create physics term -- which is actually a transport term - # -------------------------------------------------------------------- # - - assert transport_method.outflow, \ - 'Fallout requires a transport method with outflow=True' - adv_term = transport_method.form - # Add rainfall velocity by replacing transport_velocity in term - adv_term = adv_term.label_map(identity, - map_if_true=lambda t: Term( - ufl.replace(t.form, {t.get(transporting_velocity): v}), - t.labels)) - - # We don't want this term to be picked up by normal transport, so drop - # the transport label - adv_term = transport.remove(adv_term) - - adv_term = prognostic(subject(adv_term, equation.X), rain_name) - equation.residual += self.label(adv_term, self.evaluate) - - # -------------------------------------------------------------------- # - # Expressions for determining rainfall velocity - # -------------------------------------------------------------------- # - self.moments = moments - - if moments == AdvectedMoments.M0: - # all rain falls at terminal velocity - terminal_velocity = Constant(5) # in m/s - v.project(-terminal_velocity*domain.k) - elif moments == AdvectedMoments.M3: - self.explicit_only = True - # this advects the third moment M3 of the raindrop - # distribution, which corresponds to the mean mass - rho_idx = equation.field_names.index('rho') - rho = self.X.subfunctions[rho_idx] - rho_w = Constant(1000.0) # density of liquid water - # assume n(D) = n_0 * D^mu * exp(-Lambda*D) - # n_0 = N_r * Lambda^(1+mu) / gamma(1 + mu) - N_r = Constant(10**5) # number of rain droplets per m^3 - mu = 0.0 # shape constant of droplet gamma distribution - # assume V(D) = a * D^b * exp(-f*D) * (rho_0 / rho)^g - # take f = 0 - a = Constant(362.) # intercept for velocity distr. in log space - b = 0.65 # inverse scale parameter for velocity distr. - rho0 = Constant(1.22) # reference density in kg/m^3 - g = Constant(0.5) # scaling of density correction - # we keep mu in the expressions even though mu = 0 - threshold = Constant(10**-10) # only do rainfall for r > threshold - Lambda = (N_r * pi * rho_w * math.gamma(4 + mu) - / (6 * math.gamma(1 + mu) * rho * rain)) ** (1. / 3) - Lambda0 = (N_r * pi * rho_w * math.gamma(4 + mu) - / (6 * math.gamma(1 + mu) * rho * threshold)) ** (1. / 3) - v_expression = conditional(rain > threshold, - (a * math.gamma(4 + b + mu) - / (math.gamma(4 + mu) * Lambda ** b) - * (rho0 / rho) ** g), - (a * math.gamma(4 + b + mu) - / (math.gamma(4 + mu) * Lambda0 ** b) - * (rho0 / rho) ** g)) - else: - raise NotImplementedError( - 'Currently there are only implementations for zero and one ' - + 'moment schemes for rainfall. Valid options are ' - + 'AdvectedMoments.M0 and AdvectedMoments.M3') - - if moments != AdvectedMoments.M0: - # TODO: introduce reduced projector - test = TestFunction(Vu) - dx_reduced = dx(degree=4) - proj_eqn = inner(test, v + v_expression*domain.k)*dx_reduced - proj_prob = NonlinearVariationalProblem(proj_eqn, v) - self.determine_v = NonlinearVariationalSolver(proj_prob) - - def evaluate(self, x_in, dt): - """ - Evaluates the source/sink corresponding to the fallout process. - - Args: - x_in (:class:`Function`): the (mixed) field to be evolved. - dt (:class:`Constant`): the time interval for the scheme. - """ - logger.info(f'Evaluating physics parametrisation {self.label.label}') - self.X.assign(x_in) - if self.moments != AdvectedMoments.M0: - self.determine_v.solve() - - -class Coalescence(PhysicsParametrisation): - """ - Represents the coalescence of cloud droplets to form rain droplets. - - Coalescence is the process of forming rain droplets from cloud droplets. - This scheme performs that process, using two parts: accretion, which is - independent of the rain concentration, and auto-accumulation, which is - accelerated by the existence of rain. These parametrisations come from Klemp - and Wilhelmson (1978). The rate of change is limited to prevent production - of negative moisture values. - - This is only implemented for mixing ratio variables. - """ - - def __init__(self, equation, cloud_name='cloud_water', rain_name='rain', - accretion=True, accumulation=True): - """ - Args: - equation (:class:`PrognosticEquationSet`): the model's equation. - cloud_name (str, optional): name of the cloud variable. Defaults to - 'cloud_water'. - rain_name (str, optional): name of the rain variable. Defaults to - 'rain'. - accretion (bool, optional): whether to include the accretion process - in the parametrisation. Defaults to True. - accumulation (bool, optional): whether to include the accumulation - process in the parametrisation. Defaults to True. - """ - - self.explicit_only = True - label_name = "coalescence" - if accretion: - label_name += "_accretion" - if accumulation: - label_name += "_accumulation" - super().__init__(equation, label_name, parameters=None) - - # Check that fields exist - if cloud_name not in equation.field_names: - raise ValueError(f"Field {cloud_name} does not exist in the equation set") - if rain_name not in equation.field_names: - raise ValueError(f"Field {rain_name} does not exist in the equation set") - - self.cloud_idx = equation.field_names.index(cloud_name) - self.rain_idx = equation.field_names.index(rain_name) - Vcl = equation.function_space.sub(self.cloud_idx) - Vr = equation.function_space.sub(self.rain_idx) - self.cloud_water = Function(Vcl) - self.rain = Function(Vr) - - # declare function space and source field - Vt = self.cloud_water.function_space() - self.source = Function(Vt) - - # define some parameters as attributes - self.dt = Constant(0.0) - # TODO: should these parameters be hard-coded or configurable? - k_1 = Constant(0.001) # accretion rate in 1/s - k_2 = Constant(2.2) # accumulation rate in 1/s - a = Constant(0.001) # min cloud conc in kg/kg - b = Constant(0.875) # power for rain in accumulation - - # make default rates to be zero - accr_rate = Constant(0.0) - accu_rate = Constant(0.0) - - if accretion: - accr_rate = k_1 * (self.cloud_water - a) - if accumulation: - accu_rate = k_2 * self.cloud_water * self.rain ** b - - # Expression for rain increment, with conditionals to prevent negative values - rain_expr = conditional(self.rain < 0.0, # if rain is negative do only accretion - conditional(accr_rate < 0.0, - 0.0, - min_value(accr_rate, self.cloud_water / self.dt)), - # don't turn rain back into cloud - conditional(accr_rate + accu_rate < 0.0, - 0.0, - # if accretion rate is negative do only accumulation - conditional(accr_rate < 0.0, - 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) - - # Add term to equation's residual - test_cl = equation.tests[self.cloud_idx] - test_r = equation.tests[self.rain_idx] - equation.residual += self.label(subject(test_cl * self.source * dx - - test_r * self.source * dx, - equation.X), - self.evaluate) - - def evaluate(self, x_in, dt): - """ - Evaluates the source/sink for the coalescence process. - - Args: - x_in (:class:`Function`): the (mixed) field to be evolved. - dt (:class:`Constant`): the time interval for the scheme. - """ - logger.info(f'Evaluating physics parametrisation {self.label.label}') - # Update the values of internal variables - self.dt.assign(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()) - - -class EvaporationOfRain(PhysicsParametrisation): - """ - Represents the evaporation of rain into water vapour. - - This describes the evaporation of rain into water vapour, with the - associated latent heat change. This parametrisation comes from Klemp and - Wilhelmson (1978). The rate of change is limited to prevent production of - negative moisture values. - - This is only implemented for mixing ratio variables, and when the prognostic - is the virtual dry potential temperature. - """ - - def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', - latent_heat=True): - """ - Args: - equation (:class:`PrognosticEquationSet`): the model's equation. - cloud_name (str, optional): name of the rain variable. Defaults to - 'rain'. - vapour_name (str, optional): name of the water vapour variable. - Defaults to 'water_vapour'. - latent_heat (bool, optional): whether to have latent heat exchange - feeding back from the phase change. Defaults to True. - - Raises: - NotImplementedError: currently this is only implemented for the - CompressibleEulerEquations. - """ - - self.explicit_only = True - label_name = 'evaporation_of_rain' - super().__init__(equation, label_name, parameters=None) - - # TODO: make a check on the variable type of the active tracers - # if not a mixing ratio, we need to convert to mixing ratios - # this will be easier if we change equations to have dictionary of - # active tracer metadata corresponding to variable names - - # Check that fields exist - if vapour_name not in equation.field_names: - raise ValueError(f"Field {vapour_name} does not exist in the equation set") - if rain_name not in equation.field_names: - raise ValueError(f"Field {rain_name} does not exist in the equation set") - - # Make prognostic for physics scheme - self.X = Function(equation.X.function_space()) - parameters = self.parameters - self.latent_heat = latent_heat - - # Vapour and cloud variables are needed for every form of this scheme - rain_idx = equation.field_names.index(rain_name) - vap_idx = equation.field_names.index(vapour_name) - rain = self.X.subfunctions[rain_idx] - water_vapour = self.X.subfunctions[vap_idx] - - # Indices of variables in mixed function space - V_idxs = [rain_idx, vap_idx] - V = equation.function_space.sub(rain_idx) # space in which to do the calculation - - # Get variables used to calculate saturation curve - if isinstance(equation, CompressibleEulerEquations): - rho_idx = equation.field_names.index('rho') - theta_idx = equation.field_names.index('theta') - rho = self.X.subfunctions[rho_idx] - theta = self.X.subfunctions[theta_idx] - if latent_heat: - V_idxs.append(theta_idx) - - # need to evaluate rho at theta-points, and do this via recovery - boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None - rho_averaged = Function(V) - self.rho_recoverer = Recoverer(rho, rho_averaged, boundary_method=boundary_method) - - exner = thermodynamics.exner_pressure(parameters, rho_averaged, theta) - T = thermodynamics.T(parameters, theta, exner, r_v=water_vapour) - p = thermodynamics.p(parameters, exner) - - # -------------------------------------------------------------------- # - # Compute heat capacities and calculate saturation curve - # -------------------------------------------------------------------- # - # Loop through variables to extract all liquid components - liquid_water = rain - for active_tracer in equation.active_tracers: - if (active_tracer.phase == Phases.liquid - and active_tracer.chemical == 'H2O' and active_tracer.name != rain_name): - liq_idx = equation.field_names.index(active_tracer.name) - liquid_water += self.X.subfunctions[liq_idx] - - # define some parameters as attributes - self.dt = Constant(0.0) - R_d = parameters.R_d - cp = parameters.cp - cv = parameters.cv - c_pv = parameters.c_pv - c_pl = parameters.c_pl - c_vv = parameters.c_vv - R_v = parameters.R_v - - # make useful fields - L_v = thermodynamics.Lv(parameters, T) - R_m = R_d + R_v * water_vapour - c_pml = cp + c_pv * water_vapour + c_pl * liquid_water - c_vml = cv + c_vv * water_vapour + c_pl * liquid_water - - # use Teten's formula to calculate the saturation curve - sat_expr = thermodynamics.r_sat(parameters, T, p) - - # -------------------------------------------------------------------- # - # Evaporation expression - # -------------------------------------------------------------------- # - # TODO: should these parameters be hard-coded or configurable? - # expression for ventilation factor - a = Constant(1.6) - b = Constant(124.9) - c = Constant(0.2046) - C = a + b * (rho_averaged * rain) ** c - - # make appropriate condensation rate - f = Constant(5.4e5) - g = Constant(2.55e6) - h = Constant(0.525) - evap_rate = (((1 - water_vapour / sat_expr) * C * (rho_averaged * rain) ** h) - / (rho_averaged * (f + g / (p * sat_expr)))) - - # adjust evap rate so negative rain doesn't occur - evap_rate = conditional(evap_rate < 0, 0.0, - conditional(rain < 0.0, 0.0, - min_value(evap_rate, rain / self.dt))) - - # -------------------------------------------------------------------- # - # Factors for multiplying source for different variables - # -------------------------------------------------------------------- # - # Factors need to have same shape as V_idxs - factors = [Constant(-1.0), Constant(1.0)] - if latent_heat and isinstance(equation, CompressibleEulerEquations): - factors.append(-theta * (cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml / (R_m * cp * c_vml))) - - # -------------------------------------------------------------------- # - # 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)] - - tests = [equation.tests[idx] for idx in V_idxs] - - # Add source terms to residual - for test, source in zip(tests, self.source): - equation.residual += self.label(subject(test * source * dx, - equation.X), self.evaluate) - - def evaluate(self, x_in, dt): - """ - Applies the process to evaporate rain droplets. - - Args: - x_in (:class:`Function`): the (mixed) field to be evolved. - dt (:class:`Constant`): the time interval for the scheme. - """ - logger.info(f'Evaluating physics parametrisation {self.label.label}') - # Update the values of internal variables - self.dt.assign(dt) - self.X.assign(x_in) - if isinstance(self.equation, CompressibleEulerEquations): - self.rho_recoverer.project() - # Evaluate the source - for interpolator in self.source_interpolators: - interpolator.interpolate() - - -class InstantRain(PhysicsParametrisation): - """ - The process of converting vapour above the saturation curve to rain. - - A scheme to move vapour directly to rain. If convective feedback is true - then this process feeds back directly on the height equation. If rain is - accumulating then excess vapour is being tracked and stored as rain; - otherwise converted vapour is not recorded. The process can happen over the - timestep dt or over a specified time interval tau. - """ - - def __init__(self, equation, saturation_curve, - time_varying_saturation=False, - vapour_name="water_vapour", rain_name=None, gamma_r=1, - convective_feedback=False, beta1=None, tau=None, - parameters=None): - """ - Args: - equation (:class:`PrognosticEquationSet`): the model's equation. - saturation_curve (:class:`ufl.Expr` or func): the curve above which - excess moisture is converted to rain. Is either prescribed or - dependent on a prognostic field. - time_varying_saturation (bool, optional): set this to True if the - saturation curve is changing in time. Defaults to False. - vapour_name (str, optional): name of the water vapour variable. - Defaults to "water_vapour". - rain_name (str, optional): name of the rain variable. Defaults to - None. - gamma_r (float, optional): Fraction of vapour above the threshold - which is converted to rain. Defaults to one, in which case all - vapour above the threshold is converted. - convective_feedback (bool, optional): True if the conversion of - vapour affects the height equation. Defaults to False. - beta1 (float, optional): Condensation proportionality constant, - used if convection causes a response in the height equation. - Defaults to None, but must be specified if convective_feedback - is True. - tau (float, optional): Timescale for condensation. Defaults to None, - in which case the timestep dt is used. - parameters (:class:`Configuration`, optional): parameters containing - the values of gas constants. Defaults to None, in which case the - parameters are obtained from the equation. - """ - - self.explicit_only = True - label_name = 'instant_rain' - super().__init__(equation, label_name, parameters=parameters) - - self.convective_feedback = convective_feedback - self.time_varying_saturation = time_varying_saturation - - # check for the correct fields - assert vapour_name in equation.field_names, f"Field {vapour_name} does not exist in the equation set" - self.Vv_idx = equation.field_names.index(vapour_name) - - if rain_name is not None: - assert rain_name in equation.field_names, f"Field {rain_name} does not exist in the equation set " - - if self.convective_feedback: - assert "D" in equation.field_names, "Depth field must exist for convective feedback" - assert beta1 is not None, "If convective feedback is used, beta1 parameter must be specified" - - # obtain function space and functions; vapour needed for all cases - W = equation.function_space - Vv = W.sub(self.Vv_idx) - test_v = equation.tests[self.Vv_idx] - - # depth needed if convective feedback - if self.convective_feedback: - self.VD_idx = equation.field_names.index("D") - VD = W.sub(self.VD_idx) - test_D = equation.tests[self.VD_idx] - self.D = Function(VD) - - # the source function is the difference between the water vapour and - # the saturation function - self.water_v = Function(Vv) - self.source = Function(Vv) - - # tau is the timescale for conversion (may or may not be the timestep) - if tau is not None: - self.set_tau_to_dt = False - self.tau = tau - else: - self.set_tau_to_dt = True - self.tau = Constant(0) - logger.info("Timescale for rain conversion has been set to dt. If this is not the intention then provide a tau parameter as an argument to InstantRain.") - - if self.time_varying_saturation: - if isinstance(saturation_curve, FunctionType): - self.saturation_computation = saturation_curve - self.saturation_curve = Function(Vv) - else: - raise NotImplementedError( - "If time_varying_saturation is True then saturation must be a Python function of a prognostic field.") - else: - assert not isinstance(saturation_curve, FunctionType), "If time_varying_saturation is not True then saturation cannot be a Python function." - self.saturation_curve = saturation_curve - - # lose proportion of vapour above the saturation curve - equation.residual += self.label(subject(test_v * self.source * dx, - equation.X), - self.evaluate) - - # if rain is not none then the excess vapour is being tracked and is - # added to rain - if rain_name is not None: - Vr_idx = equation.field_names.index(rain_name) - test_r = equation.tests[Vr_idx] - equation.residual -= self.label(subject(test_r * self.source * dx, - equation.X), - self.evaluate) - - # if feeding back on the height adjust the height equation - if convective_feedback: - equation.residual += self.label(subject(test_D * beta1 * self.source * dx, - equation.X), - self.evaluate) - - # interpolator does the conversion of vapour to rain - self.source_interpolator = Interpolator(conditional( - self.water_v > self.saturation_curve, - (1/self.tau)*gamma_r*(self.water_v - self.saturation_curve), - 0), Vv) - - def evaluate(self, x_in, dt): - """ - Evalutes the source term generated by the physics. - - Computes the physics contributions (loss of vapour, accumulation of - rain and loss of height due to convection) at each timestep. - - 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. - """ - logger.info(f'Evaluating physics parametrisation {self.label.label}') - if self.convective_feedback: - self.D.assign(x_in.subfunctions[self.VD_idx]) - if self.time_varying_saturation: - self.saturation_curve.interpolate(self.saturation_computation(x_in)) - if self.set_tau_to_dt: - self.tau.assign(dt) - self.water_v.assign(x_in.subfunctions[self.Vv_idx]) - self.source.assign(self.source_interpolator.interpolate()) - - -class SWSaturationAdjustment(PhysicsParametrisation): - """ - Represents the process of adjusting water vapour and cloud water according - to a saturation function, via condensation and evaporation processes. - - This physics scheme follows that of Zerroukat and Allen (2015). - - """ - - def __init__(self, equation, saturation_curve, - time_varying_saturation=False, vapour_name='water_vapour', - cloud_name='cloud_water', convective_feedback=False, - beta1=None, thermal_feedback=False, beta2=None, gamma_v=1, - time_varying_gamma_v=False, tau=None, - parameters=None): - """ - Args: - equation (:class:`PrognosticEquationSet`): the model's equation - saturation_curve (:class:`ufl.Expr` or func): the curve which - dictates when phase changes occur. In a saturated atmosphere - vapour above the saturation curve becomes cloud, and if the - atmosphere is sub-saturated and there is cloud present cloud - will become vapour until the saturation curve is reached. The - saturation curve is either prescribed or dependent on a - prognostic field. - time_varying_saturation (bool, optional): set this to True if the - saturation curve is changing in time. Defaults to False. - vapour_name (str, optional): name of the water vapour variable. - Defaults to 'water_vapour'. - cloud_name (str, optional): name of the cloud variable. Defaults to - 'cloud_water'. - convective_feedback (bool, optional): True if the conversion of - vapour affects the height equation. Defaults to False. - beta1 (float, optional): Condensation proportionality constant for - height feedback, used if convection causes a response in the - height equation. Defaults to None, but must be specified if - convective_feedback is True. - thermal_feedback (bool, optional): True if moist conversions - affect the buoyancy equation. Defaults to False. - beta2 (float, optional): Condensation proportionality constant - for thermal feedback. Defaults to None, but must be specified - if thermal_feedback is True. This is equivalent to the L_v - parameter in Zerroukat and Allen (2015). - gamma_v (ufl expression or :class: `function`): The proportion of - moist species that is converted when a conversion between - vapour and cloud is taking place. Defaults to one, in which - case the full amount of species to bring vapour to the - saturation curve will undergo a conversion. Converting only a - fraction avoids a two-timestep oscillation between vapour and - cloud when saturation is tempertature/height-dependent. - time_varying_gamma_v (bool, optional): set this to True - if the fraction of moist species converted changes in time - (if gamma_v is temperature/height-dependent). - tau (float, optional): Timescale for condensation and evaporation. - Defaults to None, in which case the timestep dt is used. - parameters (:class:`Configuration`, optional): parameters containing - the values of constants. Defaults to None, in which case the - parameters are obtained from the equation. - """ - - self.explicit_only = True - label_name = 'saturation_adjustment' - super().__init__(equation, label_name, parameters=parameters) - - self.time_varying_saturation = time_varying_saturation - self.convective_feedback = convective_feedback - self.thermal_feedback = thermal_feedback - self.time_varying_gamma_v = time_varying_gamma_v - - # Check for the correct fields - assert vapour_name in equation.field_names, f"Field {vapour_name} does not exist in the equation set" - assert cloud_name in equation.field_names, f"Field {cloud_name} does not exist in the equation set" - self.Vv_idx = equation.field_names.index(vapour_name) - self.Vc_idx = equation.field_names.index(cloud_name) - - if self.convective_feedback: - assert "D" in equation.field_names, "Depth field must exist for convective feedback" - assert beta1 is not None, "If convective feedback is used, beta1 parameter must be specified" - - if self.thermal_feedback: - assert "b" in equation.field_names, "Buoyancy field must exist for thermal feedback" - assert beta2 is not None, "If thermal feedback is used, beta2 parameter must be specified" - - # Obtain function spaces and functions - W = equation.function_space - Vv = W.sub(self.Vv_idx) - Vc = W.sub(self.Vc_idx) - V_idxs = [self.Vv_idx, self.Vc_idx] - - # Source functions for both vapour and cloud - self.water_v = Function(Vv) - self.cloud = Function(Vc) - - # depth needed if convective feedback - if self.convective_feedback: - self.VD_idx = equation.field_names.index("D") - VD = W.sub(self.VD_idx) - self.D = Function(VD) - V_idxs.append(self.VD_idx) - - # buoyancy needed if thermal feedback - if self.thermal_feedback: - self.Vb_idx = equation.field_names.index("b") - Vb = W.sub(self.Vb_idx) - self.b = Function(Vb) - V_idxs.append(self.Vb_idx) - - # tau is the timescale for condensation/evaporation (may or may not be the timestep) - if tau is not None: - self.set_tau_to_dt = False - self.tau = tau - else: - self.set_tau_to_dt = True - self.tau = Constant(0) - logger.info("Timescale for moisture conversion between vapour and cloud has been set to dt. If this is not the intention then provide a tau parameter as an argument to SWSaturationAdjustment.") - - if self.time_varying_saturation: - if isinstance(saturation_curve, FunctionType): - self.saturation_computation = saturation_curve - self.saturation_curve = Function(Vv) - else: - raise NotImplementedError( - "If time_varying_saturation is True then saturation must be a Python function of at least one prognostic field.") - else: - assert not isinstance(saturation_curve, FunctionType), "If time_varying_saturation is not True then saturation cannot be a Python function." - self.saturation_curve = saturation_curve - - # Saturation adjustment expression, adjusted to stop negative values - sat_adj_expr = (self.water_v - self.saturation_curve) / self.tau - sat_adj_expr = conditional(sat_adj_expr < 0, - max_value(sat_adj_expr, - -self.cloud / self.tau), - min_value(sat_adj_expr, - self.water_v / self.tau)) - - # If gamma_v depends on variables - if self.time_varying_gamma_v: - if isinstance(gamma_v, FunctionType): - self.gamma_v_computation = gamma_v - self.gamma_v = Function(Vv) - else: - raise NotImplementedError( - "If time_varying_thermal_feedback is True then gamma_v must be a Python function of at least one prognostic field.") - else: - assert not isinstance(gamma_v, FunctionType), "If time_varying_thermal_feedback is not True then gamma_v cannot be a Python function." - self.gamma_v = gamma_v - - # Factors for multiplying source for different variables - factors = [self.gamma_v, -self.gamma_v] - if convective_feedback: - factors.append(self.gamma_v*beta1) - if thermal_feedback: - factors.append(parameters.g*self.gamma_v*beta2) - - # Add terms to equations and make interpolators - self.source = [Function(Vc) for factor in factors] - self.source_interpolators = [Interpolator(sat_adj_expr*factor, source) - for factor, source in zip(factors, self.source)] - - tests = [equation.tests[idx] for idx in V_idxs] - - # Add source terms to residual - for test, source in zip(tests, self.source): - equation.residual += self.label(subject(test * source * dx, - equation.X), self.evaluate) - - def evaluate(self, x_in, dt): - """ - Evaluates the source term generated by the physics. - - Computes the phyiscs contributions to water vapour and cloud water at - each timestep. - - 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. - """ - logger.info(f'Evaluating physics parametrisation {self.label.label}') - if self.convective_feedback: - self.D.assign(x_in.subfunctions[self.VD_idx]) - if self.thermal_feedback: - self.b.assign(x_in.subfunctions[self.Vb_idx]) - if self.time_varying_saturation: - self.saturation_curve.interpolate(self.saturation_computation(x_in)) - if self.set_tau_to_dt: - self.tau.assign(dt) - self.water_v.assign(x_in.subfunctions[self.Vv_idx]) - self.cloud.assign(x_in.subfunctions[self.Vc_idx]) - if self.time_varying_gamma_v: - self.gamma_v.interpolate(self.gamma_v_computation(x_in)) - for interpolator in self.source_interpolators: - interpolator.interpolate() - - -class SurfaceFluxes(PhysicsParametrisation): - """ - Prescribed surface temperature and moisture fluxes, to be used in aquaplanet - simulations as Sea Surface Temperature fluxes. This formulation is taken - from the DCMIP (2016) test case document. - - These can be used either with an in-built implicit formulation, or with a - generic time discretisation. - - Written to assume that dry density is unchanged by the parametrisation. - """ - - def __init__(self, equation, T_surface_expr, vapour_name=None, - implicit_formulation=False, parameters=None): - """ - Args: - equation (:class:`PrognosticEquationSet`): the model's equation. - T_surface_expr (:class:`ufl.Expr`): the surface temperature. - vapour_name (str, optional): name of the water vapour variable. - Defaults to None, in which case no moisture fluxes are applied. - implicit_formulation (bool, optional): whether the scheme is already - put into a Backwards Euler formulation (which allows this scheme - to actually be used with a Forwards Euler or other explicit time - discretisation). Otherwise, this is formulated more generally - and can be used with any time stepper. Defaults to False. - parameters (:class:`BoundaryLayerParameters`): configuration object - giving the parameters for boundary and surface level schemes. - Defaults to None, in which case default values are used. - """ - - # -------------------------------------------------------------------- # - # Check arguments and generic initialisation - # -------------------------------------------------------------------- # - if not isinstance(equation, CompressibleEulerEquations): - raise ValueError("Surface fluxes can only be used with Compressible Euler equations") - - if vapour_name is not None: - if vapour_name not in equation.field_names: - raise ValueError(f"Field {vapour_name} does not exist in the equation set") - - if parameters is None: - parameters = BoundaryLayerParameters() - - label_name = 'surface_flux' - super().__init__(equation, label_name, parameters=None) - - self.implicit_formulation = implicit_formulation - self.X = Function(equation.X.function_space()) - self.dt = Constant(0.0) - - # -------------------------------------------------------------------- # - # Extract prognostic variables - # -------------------------------------------------------------------- # - u_idx = equation.field_names.index('u') - T_idx = equation.field_names.index('theta') - rho_idx = equation.field_names.index('rho') - if vapour_name is not None: - m_v_idx = equation.field_names.index(vapour_name) - - X = self.X - tests = TestFunctions(X.function_space()) if implicit_formulation else equation.tests - - u = split(X)[u_idx] - rho = split(X)[rho_idx] - theta_vd = split(X)[T_idx] - test_theta = tests[T_idx] - - if vapour_name is not None: - m_v = split(X)[m_v_idx] - test_m_v = tests[m_v_idx] - else: - m_v = None - - if implicit_formulation: - # Need to evaluate rho at theta-points - boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None - rho_averaged = Function(equation.function_space.sub(T_idx)) - self.rho_recoverer = Recoverer(rho, rho_averaged, boundary_method=boundary_method) - exner = thermodynamics.exner_pressure(equation.parameters, rho_averaged, theta_vd) - else: - # Exner is more general expression - exner = thermodynamics.exner_pressure(equation.parameters, rho, theta_vd) - - # Alternative variables - T = thermodynamics.T(equation.parameters, theta_vd, exner, r_v=m_v) - p = thermodynamics.p(equation.parameters, exner) - - # -------------------------------------------------------------------- # - # Expressions for surface fluxes - # -------------------------------------------------------------------- # - z = equation.domain.height_above_surface - z_a = parameters.height_surface_layer - surface_expr = conditional(z < z_a, Constant(1.0), Constant(0.0)) - - u_hori = u - equation.domain.k*dot(u, equation.domain.k) - u_hori_mag = sqrt(dot(u_hori, u_hori)) - - C_H = parameters.coeff_heat - C_E = parameters.coeff_evap - - # Implicit formulation ----------------------------------------------- # - # For use with ForwardEuler only, as implicit solution is hand-written - if implicit_formulation: - self.source_interpolators = [] - - # First specify T_np1 expression - Vtheta = equation.spaces[T_idx] - T_np1_expr = ((T + C_H*u_hori_mag*T_surface_expr*self.dt/z_a) - / (1 + C_H*u_hori_mag*self.dt/z_a)) - - # If moist formulation, determine next vapour value - if vapour_name is not None: - source_mv = Function(Vtheta) - mv_sat = thermodynamics.r_sat(equation.parameters, T, p) - mv_np1_expr = ((m_v + C_E*u_hori_mag*mv_sat*self.dt/z_a) - / (1 + C_E*u_hori_mag*self.dt/z_a)) - dmv_expr = surface_expr * (mv_np1_expr - m_v) / self.dt - source_mv_expr = test_m_v * source_mv * dx - - self.source_interpolators.append(Interpolator(dmv_expr, source_mv)) - equation.residual -= self.label(subject(prognostic(source_mv_expr, vapour_name), - X), self.evaluate) - - # Moisture needs including in theta_vd expression - # NB: still using old pressure here, which implies constant p? - epsilon = equation.parameters.R_d / equation.parameters.R_v - theta_np1_expr = (thermodynamics.theta(equation.parameters, T_np1_expr, p) - * (1 + mv_np1_expr / epsilon)) - - else: - theta_np1_expr = thermodynamics.theta(equation.parameters, T_np1_expr, p) - - source_theta_vd = Function(Vtheta) - dtheta_vd_expr = surface_expr * (theta_np1_expr - theta_vd) / self.dt - source_theta_expr = test_theta * source_theta_vd * dx - self.source_interpolators.append(Interpolator(dtheta_vd_expr, source_theta_vd)) - equation.residual -= self.label(subject(prognostic(source_theta_expr, 'theta'), - X), self.evaluate) - - # General formulation ------------------------------------------------ # - else: - # Construct underlying expressions - kappa = equation.parameters.kappa - dT_dt = surface_expr * C_H * u_hori_mag * (T_surface_expr - T) / z_a - - if vapour_name is not None: - mv_sat = thermodynamics.r_sat(equation.parameters, T, p) - dmv_dt = surface_expr * C_E * u_hori_mag * (mv_sat - m_v) / z_a - source_mv_expr = test_m_v * dmv_dt * dx - equation.residual -= self.label( - prognostic(subject(source_mv_expr, X), - vapour_name), self.evaluate) - - # Theta expression depends on dmv_dt - epsilon = equation.parameters.R_d / equation.parameters.R_v - dtheta_vd_dt = (dT_dt * ((1 + m_v / epsilon) / exner - kappa * theta_vd / (rho * T)) - + dmv_dt * (T / (epsilon * exner) - kappa * theta_vd / (epsilon + m_v))) - else: - dtheta_vd_dt = dT_dt * (1 / exner - kappa * theta_vd / (rho * T)) - - dx_reduced = dx(degree=4) - source_theta_expr = test_theta * dtheta_vd_dt * dx_reduced - - equation.residual -= self.label( - subject(prognostic(source_theta_expr, 'theta'), 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. - """ - - logger.info(f'Evaluating physics parametrisation {self.label.label}') - - if self.implicit_formulation: - self.X.assign(x_in) - self.dt.assign(dt) - self.rho_recoverer.project() - for source_interpolator in self.source_interpolators: - source_interpolator.interpolate() - - -class WindDrag(PhysicsParametrisation): - """ - A simple surface wind drag scheme. This formulation is taken from the DCMIP - (2016) test case document. - - These can be used either with an in-built implicit formulation, or with a - generic time discretisation. - """ - - def __init__(self, equation, implicit_formulation=False, parameters=None): - """ - Args: - equation (:class:`PrognosticEquationSet`): the model's equation. - implicit_formulation (bool, optional): whether the scheme is already - put into a Backwards Euler formulation (which allows this scheme - to actually be used with a Forwards Euler or other explicit time - discretisation). Otherwise, this is formulated more generally - and can be used with any time stepper. Defaults to False. - parameters (:class:`BoundaryLayerParameters`): configuration object - giving the parameters for boundary and surface level schemes. - Defaults to None, in which case default values are used. - """ - - # -------------------------------------------------------------------- # - # Check arguments and generic initialisation - # -------------------------------------------------------------------- # - if not isinstance(equation, CompressibleEulerEquations): - raise ValueError("Wind drag can only be used with Compressible Euler equations") - - if parameters is None: - parameters = BoundaryLayerParameters() - - label_name = 'wind_drag' - super().__init__(equation, label_name, parameters=None) - - k = equation.domain.k - self.implicit_formulation = implicit_formulation - self.X = Function(equation.X.function_space()) - self.dt = Constant(0.0) - - # -------------------------------------------------------------------- # - # Extract prognostic variables - # -------------------------------------------------------------------- # - u_idx = equation.field_names.index('u') - - X = self.X - tests = TestFunctions(X.function_space()) if implicit_formulation else equation.tests - - test = tests[u_idx] - - u = split(X)[u_idx] - u_hori = u - k*dot(u, k) - u_hori_mag = sqrt(dot(u_hori, u_hori)) - - # -------------------------------------------------------------------- # - # Expressions for wind drag - # -------------------------------------------------------------------- # - z = equation.domain.height_above_surface - z_a = parameters.height_surface_layer - surface_expr = conditional(z < z_a, Constant(1.0), Constant(0.0)) - - C_D0 = parameters.coeff_drag_0 - C_D1 = parameters.coeff_drag_1 - C_D2 = parameters.coeff_drag_2 - - C_D = conditional(u_hori_mag < 20.0, C_D0 + C_D1*u_hori_mag, C_D2) - - # Implicit formulation ----------------------------------------------- # - # For use with ForwardEuler only, as implicit solution is hand-written - if implicit_formulation: - - # First specify T_np1 expression - Vu = equation.spaces[u_idx] - source_u = Function(Vu) - u_np1_expr = u_hori / (1 + C_D*u_hori_mag*self.dt/z_a) - - du_expr = surface_expr * (u_np1_expr - u_hori) / self.dt - - # TODO: introduce reduced projector - test_Vu = TestFunction(Vu) - dx_reduced = dx(degree=4) - proj_eqn = inner(test_Vu, source_u - du_expr)*dx_reduced - proj_prob = NonlinearVariationalProblem(proj_eqn, source_u) - self.source_projector = NonlinearVariationalSolver(proj_prob) - - source_expr = inner(test, source_u - k*dot(source_u, k)) * dx - equation.residual -= self.label(subject(prognostic(source_expr, 'u'), - X), self.evaluate) - - # General formulation ------------------------------------------------ # - else: - # Construct underlying expressions - du_dt = -surface_expr * C_D * u_hori_mag * u_hori / z_a - - dx_reduced = dx(degree=4) - source_expr = inner(test, du_dt) * dx_reduced - - equation.residual -= self.label(subject(prognostic(source_expr, '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. - """ - - logger.info(f'Evaluating physics parametrisation {self.label.label}') - - if self.implicit_formulation: - self.X.assign(x_in) - self.dt.assign(dt) - self.source_projector.solve() - - -class StaticAdjustment(PhysicsParametrisation): - """ - A scheme to provide static adjustment, by sorting the potential temperature - values in a column so that they are increasing with height. - """ - - def __init__(self, equation, theta_variable='theta_vd'): - """ - Args: - equation (:class:`PrognosticEquationSet`): the model's equation. - theta_variable (str, optional): which theta variable to sort the - profile for. Valid options are "theta" or "theta_vd". Defaults - to "theta_vd". - """ - - self.explicit_only = True - from functools import partial - - # -------------------------------------------------------------------- # - # Check arguments and generic initialisation - # -------------------------------------------------------------------- # - if not isinstance(equation, CompressibleEulerEquations): - raise ValueError("Static adjustment can only be used with Compressible Euler equations") - - if theta_variable not in ['theta', 'theta_vd']: - raise ValueError('Static adjustment: theta variable ' - + f'{theta_variable} not valid') - - label_name = 'static_adjustment' - super().__init__(equation, label_name, parameters=equation.parameters) - - self.X = Function(equation.X.function_space()) - self.dt = Constant(0.0) - - # -------------------------------------------------------------------- # - # Extract prognostic variables - # -------------------------------------------------------------------- # - - theta_idx = equation.field_names.index('theta') - Vt = equation.spaces[theta_idx] - self.theta_to_sort = Function(Vt) - sorted_theta = Function(Vt) - theta = self.X.subfunctions[theta_idx] - - if theta_variable == 'theta' and 'water_vapour' in equation.field_names: - Rv = equation.parameters.R_v - Rd = equation.parameters.R_d - mv_idx = equation.field_names.index('water_vapour') - mv = self.X.subfunctions[mv_idx] - self.get_theta_variable = Interpolator(theta / (1 + mv*Rv/Rd), self.theta_to_sort) - self.set_theta_variable = Interpolator(self.theta_to_sort * (1 + mv*Rv/Rd), sorted_theta) - else: - self.get_theta_variable = Interpolator(theta, self.theta_to_sort) - self.set_theta_variable = Interpolator(self.theta_to_sort, sorted_theta) - - # -------------------------------------------------------------------- # - # Set up routines to reshape data - # -------------------------------------------------------------------- # - - domain = equation.domain - self.get_column_data = partial(domain.coords.get_column_data, - field=self.theta_to_sort, - domain=domain) - self.set_column_data = domain.coords.set_field_from_column_data - - # -------------------------------------------------------------------- # - # Set source term - # -------------------------------------------------------------------- # - - tests = TestFunctions(self.X.function_space()) - test = tests[theta_idx] - - source_expr = inner(test, sorted_theta - theta) / self.dt * dx - - equation.residual -= self.label(subject(prognostic(source_expr, 'theta'), equation.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. - """ - - logger.info(f'Evaluating physics parametrisation {self.label.label}') - - self.X.assign(x_in) - self.dt.assign(dt) - - self.get_theta_variable.interpolate() - theta_column_data, index_data = self.get_column_data() - for col in range(theta_column_data.shape[0]): - theta_column_data[col].sort() - self.set_column_data(self.theta_to_sort, theta_column_data, index_data) - self.set_theta_variable.interpolate() - - -class SuppressVerticalWind(PhysicsParametrisation): - """ - Suppresses the model's vertical wind, reducing it to zero. This is used for - instance in a model's spin up period. - """ - - def __init__(self, equation, spin_up_period): - """ - Args: - equation (:class:`PrognosticEquationSet`): the model's equation. - spin_up_period (`ufl.Constant`): this parametrisation is applied - while the time is less than this -- corresponding to a spin up - period. - """ - - self.explicit_only = True - - # -------------------------------------------------------------------- # - # Check arguments and generic initialisation - # -------------------------------------------------------------------- # - - domain = equation.domain - - if not domain.mesh.extruded: - raise RuntimeError("Suppress vertical wind can only be used with " - + "extruded meshes") - - label_name = 'suppress_vertical_wind' - super().__init__(equation, label_name, parameters=equation.parameters) - - self.X = Function(equation.X.function_space()) - self.dt = Constant(0.0) - self.t = domain.t - self.spin_up_period = Constant(spin_up_period) - self.strength = Constant(1.0) - self.spin_up_done = False - - # -------------------------------------------------------------------- # - # Extract prognostic variables - # -------------------------------------------------------------------- # - - u_idx = equation.field_names.index('u') - wind = self.X.subfunctions[u_idx] - - # -------------------------------------------------------------------- # - # Set source term - # -------------------------------------------------------------------- # - - tests = TestFunctions(self.X.function_space()) - test = tests[u_idx] - - # The sink should be just the value of the current vertical wind - source_expr = -self.strength * inner(test, domain.k*dot(domain.k, wind)) / self.dt * dx - - equation.residual -= self.label(subject(prognostic(source_expr, 'u'), equation.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. - """ - - if float(self.t) < float(self.spin_up_period): - logger.info(f'Evaluating physics parametrisation {self.label.label}') - - self.X.assign(x_in) - self.dt.assign(dt) - - elif not self.spin_up_done: - self.strength.assign(0.0) - self.spin_up_done = True - - -class BoundaryLayerMixing(PhysicsParametrisation): - """ - A simple boundary layer mixing scheme. This acts like a vertical diffusion, - using an interior penalty method. - """ - - def __init__(self, equation, field_name, parameters=None): - """ - Args: - equation (:class:`PrognosticEquationSet`): the model's equation. - field_name (str): name of the field to apply the diffusion to. - ibp (:class:`IntegrateByParts`, optional): how many times to - integrate the term by parts. Defaults to IntegrateByParts.ONCE. - parameters (:class:`BoundaryLayerParameters`): configuration object - giving the parameters for boundary and surface level schemes. - Defaults to None, in which case default values are used. - """ - - # -------------------------------------------------------------------- # - # Check arguments and generic initialisation - # -------------------------------------------------------------------- # - - if not isinstance(equation, CompressibleEulerEquations): - raise ValueError("Boundary layer mixing can only be used with Compressible Euler equations") - - if field_name not in equation.field_names: - raise ValueError(f'field {field_name} not found in equation') - - if parameters is None: - parameters = BoundaryLayerParameters() - - label_name = f'boundary_layer_mixing_{field_name}' - super().__init__(equation, label_name, parameters=None) - - self.X = Function(equation.X.function_space()) - - # -------------------------------------------------------------------- # - # Extract prognostic variables - # -------------------------------------------------------------------- # - - u_idx = equation.field_names.index('u') - T_idx = equation.field_names.index('theta') - rho_idx = equation.field_names.index('rho') - - u = split(self.X)[u_idx] - rho = split(self.X)[rho_idx] - theta_vd = split(self.X)[T_idx] - - boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None - rho_averaged = Function(equation.function_space.sub(T_idx)) - self.rho_recoverer = Recoverer(rho, rho_averaged, boundary_method=boundary_method) - exner = thermodynamics.exner_pressure(equation.parameters, rho_averaged, theta_vd) - - # Alternative variables - p = thermodynamics.p(equation.parameters, exner) - p_top = Constant(85000.) - p_strato = Constant(10000.) - - # -------------------------------------------------------------------- # - # Expressions for diffusivity coefficients - # -------------------------------------------------------------------- # - z_a = parameters.height_surface_layer - - domain = equation.domain - u_hori = u - domain.k*dot(u, domain.k) - u_hori_mag = sqrt(dot(u_hori, u_hori)) - - if field_name == 'u': - C_D0 = parameters.coeff_drag_0 - C_D1 = parameters.coeff_drag_1 - C_D2 = parameters.coeff_drag_2 - - C_D = conditional(u_hori_mag < 20.0, C_D0 + C_D1*u_hori_mag, C_D2) - K = conditional(p > p_top, - C_D*u_hori_mag*z_a, - C_D*u_hori_mag*z_a*exp(-((p_top - p)/p_strato)**2)) - - else: - C_E = parameters.coeff_evap - K = conditional(p > p_top, - C_E*u_hori_mag*z_a, - C_E*u_hori_mag*z_a*exp(-((p_top - p)/p_strato)**2)) - - # -------------------------------------------------------------------- # - # Make source expression - # -------------------------------------------------------------------- # - - dx_reduced = dx(degree=4) - dS_v_reduced = dS_v(degree=4) - # Need to be careful with order of operations here, to correctly index - # when the field is vector-valued - d_dz = lambda q: outer(domain.k, dot(grad(q), domain.k)) - n = FacetNormal(domain.mesh) - # Work out vertical height - xyz = SpatialCoordinate(domain.mesh) - Vr = domain.spaces('L2') - dz = Function(Vr) - dz.interpolate(dot(domain.k, d_dz(dot(domain.k, xyz)))) - mu = parameters.mu - - X = self.X - tests = equation.tests - - idx = equation.field_names.index(field_name) - test = tests[idx] - field = X.subfunctions[idx] - - if field_name == 'u': - # Horizontal diffusion only - field = field - domain.k*dot(field, domain.k) - - # Interior penalty discretisation of vertical diffusion - source_expr = ( - # Volume term - rho_averaged*K*inner(d_dz(test/rho_averaged), d_dz(field))*dx_reduced - # Interior penalty surface term - - 2*inner(avg(outer(K*field, n)), avg(d_dz(test)))*dS_v_reduced - - 2*inner(avg(outer(test, n)), avg(d_dz(K*field)))*dS_v_reduced - + 4*mu*avg(dz)*inner(avg(outer(K*field, n)), avg(outer(test, n)))*dS_v_reduced - ) - - equation.residual += self.label( - subject(prognostic(source_expr, field_name), X), self.evaluate) - - def evaluate(self, x_in, dt): - """ - Evaluates the source term generated by the physics. This only recovers - the density field. - - 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. - """ - - logger.info(f'Evaluating physics parametrisation {self.label.label}') - - self.X.assign(x_in) - self.rho_recoverer.project() - - -class TerminatorToy(PhysicsParametrisation): - """ - Setup the Terminator Toy chemistry interaction - as specified in 'The terminator toy chemistry test ...' - Lauritzen et. al. (2014). - - The coupled equations for the two species are given by: - - D/Dt (X) = 2Kx - D/Dt (X2) = -Kx - - where Kx = k1*X2 - k2*(X**2) - """ - - def __init__(self, equation, k1=1, k2=1, - species1_name='X', species2_name='X2'): - """ - Args: - equation (:class: 'PrognosticEquationSet'): the model's equation - k1(float, optional): Reaction rate for species 1 (X). Defaults to a constant - over the domain. - k2(float, optional): Reaction rate for species 2 (X2). Defaults to a constant - over the domain. - species1_name(str, optional): Name of the first interacting species. Defaults - to 'X'. - species2_name(str, optional): Name of the second interacting species. Defaults - to 'X2'. - """ - - label_name = 'terminator_toy' - super().__init__(equation, label_name, parameters=None) - - if species1_name not in equation.field_names: - raise ValueError(f"Field {species1_name} does not exist in the equation set") - if species2_name not in equation.field_names: - raise ValueError(f"Field {species2_name} does not exist in the equation set") - - self.species1_idx = equation.field_names.index(species1_name) - self.species2_idx = equation.field_names.index(species2_name) - - assert equation.function_space.sub(self.species1_idx) == equation.function_space.sub(self.species2_idx), "The function spaces for the two species need to be the same" - - self.Xq = Function(equation.X.function_space()) - Xq = self.Xq - - species1 = split(Xq)[self.species1_idx] - species2 = split(Xq)[self.species2_idx] - - test_1 = equation.tests[self.species1_idx] - test_2 = equation.tests[self.species2_idx] - - Kx = k1*species2 - k2*(species1**2) - - source1_expr = test_1 * 2*Kx * dx - source2_expr = test_2 * -Kx * dx - - equation.residual -= self.label(subject(prognostic(source1_expr, 'X'), Xq), self.evaluate) - equation.residual -= self.label(subject(prognostic(source2_expr, 'X2'), Xq), self.evaluate) - - def evaluate(self, x_in, dt): - """ - Evaluates the source/sink for the coalescence process. - - Args: - x_in (:class:`Function`): the (mixed) field to be evolved. - dt (:class:`Constant`): the time interval for the scheme. - """ - - logger.info(f'Evaluating physics parametrisation {self.label.label}') - - pass diff --git a/gusto/physics/__init__.py b/gusto/physics/__init__.py new file mode 100644 index 000000000..91f09c2ec --- /dev/null +++ b/gusto/physics/__init__.py @@ -0,0 +1,5 @@ +from gusto.physics.physics_parametrisation import * # noqa +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 \ No newline at end of file diff --git a/gusto/physics/boundary_and_turbulence.py b/gusto/physics/boundary_and_turbulence.py new file mode 100644 index 000000000..317cc4ac1 --- /dev/null +++ b/gusto/physics/boundary_and_turbulence.py @@ -0,0 +1,645 @@ +""" +Objects to describe physics parametrisations for the boundary layer, such as +drag and turbulence.""" + +from firedrake import ( + Interpolator, conditional, Function, dx, sqrt, dot, Constant, grad, + TestFunctions, split, inner, TestFunction, exp, avg, outer, FacetNormal, + SpatialCoordinate, dS_v, NonlinearVariationalProblem, + NonlinearVariationalSolver +) +from firedrake.fml import subject +from gusto.core.configuration import BoundaryLayerParameters +from gusto.recovery import Recoverer, BoundaryMethod +from gusto.equations import CompressibleEulerEquations +from gusto.core.labels import prognostic +from gusto.core.logging import logger +from gusto.equations import thermodynamics +from gusto.physics.physics_parametrisation import PhysicsParametrisation + +__all__ = ["SurfaceFluxes", "WindDrag", "StaticAdjustment", + "SuppressVerticalWind", "BoundaryLayerMixing"] + + +class SurfaceFluxes(PhysicsParametrisation): + """ + Prescribed surface temperature and moisture fluxes, to be used in aquaplanet + simulations as Sea Surface Temperature fluxes. This formulation is taken + from the DCMIP (2016) test case document. + + These can be used either with an in-built implicit formulation, or with a + generic time discretisation. + + Written to assume that dry density is unchanged by the parametrisation. + """ + + def __init__(self, equation, T_surface_expr, vapour_name=None, + implicit_formulation=False, parameters=None): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + T_surface_expr (:class:`ufl.Expr`): the surface temperature. + vapour_name (str, optional): name of the water vapour variable. + Defaults to None, in which case no moisture fluxes are applied. + implicit_formulation (bool, optional): whether the scheme is already + put into a Backwards Euler formulation (which allows this scheme + to actually be used with a Forwards Euler or other explicit time + discretisation). Otherwise, this is formulated more generally + and can be used with any time stepper. Defaults to False. + parameters (:class:`BoundaryLayerParameters`): configuration object + giving the parameters for boundary and surface level schemes. + Defaults to None, in which case default values are used. + """ + + # -------------------------------------------------------------------- # + # Check arguments and generic initialisation + # -------------------------------------------------------------------- # + if not isinstance(equation, CompressibleEulerEquations): + raise ValueError("Surface fluxes can only be used with Compressible Euler equations") + + if vapour_name is not None: + if vapour_name not in equation.field_names: + raise ValueError(f"Field {vapour_name} does not exist in the equation set") + + if parameters is None: + parameters = BoundaryLayerParameters() + + label_name = 'surface_flux' + super().__init__(equation, label_name, parameters=None) + + self.implicit_formulation = implicit_formulation + self.X = Function(equation.X.function_space()) + self.dt = Constant(0.0) + + # -------------------------------------------------------------------- # + # Extract prognostic variables + # -------------------------------------------------------------------- # + u_idx = equation.field_names.index('u') + T_idx = equation.field_names.index('theta') + rho_idx = equation.field_names.index('rho') + if vapour_name is not None: + m_v_idx = equation.field_names.index(vapour_name) + + X = self.X + tests = TestFunctions(X.function_space()) if implicit_formulation else equation.tests + + u = split(X)[u_idx] + rho = split(X)[rho_idx] + theta_vd = split(X)[T_idx] + test_theta = tests[T_idx] + + if vapour_name is not None: + m_v = split(X)[m_v_idx] + test_m_v = tests[m_v_idx] + else: + m_v = None + + if implicit_formulation: + # Need to evaluate rho at theta-points + boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None + rho_averaged = Function(equation.function_space.sub(T_idx)) + self.rho_recoverer = Recoverer(rho, rho_averaged, boundary_method=boundary_method) + exner = thermodynamics.exner_pressure(equation.parameters, rho_averaged, theta_vd) + else: + # Exner is more general expression + exner = thermodynamics.exner_pressure(equation.parameters, rho, theta_vd) + + # Alternative variables + T = thermodynamics.T(equation.parameters, theta_vd, exner, r_v=m_v) + p = thermodynamics.p(equation.parameters, exner) + + # -------------------------------------------------------------------- # + # Expressions for surface fluxes + # -------------------------------------------------------------------- # + z = equation.domain.height_above_surface + z_a = parameters.height_surface_layer + surface_expr = conditional(z < z_a, Constant(1.0), Constant(0.0)) + + u_hori = u - equation.domain.k*dot(u, equation.domain.k) + u_hori_mag = sqrt(dot(u_hori, u_hori)) + + C_H = parameters.coeff_heat + C_E = parameters.coeff_evap + + # Implicit formulation ----------------------------------------------- # + # For use with ForwardEuler only, as implicit solution is hand-written + if implicit_formulation: + self.source_interpolators = [] + + # First specify T_np1 expression + Vtheta = equation.spaces[T_idx] + T_np1_expr = ((T + C_H*u_hori_mag*T_surface_expr*self.dt/z_a) + / (1 + C_H*u_hori_mag*self.dt/z_a)) + + # If moist formulation, determine next vapour value + if vapour_name is not None: + source_mv = Function(Vtheta) + mv_sat = thermodynamics.r_sat(equation.parameters, T, p) + mv_np1_expr = ((m_v + C_E*u_hori_mag*mv_sat*self.dt/z_a) + / (1 + C_E*u_hori_mag*self.dt/z_a)) + dmv_expr = surface_expr * (mv_np1_expr - m_v) / self.dt + source_mv_expr = test_m_v * source_mv * dx + + self.source_interpolators.append(Interpolator(dmv_expr, source_mv)) + equation.residual -= self.label(subject(prognostic(source_mv_expr, vapour_name), + X), self.evaluate) + + # Moisture needs including in theta_vd expression + # NB: still using old pressure here, which implies constant p? + epsilon = equation.parameters.R_d / equation.parameters.R_v + theta_np1_expr = (thermodynamics.theta(equation.parameters, T_np1_expr, p) + * (1 + mv_np1_expr / epsilon)) + + else: + theta_np1_expr = thermodynamics.theta(equation.parameters, T_np1_expr, p) + + source_theta_vd = Function(Vtheta) + dtheta_vd_expr = surface_expr * (theta_np1_expr - theta_vd) / self.dt + source_theta_expr = test_theta * source_theta_vd * dx + self.source_interpolators.append(Interpolator(dtheta_vd_expr, source_theta_vd)) + equation.residual -= self.label(subject(prognostic(source_theta_expr, 'theta'), + X), self.evaluate) + + # General formulation ------------------------------------------------ # + else: + # Construct underlying expressions + kappa = equation.parameters.kappa + dT_dt = surface_expr * C_H * u_hori_mag * (T_surface_expr - T) / z_a + + if vapour_name is not None: + mv_sat = thermodynamics.r_sat(equation.parameters, T, p) + dmv_dt = surface_expr * C_E * u_hori_mag * (mv_sat - m_v) / z_a + source_mv_expr = test_m_v * dmv_dt * dx + equation.residual -= self.label( + prognostic(subject(source_mv_expr, X), + vapour_name), self.evaluate) + + # Theta expression depends on dmv_dt + epsilon = equation.parameters.R_d / equation.parameters.R_v + dtheta_vd_dt = (dT_dt * ((1 + m_v / epsilon) / exner - kappa * theta_vd / (rho * T)) + + dmv_dt * (T / (epsilon * exner) - kappa * theta_vd / (epsilon + m_v))) + else: + dtheta_vd_dt = dT_dt * (1 / exner - kappa * theta_vd / (rho * T)) + + dx_reduced = dx(degree=4) + source_theta_expr = test_theta * dtheta_vd_dt * dx_reduced + + equation.residual -= self.label( + subject(prognostic(source_theta_expr, 'theta'), 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. + """ + + logger.info(f'Evaluating physics parametrisation {self.label.label}') + + if self.implicit_formulation: + self.X.assign(x_in) + self.dt.assign(dt) + self.rho_recoverer.project() + for source_interpolator in self.source_interpolators: + source_interpolator.interpolate() + + +class WindDrag(PhysicsParametrisation): + """ + A simple surface wind drag scheme. This formulation is taken from the DCMIP + (2016) test case document. + + These can be used either with an in-built implicit formulation, or with a + generic time discretisation. + """ + + def __init__(self, equation, implicit_formulation=False, parameters=None): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + implicit_formulation (bool, optional): whether the scheme is already + put into a Backwards Euler formulation (which allows this scheme + to actually be used with a Forwards Euler or other explicit time + discretisation). Otherwise, this is formulated more generally + and can be used with any time stepper. Defaults to False. + parameters (:class:`BoundaryLayerParameters`): configuration object + giving the parameters for boundary and surface level schemes. + Defaults to None, in which case default values are used. + """ + + # -------------------------------------------------------------------- # + # Check arguments and generic initialisation + # -------------------------------------------------------------------- # + if not isinstance(equation, CompressibleEulerEquations): + raise ValueError("Wind drag can only be used with Compressible Euler equations") + + if parameters is None: + parameters = BoundaryLayerParameters() + + label_name = 'wind_drag' + super().__init__(equation, label_name, parameters=None) + + k = equation.domain.k + self.implicit_formulation = implicit_formulation + self.X = Function(equation.X.function_space()) + self.dt = Constant(0.0) + + # -------------------------------------------------------------------- # + # Extract prognostic variables + # -------------------------------------------------------------------- # + u_idx = equation.field_names.index('u') + + X = self.X + tests = TestFunctions(X.function_space()) if implicit_formulation else equation.tests + + test = tests[u_idx] + + u = split(X)[u_idx] + u_hori = u - k*dot(u, k) + u_hori_mag = sqrt(dot(u_hori, u_hori)) + + # -------------------------------------------------------------------- # + # Expressions for wind drag + # -------------------------------------------------------------------- # + z = equation.domain.height_above_surface + z_a = parameters.height_surface_layer + surface_expr = conditional(z < z_a, Constant(1.0), Constant(0.0)) + + C_D0 = parameters.coeff_drag_0 + C_D1 = parameters.coeff_drag_1 + C_D2 = parameters.coeff_drag_2 + + C_D = conditional(u_hori_mag < 20.0, C_D0 + C_D1*u_hori_mag, C_D2) + + # Implicit formulation ----------------------------------------------- # + # For use with ForwardEuler only, as implicit solution is hand-written + if implicit_formulation: + + # First specify T_np1 expression + Vu = equation.spaces[u_idx] + source_u = Function(Vu) + u_np1_expr = u_hori / (1 + C_D*u_hori_mag*self.dt/z_a) + + du_expr = surface_expr * (u_np1_expr - u_hori) / self.dt + + # TODO: introduce reduced projector + test_Vu = TestFunction(Vu) + dx_reduced = dx(degree=4) + proj_eqn = inner(test_Vu, source_u - du_expr)*dx_reduced + proj_prob = NonlinearVariationalProblem(proj_eqn, source_u) + self.source_projector = NonlinearVariationalSolver(proj_prob) + + source_expr = inner(test, source_u - k*dot(source_u, k)) * dx + equation.residual -= self.label(subject(prognostic(source_expr, 'u'), + X), self.evaluate) + + # General formulation ------------------------------------------------ # + else: + # Construct underlying expressions + du_dt = -surface_expr * C_D * u_hori_mag * u_hori / z_a + + dx_reduced = dx(degree=4) + source_expr = inner(test, du_dt) * dx_reduced + + equation.residual -= self.label(subject(prognostic(source_expr, '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. + """ + + logger.info(f'Evaluating physics parametrisation {self.label.label}') + + if self.implicit_formulation: + self.X.assign(x_in) + self.dt.assign(dt) + self.source_projector.solve() + + +class StaticAdjustment(PhysicsParametrisation): + """ + A scheme to provide static adjustment, by sorting the potential temperature + values in a column so that they are increasing with height. + """ + + def __init__(self, equation, theta_variable='theta_vd'): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + theta_variable (str, optional): which theta variable to sort the + profile for. Valid options are "theta" or "theta_vd". Defaults + to "theta_vd". + """ + + self.explicit_only = True + from functools import partial + + # -------------------------------------------------------------------- # + # Check arguments and generic initialisation + # -------------------------------------------------------------------- # + if not isinstance(equation, CompressibleEulerEquations): + raise ValueError("Static adjustment can only be used with Compressible Euler equations") + + if theta_variable not in ['theta', 'theta_vd']: + raise ValueError('Static adjustment: theta variable ' + + f'{theta_variable} not valid') + + label_name = 'static_adjustment' + super().__init__(equation, label_name, parameters=equation.parameters) + + self.X = Function(equation.X.function_space()) + self.dt = Constant(0.0) + + # -------------------------------------------------------------------- # + # Extract prognostic variables + # -------------------------------------------------------------------- # + + theta_idx = equation.field_names.index('theta') + Vt = equation.spaces[theta_idx] + self.theta_to_sort = Function(Vt) + sorted_theta = Function(Vt) + theta = self.X.subfunctions[theta_idx] + + if theta_variable == 'theta' and 'water_vapour' in equation.field_names: + Rv = equation.parameters.R_v + Rd = equation.parameters.R_d + mv_idx = equation.field_names.index('water_vapour') + mv = self.X.subfunctions[mv_idx] + self.get_theta_variable = Interpolator(theta / (1 + mv*Rv/Rd), self.theta_to_sort) + self.set_theta_variable = Interpolator(self.theta_to_sort * (1 + mv*Rv/Rd), sorted_theta) + else: + self.get_theta_variable = Interpolator(theta, self.theta_to_sort) + self.set_theta_variable = Interpolator(self.theta_to_sort, sorted_theta) + + # -------------------------------------------------------------------- # + # Set up routines to reshape data + # -------------------------------------------------------------------- # + + domain = equation.domain + self.get_column_data = partial(domain.coords.get_column_data, + field=self.theta_to_sort, + domain=domain) + self.set_column_data = domain.coords.set_field_from_column_data + + # -------------------------------------------------------------------- # + # Set source term + # -------------------------------------------------------------------- # + + tests = TestFunctions(self.X.function_space()) + test = tests[theta_idx] + + source_expr = inner(test, sorted_theta - theta) / self.dt * dx + + equation.residual -= self.label(subject(prognostic(source_expr, 'theta'), equation.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. + """ + + logger.info(f'Evaluating physics parametrisation {self.label.label}') + + self.X.assign(x_in) + self.dt.assign(dt) + + self.get_theta_variable.interpolate() + theta_column_data, index_data = self.get_column_data() + for col in range(theta_column_data.shape[0]): + theta_column_data[col].sort() + self.set_column_data(self.theta_to_sort, theta_column_data, index_data) + self.set_theta_variable.interpolate() + + +class SuppressVerticalWind(PhysicsParametrisation): + """ + Suppresses the model's vertical wind, reducing it to zero. This is used for + instance in a model's spin up period. + """ + + def __init__(self, equation, spin_up_period): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + spin_up_period (`ufl.Constant`): this parametrisation is applied + while the time is less than this -- corresponding to a spin up + period. + """ + + self.explicit_only = True + + # -------------------------------------------------------------------- # + # Check arguments and generic initialisation + # -------------------------------------------------------------------- # + + domain = equation.domain + + if not domain.mesh.extruded: + raise RuntimeError("Suppress vertical wind can only be used with " + + "extruded meshes") + + label_name = 'suppress_vertical_wind' + super().__init__(equation, label_name, parameters=equation.parameters) + + self.X = Function(equation.X.function_space()) + self.dt = Constant(0.0) + self.t = domain.t + self.spin_up_period = Constant(spin_up_period) + self.strength = Constant(1.0) + self.spin_up_done = False + + # -------------------------------------------------------------------- # + # Extract prognostic variables + # -------------------------------------------------------------------- # + + u_idx = equation.field_names.index('u') + wind = self.X.subfunctions[u_idx] + + # -------------------------------------------------------------------- # + # Set source term + # -------------------------------------------------------------------- # + + tests = TestFunctions(self.X.function_space()) + test = tests[u_idx] + + # The sink should be just the value of the current vertical wind + source_expr = -self.strength * inner(test, domain.k*dot(domain.k, wind)) / self.dt * dx + + equation.residual -= self.label(subject(prognostic(source_expr, 'u'), equation.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. + """ + + if float(self.t) < float(self.spin_up_period): + logger.info(f'Evaluating physics parametrisation {self.label.label}') + + self.X.assign(x_in) + self.dt.assign(dt) + + elif not self.spin_up_done: + self.strength.assign(0.0) + self.spin_up_done = True + + +class BoundaryLayerMixing(PhysicsParametrisation): + """ + A simple boundary layer mixing scheme. This acts like a vertical diffusion, + using an interior penalty method. + """ + + def __init__(self, equation, field_name, parameters=None): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + field_name (str): name of the field to apply the diffusion to. + ibp (:class:`IntegrateByParts`, optional): how many times to + integrate the term by parts. Defaults to IntegrateByParts.ONCE. + parameters (:class:`BoundaryLayerParameters`): configuration object + giving the parameters for boundary and surface level schemes. + Defaults to None, in which case default values are used. + """ + + # -------------------------------------------------------------------- # + # Check arguments and generic initialisation + # -------------------------------------------------------------------- # + + if not isinstance(equation, CompressibleEulerEquations): + raise ValueError("Boundary layer mixing can only be used with Compressible Euler equations") + + if field_name not in equation.field_names: + raise ValueError(f'field {field_name} not found in equation') + + if parameters is None: + parameters = BoundaryLayerParameters() + + label_name = f'boundary_layer_mixing_{field_name}' + super().__init__(equation, label_name, parameters=None) + + self.X = Function(equation.X.function_space()) + + # -------------------------------------------------------------------- # + # Extract prognostic variables + # -------------------------------------------------------------------- # + + u_idx = equation.field_names.index('u') + T_idx = equation.field_names.index('theta') + rho_idx = equation.field_names.index('rho') + + u = split(self.X)[u_idx] + rho = split(self.X)[rho_idx] + theta_vd = split(self.X)[T_idx] + + boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None + rho_averaged = Function(equation.function_space.sub(T_idx)) + self.rho_recoverer = Recoverer(rho, rho_averaged, boundary_method=boundary_method) + exner = thermodynamics.exner_pressure(equation.parameters, rho_averaged, theta_vd) + + # Alternative variables + p = thermodynamics.p(equation.parameters, exner) + p_top = Constant(85000.) + p_strato = Constant(10000.) + + # -------------------------------------------------------------------- # + # Expressions for diffusivity coefficients + # -------------------------------------------------------------------- # + z_a = parameters.height_surface_layer + + domain = equation.domain + u_hori = u - domain.k*dot(u, domain.k) + u_hori_mag = sqrt(dot(u_hori, u_hori)) + + if field_name == 'u': + C_D0 = parameters.coeff_drag_0 + C_D1 = parameters.coeff_drag_1 + C_D2 = parameters.coeff_drag_2 + + C_D = conditional(u_hori_mag < 20.0, C_D0 + C_D1*u_hori_mag, C_D2) + K = conditional(p > p_top, + C_D*u_hori_mag*z_a, + C_D*u_hori_mag*z_a*exp(-((p_top - p)/p_strato)**2)) + + else: + C_E = parameters.coeff_evap + K = conditional(p > p_top, + C_E*u_hori_mag*z_a, + C_E*u_hori_mag*z_a*exp(-((p_top - p)/p_strato)**2)) + + # -------------------------------------------------------------------- # + # Make source expression + # -------------------------------------------------------------------- # + + dx_reduced = dx(degree=4) + dS_v_reduced = dS_v(degree=4) + # Need to be careful with order of operations here, to correctly index + # when the field is vector-valued + d_dz = lambda q: outer(domain.k, dot(grad(q), domain.k)) + n = FacetNormal(domain.mesh) + # Work out vertical height + xyz = SpatialCoordinate(domain.mesh) + Vr = domain.spaces('L2') + dz = Function(Vr) + dz.interpolate(dot(domain.k, d_dz(dot(domain.k, xyz)))) + mu = parameters.mu + + X = self.X + tests = equation.tests + + idx = equation.field_names.index(field_name) + test = tests[idx] + field = X.subfunctions[idx] + + if field_name == 'u': + # Horizontal diffusion only + field = field - domain.k*dot(field, domain.k) + + # Interior penalty discretisation of vertical diffusion + source_expr = ( + # Volume term + rho_averaged*K*inner(d_dz(test/rho_averaged), d_dz(field))*dx_reduced + # Interior penalty surface term + - 2*inner(avg(outer(K*field, n)), avg(d_dz(test)))*dS_v_reduced + - 2*inner(avg(outer(test, n)), avg(d_dz(K*field)))*dS_v_reduced + + 4*mu*avg(dz)*inner(avg(outer(K*field, n)), avg(outer(test, n)))*dS_v_reduced + ) + + equation.residual += self.label( + subject(prognostic(source_expr, field_name), X), self.evaluate) + + def evaluate(self, x_in, dt): + """ + Evaluates the source term generated by the physics. This only recovers + the density field. + + 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. + """ + + logger.info(f'Evaluating physics parametrisation {self.label.label}') + + self.X.assign(x_in) + self.rho_recoverer.project() diff --git a/gusto/physics/chemistry.py b/gusto/physics/chemistry.py new file mode 100644 index 000000000..2304e64a3 --- /dev/null +++ b/gusto/physics/chemistry.py @@ -0,0 +1,83 @@ +"""Objects describe chemical conversion and reaction processes.""" + +from firedrake import dx, split, Function +from firedrake.fml import subject +from gusto.core.labels import prognostic +from gusto.core.logging import logger +from gusto.physics.physics_parametrisation import PhysicsParametrisation + +__all__ = ["TerminatorToy"] + + +class TerminatorToy(PhysicsParametrisation): + """ + Setup the Terminator Toy chemistry interaction + as specified in 'The terminator toy chemistry test ...' + Lauritzen et. al. (2014). + + The coupled equations for the two species are given by: + + D/Dt (X) = 2Kx + D/Dt (X2) = -Kx + + where Kx = k1*X2 - k2*(X**2) + """ + + def __init__(self, equation, k1=1, k2=1, + species1_name='X', species2_name='X2'): + """ + Args: + equation (:class: 'PrognosticEquationSet'): the model's equation + k1(float, optional): Reaction rate for species 1 (X). Defaults to a + constant 1 over the domain. + k2(float, optional): Reaction rate for species 2 (X2). Defaults to a + constant 1 over the domain. + species1_name(str, optional): Name of the first interacting species. + Defaults to 'X'. + species2_name(str, optional): Name of the second interacting + species. Defaults to 'X2'. + """ + + label_name = 'terminator_toy' + super().__init__(equation, label_name, parameters=None) + + if species1_name not in equation.field_names: + raise ValueError(f"Field {species1_name} does not exist in the equation set") + if species2_name not in equation.field_names: + raise ValueError(f"Field {species2_name} does not exist in the equation set") + + self.species1_idx = equation.field_names.index(species1_name) + self.species2_idx = equation.field_names.index(species2_name) + + assert equation.function_space.sub(self.species1_idx) == equation.function_space.sub(self.species2_idx), \ + "The function spaces for the two species need to be the same" + + self.Xq = Function(equation.X.function_space()) + Xq = self.Xq + + species1 = split(Xq)[self.species1_idx] + species2 = split(Xq)[self.species2_idx] + + test_1 = equation.tests[self.species1_idx] + test_2 = equation.tests[self.species2_idx] + + Kx = k1*species2 - k2*(species1**2) + + source1_expr = test_1 * 2*Kx * dx + source2_expr = test_2 * -Kx * dx + + equation.residual -= self.label(subject(prognostic(source1_expr, 'X'), Xq), self.evaluate) + equation.residual -= self.label(subject(prognostic(source2_expr, 'X2'), Xq), self.evaluate) + + def evaluate(self, x_in, dt): + """ + Evaluates the source/sink for the coalescence process. + + Args: + x_in (:class:`Function`): the (mixed) field to be evolved. + dt (:class:`Constant`): the time interval for the scheme. + """ + + logger.info(f'Evaluating physics parametrisation {self.label.label}') + + pass diff --git a/gusto/physics/microphysics.py b/gusto/physics/microphysics.py new file mode 100644 index 000000000..c1d3a440d --- /dev/null +++ b/gusto/physics/microphysics.py @@ -0,0 +1,640 @@ +""" +Defines microphysics routines to parametrise moist processes for the +compressible Euler equations. +""" + +from firedrake import ( + Interpolator, conditional, Function, dx, min_value, max_value, Constant, pi, + inner, TestFunction, NonlinearVariationalProblem, NonlinearVariationalSolver +) +from firedrake.fml import identity, Term, subject +from gusto.equations import Phases, TracerVariableType +from gusto.recovery import Recoverer, BoundaryMethod +from gusto.equations import CompressibleEulerEquations +from gusto.core.labels import transporting_velocity, transport, prognostic +from gusto.core.logging import logger +from gusto.equations import thermodynamics +from gusto.physics.physics_parametrisation import PhysicsParametrisation +import ufl +import math +from enum import Enum + + +__all__ = ["SaturationAdjustment", "Fallout", "Coalescence", + "EvaporationOfRain", "AdvectedMoments"] + + +class SaturationAdjustment(PhysicsParametrisation): + """ + Represents the phase change between water vapour and cloud liquid. + + This class computes updates to water vapour, cloud liquid and (virtual dry) + potential temperature, representing the action of condensation of water + vapour and/or evaporation of cloud liquid, with the associated latent heat + change. The parametrisation follows the saturation adjustment used in Bryan + and Fritsch (2002). + + Currently this is only implemented for use with mixing ratio variables, and + with "theta" assumed to be the virtual dry potential temperature. Latent + heating effects are always assumed, and the total mixing ratio is conserved. + A filter is applied to prevent generation of negative mixing ratios. + """ + + def __init__(self, equation, vapour_name='water_vapour', + cloud_name='cloud_water', latent_heat=True, parameters=None): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + vapour_name (str, optional): name of the water vapour variable. + Defaults to 'water_vapour'. + cloud_name (str, optional): name of the cloud water variable. + Defaults to 'cloud_water'. + latent_heat (bool, optional): whether to have latent heat exchange + feeding back from the phase change. Defaults to True. + parameters (:class:`Configuration`, optional): parameters containing + the values of gas constants. Defaults to None, in which case the + parameters are obtained from the equation. + + Raises: + NotImplementedError: currently this is only implemented for the + CompressibleEulerEquations. + """ + + label_name = 'saturation_adjustment' + self.explicit_only = True + super().__init__(equation, label_name, parameters=parameters) + + # TODO: make a check on the variable type of the active tracers + # if not a mixing ratio, we need to convert to mixing ratios + # this will be easier if we change equations to have dictionary of + # active tracer metadata corresponding to variable names + + # Check that fields exist + if vapour_name not in equation.field_names: + raise ValueError(f"Field {vapour_name} does not exist in the equation set") + if cloud_name not in equation.field_names: + raise ValueError(f"Field {cloud_name} does not exist in the equation set") + + # Make prognostic for physics scheme + parameters = self.parameters + self.X = Function(equation.X.function_space()) + self.latent_heat = latent_heat + + # Vapour and cloud variables are needed for every form of this scheme + cloud_idx = equation.field_names.index(cloud_name) + vap_idx = equation.field_names.index(vapour_name) + cloud_water = self.X.subfunctions[cloud_idx] + water_vapour = self.X.subfunctions[vap_idx] + + # Indices of variables in mixed function space + V_idxs = [vap_idx, cloud_idx] + V = equation.function_space.sub(vap_idx) # space in which to do the calculation + + # Get variables used to calculate saturation curve + if isinstance(equation, CompressibleEulerEquations): + rho_idx = equation.field_names.index('rho') + theta_idx = equation.field_names.index('theta') + rho = self.X.subfunctions[rho_idx] + theta = self.X.subfunctions[theta_idx] + if latent_heat: + V_idxs.append(theta_idx) + + # need to evaluate rho at theta-points, and do this via recovery + boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None + rho_averaged = Function(V) + self.rho_recoverer = Recoverer(rho, rho_averaged, boundary_method=boundary_method) + + exner = thermodynamics.exner_pressure(parameters, rho_averaged, theta) + T = thermodynamics.T(parameters, theta, exner, r_v=water_vapour) + p = thermodynamics.p(parameters, exner) + + else: + raise NotImplementedError( + 'Saturation adjustment only implemented for the Compressible Euler equations') + + # -------------------------------------------------------------------- # + # Compute heat capacities and calculate saturation curve + # -------------------------------------------------------------------- # + # Loop through variables to extract all liquid components + liquid_water = cloud_water + for active_tracer in equation.active_tracers: + if (active_tracer.phase == Phases.liquid + and active_tracer.chemical == 'H2O' and active_tracer.name != cloud_name): + liq_idx = equation.field_names.index(active_tracer.name) + liquid_water += self.X.subfunctions[liq_idx] + + # define some parameters as attributes + self.dt = Constant(0.0) + R_d = parameters.R_d + cp = parameters.cp + cv = parameters.cv + c_pv = parameters.c_pv + c_pl = parameters.c_pl + c_vv = parameters.c_vv + R_v = parameters.R_v + + # make useful fields + L_v = thermodynamics.Lv(parameters, T) + R_m = R_d + R_v * water_vapour + c_pml = cp + c_pv * water_vapour + c_pl * liquid_water + c_vml = cv + c_vv * water_vapour + c_pl * liquid_water + + # use Teten's formula to calculate the saturation curve + sat_expr = thermodynamics.r_sat(parameters, T, p) + + # -------------------------------------------------------------------- # + # Saturation adjustment expression + # -------------------------------------------------------------------- # + # make appropriate condensation rate + sat_adj_expr = (water_vapour - sat_expr) / self.dt + if latent_heat: + # As condensation/evaporation happens, the temperature changes + # so need to take this into account with an extra factor + sat_adj_expr = sat_adj_expr / (1.0 + ((L_v ** 2.0 * sat_expr) + / (cp * R_v * T ** 2.0))) + + # adjust the rate so that so negative values don't occur + sat_adj_expr = conditional(sat_adj_expr < 0, + max_value(sat_adj_expr, -cloud_water / self.dt), + min_value(sat_adj_expr, water_vapour / self.dt)) + + # -------------------------------------------------------------------- # + # Factors for multiplying source for different variables + # -------------------------------------------------------------------- # + # Factors need to have same shape as V_idxs + factors = [Constant(1.0), Constant(-1.0)] + if latent_heat and isinstance(equation, CompressibleEulerEquations): + factors.append(-theta * (cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml / (R_m * cp * c_vml))) + + # -------------------------------------------------------------------- # + # 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)] + + tests = [equation.tests[idx] for idx in V_idxs] + + # Add source terms to residual + for test, source in zip(tests, self.source): + equation.residual += self.label(subject(test * source * dx, + equation.X), self.evaluate) + + def evaluate(self, x_in, dt): + """ + Evaluates the source/sink for the saturation adjustment process. + + Args: + x_in (:class:`Function`): the (mixed) field to be evolved. + dt (:class:`Constant`): the time interval for the scheme. + """ + logger.info(f'Evaluating physics parametrisation {self.label.label}') + # Update the values of internal variables + self.dt.assign(dt) + self.X.assign(x_in) + if isinstance(self.equation, CompressibleEulerEquations): + self.rho_recoverer.project() + # Evaluate the source + for interpolator in self.source_interpolators: + interpolator.interpolate() + + +class AdvectedMoments(Enum): + """ + Enumerator describing the moments in the raindrop size distribution. + + This stores enumerators for the number of moments used to describe the + size distribution curve of raindrops. This can be used for deciding which + moments to advect in a precipitation scheme. + """ + + M0 = 0 # don't advect the distribution -- advect all rain at the same speed + M3 = 1 # advect the mass of the distribution + + +class Fallout(PhysicsParametrisation): + """ + Represents the fallout process of hydrometeors. + + Precipitation is described by downwards transport of tracer species. This + process determines the precipitation velocity from the `AdvectedMoments` + enumerator, which either: + (a) sets a terminal velocity of 5 m/s + (b) determines a rainfall size distribution based on a Gamma distribution, + as in Paluch (1979). The droplets are based on the mean mass of the rain + droplets (aka a single-moment scheme). + + Outflow boundary conditions are applied to the transport, so the rain will + flow out of the bottom of the domain. + + This is currently only implemented for "rain" in the compressible Euler + equation set. This variable must be a mixing ratio, It is only implemented + for Cartesian geometry. + """ + + def __init__(self, equation, rain_name, domain, transport_method, + moments=AdvectedMoments.M3): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + rain_name (str, optional): name of the rain variable. Defaults to + 'rain'. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + transport_method (:class:`TransportMethod`): the spatial method + used for transporting the rain. + moments (int, optional): an :class:`AdvectedMoments` enumerator, + representing the number of moments of the size distribution of + raindrops to be transported. Defaults to `AdvectedMoments.M3`. + """ + + label_name = f'fallout_{rain_name}' + super().__init__(equation, label_name, parameters=None) + + # Check that fields exist + if rain_name not in equation.field_names: + raise ValueError(f"Field {rain_name} does not exist in the equation set") + + # Check if variable is a mixing ratio + rain_tracer = equation.get_active_tracer(rain_name) + if rain_tracer.variable_type != TracerVariableType.mixing_ratio: + raise NotImplementedError('Fallout only implemented when rain ' + + 'variable is a mixing ratio') + + # Set up rain and velocity + self.X = Function(equation.X.function_space()) + + rain_idx = equation.field_names.index(rain_name) + rain = self.X.subfunctions[rain_idx] + + Vu = domain.spaces("HDiv") + # TODO: there must be a better way than forcing this into the equation + v = equation.prescribed_fields(name='rainfall_velocity', space=Vu) + + # -------------------------------------------------------------------- # + # Create physics term -- which is actually a transport term + # -------------------------------------------------------------------- # + + assert transport_method.outflow, \ + 'Fallout requires a transport method with outflow=True' + adv_term = transport_method.form + # Add rainfall velocity by replacing transport_velocity in term + adv_term = adv_term.label_map(identity, + map_if_true=lambda t: Term( + ufl.replace(t.form, {t.get(transporting_velocity): v}), + t.labels)) + + # We don't want this term to be picked up by normal transport, so drop + # the transport label + adv_term = transport.remove(adv_term) + + adv_term = prognostic(subject(adv_term, equation.X), rain_name) + equation.residual += self.label(adv_term, self.evaluate) + + # -------------------------------------------------------------------- # + # Expressions for determining rainfall velocity + # -------------------------------------------------------------------- # + self.moments = moments + + if moments == AdvectedMoments.M0: + # all rain falls at terminal velocity + terminal_velocity = Constant(5) # in m/s + v.project(-terminal_velocity*domain.k) + elif moments == AdvectedMoments.M3: + self.explicit_only = True + # this advects the third moment M3 of the raindrop + # distribution, which corresponds to the mean mass + rho_idx = equation.field_names.index('rho') + rho = self.X.subfunctions[rho_idx] + rho_w = Constant(1000.0) # density of liquid water + # assume n(D) = n_0 * D^mu * exp(-Lambda*D) + # n_0 = N_r * Lambda^(1+mu) / gamma(1 + mu) + N_r = Constant(10**5) # number of rain droplets per m^3 + mu = 0.0 # shape constant of droplet gamma distribution + # assume V(D) = a * D^b * exp(-f*D) * (rho_0 / rho)^g + # take f = 0 + a = Constant(362.) # intercept for velocity distr. in log space + b = 0.65 # inverse scale parameter for velocity distr. + rho0 = Constant(1.22) # reference density in kg/m^3 + g = Constant(0.5) # scaling of density correction + # we keep mu in the expressions even though mu = 0 + threshold = Constant(10**-10) # only do rainfall for r > threshold + Lambda = (N_r * pi * rho_w * math.gamma(4 + mu) + / (6 * math.gamma(1 + mu) * rho * rain)) ** (1. / 3) + Lambda0 = (N_r * pi * rho_w * math.gamma(4 + mu) + / (6 * math.gamma(1 + mu) * rho * threshold)) ** (1. / 3) + v_expression = conditional(rain > threshold, + (a * math.gamma(4 + b + mu) + / (math.gamma(4 + mu) * Lambda ** b) + * (rho0 / rho) ** g), + (a * math.gamma(4 + b + mu) + / (math.gamma(4 + mu) * Lambda0 ** b) + * (rho0 / rho) ** g)) + else: + raise NotImplementedError( + 'Currently there are only implementations for zero and one ' + + 'moment schemes for rainfall. Valid options are ' + + 'AdvectedMoments.M0 and AdvectedMoments.M3') + + if moments != AdvectedMoments.M0: + # TODO: introduce reduced projector + test = TestFunction(Vu) + dx_reduced = dx(degree=4) + proj_eqn = inner(test, v + v_expression*domain.k)*dx_reduced + proj_prob = NonlinearVariationalProblem(proj_eqn, v) + self.determine_v = NonlinearVariationalSolver(proj_prob) + + def evaluate(self, x_in, dt): + """ + Evaluates the source/sink corresponding to the fallout process. + + Args: + x_in (:class:`Function`): the (mixed) field to be evolved. + dt (:class:`Constant`): the time interval for the scheme. + """ + logger.info(f'Evaluating physics parametrisation {self.label.label}') + self.X.assign(x_in) + if self.moments != AdvectedMoments.M0: + self.determine_v.solve() + + +class Coalescence(PhysicsParametrisation): + """ + Represents the coalescence of cloud droplets to form rain droplets. + + Coalescence is the process of forming rain droplets from cloud droplets. + This scheme performs that process, using two parts: accretion, which is + independent of the rain concentration, and auto-accumulation, which is + accelerated by the existence of rain. These parametrisations come from Klemp + and Wilhelmson (1978). The rate of change is limited to prevent production + of negative moisture values. + + This is only implemented for mixing ratio variables. + """ + + def __init__(self, equation, cloud_name='cloud_water', rain_name='rain', + accretion=True, accumulation=True): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + cloud_name (str, optional): name of the cloud variable. Defaults to + 'cloud_water'. + rain_name (str, optional): name of the rain variable. Defaults to + 'rain'. + accretion (bool, optional): whether to include the accretion process + in the parametrisation. Defaults to True. + accumulation (bool, optional): whether to include the accumulation + process in the parametrisation. Defaults to True. + """ + + self.explicit_only = True + label_name = "coalescence" + if accretion: + label_name += "_accretion" + if accumulation: + label_name += "_accumulation" + super().__init__(equation, label_name, parameters=None) + + # Check that fields exist + if cloud_name not in equation.field_names: + raise ValueError(f"Field {cloud_name} does not exist in the equation set") + if rain_name not in equation.field_names: + raise ValueError(f"Field {rain_name} does not exist in the equation set") + + self.cloud_idx = equation.field_names.index(cloud_name) + self.rain_idx = equation.field_names.index(rain_name) + Vcl = equation.function_space.sub(self.cloud_idx) + Vr = equation.function_space.sub(self.rain_idx) + self.cloud_water = Function(Vcl) + self.rain = Function(Vr) + + # declare function space and source field + Vt = self.cloud_water.function_space() + self.source = Function(Vt) + + # define some parameters as attributes + self.dt = Constant(0.0) + # TODO: should these parameters be hard-coded or configurable? + k_1 = Constant(0.001) # accretion rate in 1/s + k_2 = Constant(2.2) # accumulation rate in 1/s + a = Constant(0.001) # min cloud conc in kg/kg + b = Constant(0.875) # power for rain in accumulation + + # make default rates to be zero + accr_rate = Constant(0.0) + accu_rate = Constant(0.0) + + if accretion: + accr_rate = k_1 * (self.cloud_water - a) + if accumulation: + accu_rate = k_2 * self.cloud_water * self.rain ** b + + # Expression for rain increment, with conditionals to prevent negative values + rain_expr = conditional(self.rain < 0.0, # if rain is negative do only accretion + conditional(accr_rate < 0.0, + 0.0, + min_value(accr_rate, self.cloud_water / self.dt)), + # don't turn rain back into cloud + conditional(accr_rate + accu_rate < 0.0, + 0.0, + # if accretion rate is negative do only accumulation + conditional(accr_rate < 0.0, + 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) + + # Add term to equation's residual + test_cl = equation.tests[self.cloud_idx] + test_r = equation.tests[self.rain_idx] + equation.residual += self.label(subject(test_cl * self.source * dx + - test_r * self.source * dx, + equation.X), + self.evaluate) + + def evaluate(self, x_in, dt): + """ + Evaluates the source/sink for the coalescence process. + + Args: + x_in (:class:`Function`): the (mixed) field to be evolved. + dt (:class:`Constant`): the time interval for the scheme. + """ + logger.info(f'Evaluating physics parametrisation {self.label.label}') + # Update the values of internal variables + self.dt.assign(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()) + + +class EvaporationOfRain(PhysicsParametrisation): + """ + Represents the evaporation of rain into water vapour. + + This describes the evaporation of rain into water vapour, with the + associated latent heat change. This parametrisation comes from Klemp and + Wilhelmson (1978). The rate of change is limited to prevent production of + negative moisture values. + + This is only implemented for mixing ratio variables, and when the prognostic + is the virtual dry potential temperature. + """ + + def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', + latent_heat=True): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + cloud_name (str, optional): name of the rain variable. Defaults to + 'rain'. + vapour_name (str, optional): name of the water vapour variable. + Defaults to 'water_vapour'. + latent_heat (bool, optional): whether to have latent heat exchange + feeding back from the phase change. Defaults to True. + + Raises: + NotImplementedError: currently this is only implemented for the + CompressibleEulerEquations. + """ + + self.explicit_only = True + label_name = 'evaporation_of_rain' + super().__init__(equation, label_name, parameters=None) + + # TODO: make a check on the variable type of the active tracers + # if not a mixing ratio, we need to convert to mixing ratios + # this will be easier if we change equations to have dictionary of + # active tracer metadata corresponding to variable names + + # Check that fields exist + if vapour_name not in equation.field_names: + raise ValueError(f"Field {vapour_name} does not exist in the equation set") + if rain_name not in equation.field_names: + raise ValueError(f"Field {rain_name} does not exist in the equation set") + + # Make prognostic for physics scheme + self.X = Function(equation.X.function_space()) + parameters = self.parameters + self.latent_heat = latent_heat + + # Vapour and cloud variables are needed for every form of this scheme + rain_idx = equation.field_names.index(rain_name) + vap_idx = equation.field_names.index(vapour_name) + rain = self.X.subfunctions[rain_idx] + water_vapour = self.X.subfunctions[vap_idx] + + # Indices of variables in mixed function space + V_idxs = [rain_idx, vap_idx] + V = equation.function_space.sub(rain_idx) # space in which to do the calculation + + # Get variables used to calculate saturation curve + if isinstance(equation, CompressibleEulerEquations): + rho_idx = equation.field_names.index('rho') + theta_idx = equation.field_names.index('theta') + rho = self.X.subfunctions[rho_idx] + theta = self.X.subfunctions[theta_idx] + if latent_heat: + V_idxs.append(theta_idx) + + # need to evaluate rho at theta-points, and do this via recovery + boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None + rho_averaged = Function(V) + self.rho_recoverer = Recoverer(rho, rho_averaged, boundary_method=boundary_method) + + exner = thermodynamics.exner_pressure(parameters, rho_averaged, theta) + T = thermodynamics.T(parameters, theta, exner, r_v=water_vapour) + p = thermodynamics.p(parameters, exner) + + # -------------------------------------------------------------------- # + # Compute heat capacities and calculate saturation curve + # -------------------------------------------------------------------- # + # Loop through variables to extract all liquid components + liquid_water = rain + for active_tracer in equation.active_tracers: + if (active_tracer.phase == Phases.liquid + and active_tracer.chemical == 'H2O' and active_tracer.name != rain_name): + liq_idx = equation.field_names.index(active_tracer.name) + liquid_water += self.X.subfunctions[liq_idx] + + # define some parameters as attributes + self.dt = Constant(0.0) + R_d = parameters.R_d + cp = parameters.cp + cv = parameters.cv + c_pv = parameters.c_pv + c_pl = parameters.c_pl + c_vv = parameters.c_vv + R_v = parameters.R_v + + # make useful fields + L_v = thermodynamics.Lv(parameters, T) + R_m = R_d + R_v * water_vapour + c_pml = cp + c_pv * water_vapour + c_pl * liquid_water + c_vml = cv + c_vv * water_vapour + c_pl * liquid_water + + # use Teten's formula to calculate the saturation curve + sat_expr = thermodynamics.r_sat(parameters, T, p) + + # -------------------------------------------------------------------- # + # Evaporation expression + # -------------------------------------------------------------------- # + # TODO: should these parameters be hard-coded or configurable? + # expression for ventilation factor + a = Constant(1.6) + b = Constant(124.9) + c = Constant(0.2046) + C = a + b * (rho_averaged * rain) ** c + + # make appropriate condensation rate + f = Constant(5.4e5) + g = Constant(2.55e6) + h = Constant(0.525) + evap_rate = (((1 - water_vapour / sat_expr) * C * (rho_averaged * rain) ** h) + / (rho_averaged * (f + g / (p * sat_expr)))) + + # adjust evap rate so negative rain doesn't occur + evap_rate = conditional(evap_rate < 0, 0.0, + conditional(rain < 0.0, 0.0, + min_value(evap_rate, rain / self.dt))) + + # -------------------------------------------------------------------- # + # Factors for multiplying source for different variables + # -------------------------------------------------------------------- # + # Factors need to have same shape as V_idxs + factors = [Constant(-1.0), Constant(1.0)] + if latent_heat and isinstance(equation, CompressibleEulerEquations): + factors.append(-theta * (cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml / (R_m * cp * c_vml))) + + # -------------------------------------------------------------------- # + # 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)] + + tests = [equation.tests[idx] for idx in V_idxs] + + # Add source terms to residual + for test, source in zip(tests, self.source): + equation.residual += self.label(subject(test * source * dx, + equation.X), self.evaluate) + + def evaluate(self, x_in, dt): + """ + Applies the process to evaporate rain droplets. + + Args: + x_in (:class:`Function`): the (mixed) field to be evolved. + dt (:class:`Constant`): the time interval for the scheme. + """ + logger.info(f'Evaluating physics parametrisation {self.label.label}') + # Update the values of internal variables + self.dt.assign(dt) + self.X.assign(x_in) + if isinstance(self.equation, CompressibleEulerEquations): + self.rho_recoverer.project() + # Evaluate the source + for interpolator in self.source_interpolators: + interpolator.interpolate() diff --git a/gusto/physics/physics_parametrisation.py b/gusto/physics/physics_parametrisation.py new file mode 100644 index 000000000..2d43dae48 --- /dev/null +++ b/gusto/physics/physics_parametrisation.py @@ -0,0 +1,147 @@ +""" +Defines objects to perform parametrisations of physical processes, or "physics". + +"PhysicsParametrisation" schemes are routines to compute updates to prognostic +fields that represent the action of non-fluid processes, or those fluid +processes that are unresolved. This module contains a set of these processes in +the form of classes with "evaluate" methods. +""" + +from abc import ABCMeta, abstractmethod +from firedrake import Interpolator, Function, dx, Projector +from firedrake.fml import subject +from gusto.core.labels import PhysicsLabel +from gusto.core.logging import logger + + +__all__ = ["PhysicsParametrisation", "SourceSink"] + + +class PhysicsParametrisation(object, metaclass=ABCMeta): + """Base class for the parametrisation of physical processes for Gusto.""" + + def __init__(self, equation, label_name, parameters=None): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + label_name (str): name of physics scheme, to be passed to its label. + parameters (:class:`Configuration`, optional): parameters containing + the values of gas constants. Defaults to None, in which case the + parameters are obtained from the equation. + """ + + self.label = PhysicsLabel(label_name) + self.equation = equation + if parameters is None and hasattr(equation, 'parameters'): + self.parameters = equation.parameters + else: + self.parameters = parameters + + @abstractmethod + def evaluate(self): + """ + Computes the value of physics source and sink terms. + """ + pass + + +class SourceSink(PhysicsParametrisation): + """ + The source or sink of some variable, described through a UFL expression. + + A scheme representing the general source or sink of a variable, described + through a UFL expression. The expression should be for the rate of change + of the variable. It is intended that the source/sink is independent of the + prognostic variables. + + The expression can also be a time-varying expression. In which case a + function should be provided, taking a :class:`Constant` as an argument (to + represent the time.) + """ + + def __init__(self, equation, variable_name, rate_expression, + time_varying=False, method='interpolate'): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + variable_name (str): the name of the variable + rate_expression (:class:`ufl.Expr` or func): an expression giving + the rate of change of the variable. If a time-varying expression + is needed, this should be a function taking a single argument + representing the time. Then the `time_varying` argument must + be set to True. + time_varying (bool, optional): whether the source/sink expression + varies with time. Defaults to False. + method (str, optional): the method to use to evaluate the expression + for the source. Valid options are 'interpolate' or 'project'. + Default is 'interpolate'. + """ + + label_name = f'source_sink_{variable_name}' + super().__init__(equation, label_name, parameters=None) + + if method not in ['interpolate', 'project']: + raise ValueError(f'Method {method} for source/sink evaluation not valid') + self.method = method + self.time_varying = time_varying + self.variable_name = variable_name + + # Check the variable exists + if hasattr(equation, "field_names"): + assert variable_name in equation.field_names, \ + f'Field {variable_name} does not exist in the equation set' + else: + assert variable_name == equation.field_name, \ + f'Field {variable_name} does not exist in the equation' + + # Work out the appropriate function space + if hasattr(equation, "field_names"): + V_idx = equation.field_names.index(variable_name) + W = equation.function_space + V = W.sub(V_idx) + test = equation.tests[V_idx] + else: + V = equation.function_space + test = equation.test + + # Make source/sink term + self.source = Function(V) + equation.residual += self.label(subject(test * self.source * dx, equation.X), + self.evaluate) + + # Handle whether the expression is time-varying or not + if self.time_varying: + expression = rate_expression(equation.domain.t) + else: + expression = rate_expression + + # Handle method of evaluating source/sink + if self.method == 'interpolate': + self.source_interpolator = Interpolator(expression, V) + else: + self.source_projector = Projector(expression, V) + + # If not time-varying, evaluate for the first time here + if not self.time_varying: + if self.method == 'interpolate': + self.source.assign(self.source_interpolator.interpolate()) + else: + self.source.assign(self.source_projector.project()) + + 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. Unused. + dt: (:class:`Constant`): the timestep, which can be the time + interval for the scheme. Unused. + """ + if self.time_varying: + logger.info(f'Evaluating physics parametrisation {self.label.label}') + if self.method == 'interpolate': + self.source.assign(self.source_interpolator.interpolate()) + else: + self.source.assign(self.source_projector.project()) + else: + pass diff --git a/gusto/physics/shallow_water_microphysics.py b/gusto/physics/shallow_water_microphysics.py new file mode 100644 index 000000000..8ffdcf7da --- /dev/null +++ b/gusto/physics/shallow_water_microphysics.py @@ -0,0 +1,356 @@ +""" +Defines microphysics routines to be used with the moist shallow water equations. +""" + +from firedrake import ( + Interpolator, conditional, Function, dx, min_value, max_value, Constant +) +from firedrake.fml import subject +from gusto.core.logging import logger +from gusto.physics.physics_parametrisation import PhysicsParametrisation +from types import FunctionType + +__all__ = ["InstantRain", "SWSaturationAdjustment"] + + +class InstantRain(PhysicsParametrisation): + """ + The process of converting vapour above the saturation curve to rain. + + A scheme to move vapour directly to rain. If convective feedback is true + then this process feeds back directly on the height equation. If rain is + accumulating then excess vapour is being tracked and stored as rain; + otherwise converted vapour is not recorded. The process can happen over the + timestep dt or over a specified time interval tau. + """ + + def __init__(self, equation, saturation_curve, + time_varying_saturation=False, + vapour_name="water_vapour", rain_name=None, gamma_r=1, + convective_feedback=False, beta1=None, tau=None, + parameters=None): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + saturation_curve (:class:`ufl.Expr` or func): the curve above which + excess moisture is converted to rain. Is either prescribed or + dependent on a prognostic field. + time_varying_saturation (bool, optional): set this to True if the + saturation curve is changing in time. Defaults to False. + vapour_name (str, optional): name of the water vapour variable. + Defaults to "water_vapour". + rain_name (str, optional): name of the rain variable. Defaults to + None. + gamma_r (float, optional): Fraction of vapour above the threshold + which is converted to rain. Defaults to one, in which case all + vapour above the threshold is converted. + convective_feedback (bool, optional): True if the conversion of + vapour affects the height equation. Defaults to False. + beta1 (float, optional): Condensation proportionality constant, + used if convection causes a response in the height equation. + Defaults to None, but must be specified if convective_feedback + is True. + tau (float, optional): Timescale for condensation. Defaults to None, + in which case the timestep dt is used. + parameters (:class:`Configuration`, optional): parameters containing + the values of gas constants. Defaults to None, in which case the + parameters are obtained from the equation. + """ + + self.explicit_only = True + label_name = 'instant_rain' + super().__init__(equation, label_name, parameters=parameters) + + self.convective_feedback = convective_feedback + self.time_varying_saturation = time_varying_saturation + + # check for the correct fields + assert vapour_name in equation.field_names, f"Field {vapour_name} does not exist in the equation set" + self.Vv_idx = equation.field_names.index(vapour_name) + + if rain_name is not None: + assert rain_name in equation.field_names, f"Field {rain_name} does not exist in the equation set " + + if self.convective_feedback: + assert "D" in equation.field_names, "Depth field must exist for convective feedback" + assert beta1 is not None, "If convective feedback is used, beta1 parameter must be specified" + + # obtain function space and functions; vapour needed for all cases + W = equation.function_space + Vv = W.sub(self.Vv_idx) + test_v = equation.tests[self.Vv_idx] + + # depth needed if convective feedback + if self.convective_feedback: + self.VD_idx = equation.field_names.index("D") + VD = W.sub(self.VD_idx) + test_D = equation.tests[self.VD_idx] + self.D = Function(VD) + + # the source function is the difference between the water vapour and + # the saturation function + self.water_v = Function(Vv) + self.source = Function(Vv) + + # tau is the timescale for conversion (may or may not be the timestep) + if tau is not None: + self.set_tau_to_dt = False + self.tau = tau + else: + self.set_tau_to_dt = True + self.tau = Constant(0) + logger.info("Timescale for rain conversion has been set to dt. If this is not the intention then provide a tau parameter as an argument to InstantRain.") + + if self.time_varying_saturation: + if isinstance(saturation_curve, FunctionType): + self.saturation_computation = saturation_curve + self.saturation_curve = Function(Vv) + else: + raise NotImplementedError( + "If time_varying_saturation is True then saturation must be a Python function of a prognostic field.") + else: + assert not isinstance(saturation_curve, FunctionType), "If time_varying_saturation is not True then saturation cannot be a Python function." + self.saturation_curve = saturation_curve + + # lose proportion of vapour above the saturation curve + equation.residual += self.label(subject(test_v * self.source * dx, + equation.X), + self.evaluate) + + # if rain is not none then the excess vapour is being tracked and is + # added to rain + if rain_name is not None: + Vr_idx = equation.field_names.index(rain_name) + test_r = equation.tests[Vr_idx] + equation.residual -= self.label(subject(test_r * self.source * dx, + equation.X), + self.evaluate) + + # if feeding back on the height adjust the height equation + if convective_feedback: + equation.residual += self.label(subject(test_D * beta1 * self.source * dx, + equation.X), + self.evaluate) + + # interpolator does the conversion of vapour to rain + self.source_interpolator = Interpolator(conditional( + self.water_v > self.saturation_curve, + (1/self.tau)*gamma_r*(self.water_v - self.saturation_curve), + 0), Vv) + + def evaluate(self, x_in, dt): + """ + Evalutes the source term generated by the physics. + + Computes the physics contributions (loss of vapour, accumulation of + rain and loss of height due to convection) at each timestep. + + 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. + """ + logger.info(f'Evaluating physics parametrisation {self.label.label}') + if self.convective_feedback: + self.D.assign(x_in.subfunctions[self.VD_idx]) + if self.time_varying_saturation: + self.saturation_curve.interpolate(self.saturation_computation(x_in)) + if self.set_tau_to_dt: + self.tau.assign(dt) + self.water_v.assign(x_in.subfunctions[self.Vv_idx]) + self.source.assign(self.source_interpolator.interpolate()) + + +class SWSaturationAdjustment(PhysicsParametrisation): + """ + Represents the process of adjusting water vapour and cloud water according + to a saturation function, via condensation and evaporation processes. + + This physics scheme follows that of Zerroukat and Allen (2015). + + """ + + def __init__(self, equation, saturation_curve, + time_varying_saturation=False, vapour_name='water_vapour', + cloud_name='cloud_water', convective_feedback=False, + beta1=None, thermal_feedback=False, beta2=None, gamma_v=1, + time_varying_gamma_v=False, tau=None, + parameters=None): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation + saturation_curve (:class:`ufl.Expr` or func): the curve which + dictates when phase changes occur. In a saturated atmosphere + vapour above the saturation curve becomes cloud, and if the + atmosphere is sub-saturated and there is cloud present cloud + will become vapour until the saturation curve is reached. The + saturation curve is either prescribed or dependent on a + prognostic field. + time_varying_saturation (bool, optional): set this to True if the + saturation curve is changing in time. Defaults to False. + vapour_name (str, optional): name of the water vapour variable. + Defaults to 'water_vapour'. + cloud_name (str, optional): name of the cloud variable. Defaults to + 'cloud_water'. + convective_feedback (bool, optional): True if the conversion of + vapour affects the height equation. Defaults to False. + beta1 (float, optional): Condensation proportionality constant for + height feedback, used if convection causes a response in the + height equation. Defaults to None, but must be specified if + convective_feedback is True. + thermal_feedback (bool, optional): True if moist conversions + affect the buoyancy equation. Defaults to False. + beta2 (float, optional): Condensation proportionality constant + for thermal feedback. Defaults to None, but must be specified + if thermal_feedback is True. This is equivalent to the L_v + parameter in Zerroukat and Allen (2015). + gamma_v (ufl expression or :class: `function`): The proportion of + moist species that is converted when a conversion between + vapour and cloud is taking place. Defaults to one, in which + case the full amount of species to bring vapour to the + saturation curve will undergo a conversion. Converting only a + fraction avoids a two-timestep oscillation between vapour and + cloud when saturation is tempertature/height-dependent. + time_varying_gamma_v (bool, optional): set this to True + if the fraction of moist species converted changes in time + (if gamma_v is temperature/height-dependent). + tau (float, optional): Timescale for condensation and evaporation. + Defaults to None, in which case the timestep dt is used. + parameters (:class:`Configuration`, optional): parameters containing + the values of constants. Defaults to None, in which case the + parameters are obtained from the equation. + """ + + self.explicit_only = True + label_name = 'saturation_adjustment' + super().__init__(equation, label_name, parameters=parameters) + + self.time_varying_saturation = time_varying_saturation + self.convective_feedback = convective_feedback + self.thermal_feedback = thermal_feedback + self.time_varying_gamma_v = time_varying_gamma_v + + # Check for the correct fields + assert vapour_name in equation.field_names, f"Field {vapour_name} does not exist in the equation set" + assert cloud_name in equation.field_names, f"Field {cloud_name} does not exist in the equation set" + self.Vv_idx = equation.field_names.index(vapour_name) + self.Vc_idx = equation.field_names.index(cloud_name) + + if self.convective_feedback: + assert "D" in equation.field_names, "Depth field must exist for convective feedback" + assert beta1 is not None, "If convective feedback is used, beta1 parameter must be specified" + + if self.thermal_feedback: + assert "b" in equation.field_names, "Buoyancy field must exist for thermal feedback" + assert beta2 is not None, "If thermal feedback is used, beta2 parameter must be specified" + + # Obtain function spaces and functions + W = equation.function_space + Vv = W.sub(self.Vv_idx) + Vc = W.sub(self.Vc_idx) + V_idxs = [self.Vv_idx, self.Vc_idx] + + # Source functions for both vapour and cloud + self.water_v = Function(Vv) + self.cloud = Function(Vc) + + # depth needed if convective feedback + if self.convective_feedback: + self.VD_idx = equation.field_names.index("D") + VD = W.sub(self.VD_idx) + self.D = Function(VD) + V_idxs.append(self.VD_idx) + + # buoyancy needed if thermal feedback + if self.thermal_feedback: + self.Vb_idx = equation.field_names.index("b") + Vb = W.sub(self.Vb_idx) + self.b = Function(Vb) + V_idxs.append(self.Vb_idx) + + # tau is the timescale for condensation/evaporation (may or may not be the timestep) + if tau is not None: + self.set_tau_to_dt = False + self.tau = tau + else: + self.set_tau_to_dt = True + self.tau = Constant(0) + logger.info("Timescale for moisture conversion between vapour and cloud has been set to dt. If this is not the intention then provide a tau parameter as an argument to SWSaturationAdjustment.") + + if self.time_varying_saturation: + if isinstance(saturation_curve, FunctionType): + self.saturation_computation = saturation_curve + self.saturation_curve = Function(Vv) + else: + raise NotImplementedError( + "If time_varying_saturation is True then saturation must be a Python function of at least one prognostic field.") + else: + assert not isinstance(saturation_curve, FunctionType), "If time_varying_saturation is not True then saturation cannot be a Python function." + self.saturation_curve = saturation_curve + + # Saturation adjustment expression, adjusted to stop negative values + sat_adj_expr = (self.water_v - self.saturation_curve) / self.tau + sat_adj_expr = conditional(sat_adj_expr < 0, + max_value(sat_adj_expr, + -self.cloud / self.tau), + min_value(sat_adj_expr, + self.water_v / self.tau)) + + # If gamma_v depends on variables + if self.time_varying_gamma_v: + if isinstance(gamma_v, FunctionType): + self.gamma_v_computation = gamma_v + self.gamma_v = Function(Vv) + else: + raise NotImplementedError( + "If time_varying_thermal_feedback is True then gamma_v must be a Python function of at least one prognostic field.") + else: + assert not isinstance(gamma_v, FunctionType), "If time_varying_thermal_feedback is not True then gamma_v cannot be a Python function." + self.gamma_v = gamma_v + + # Factors for multiplying source for different variables + factors = [self.gamma_v, -self.gamma_v] + if convective_feedback: + factors.append(self.gamma_v*beta1) + if thermal_feedback: + factors.append(parameters.g*self.gamma_v*beta2) + + # Add terms to equations and make interpolators + self.source = [Function(Vc) for factor in factors] + self.source_interpolators = [Interpolator(sat_adj_expr*factor, source) + for factor, source in zip(factors, self.source)] + + tests = [equation.tests[idx] for idx in V_idxs] + + # Add source terms to residual + for test, source in zip(tests, self.source): + equation.residual += self.label(subject(test * source * dx, + equation.X), self.evaluate) + + def evaluate(self, x_in, dt): + """ + Evaluates the source term generated by the physics. + + Computes the phyiscs contributions to water vapour and cloud water at + each timestep. + + 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. + """ + logger.info(f'Evaluating physics parametrisation {self.label.label}') + if self.convective_feedback: + self.D.assign(x_in.subfunctions[self.VD_idx]) + if self.thermal_feedback: + self.b.assign(x_in.subfunctions[self.Vb_idx]) + if self.time_varying_saturation: + self.saturation_curve.interpolate(self.saturation_computation(x_in)) + if self.set_tau_to_dt: + self.tau.assign(dt) + self.water_v.assign(x_in.subfunctions[self.Vv_idx]) + self.cloud.assign(x_in.subfunctions[self.Vc_idx]) + if self.time_varying_gamma_v: + self.gamma_v.interpolate(self.gamma_v_computation(x_in)) + for interpolator in self.source_interpolators: + interpolator.interpolate() diff --git a/gusto/solvers/__init__.py b/gusto/solvers/__init__.py new file mode 100644 index 000000000..efb4a5af4 --- /dev/null +++ b/gusto/solvers/__init__.py @@ -0,0 +1,2 @@ +from gusto.solvers.linear_solvers import * # noqa +from gusto.solvers.preconditioners import * # noqa \ No newline at end of file diff --git a/gusto/linear_solvers.py b/gusto/solvers/linear_solvers.py similarity index 99% rename from gusto/linear_solvers.py rename to gusto/solvers/linear_solvers.py index be40ed36f..c21aad09f 100644 --- a/gusto/linear_solvers.py +++ b/gusto/solvers/linear_solvers.py @@ -16,10 +16,10 @@ from firedrake.petsc import flatten_parameters from pyop2.profiling import timed_function, timed_region -from gusto.active_tracers import TracerVariableType -from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual -from gusto.labels import linearisation, time_derivative, hydrostatic -from gusto import thermodynamics +from gusto.equations.active_tracers import TracerVariableType +from gusto.core.logging import logger, DEBUG, logging_ksp_monitor_true_residual +from gusto.core.labels import linearisation, time_derivative, hydrostatic +from gusto.equations import thermodynamics from gusto.recovery.recovery_kernels import AverageWeightings, AverageKernel from abc import ABCMeta, abstractmethod, abstractproperty @@ -353,7 +353,7 @@ def L_tr(f): # Log residuals on hybridized solver self.log_ksp_residuals(self.hybridized_solver.snes.ksp) # Log residuals on the trace system too - from gusto.logging import attach_custom_monitor + from gusto.core.logging import attach_custom_monitor python_context = self.hybridized_solver.snes.ksp.pc.getPythonContext() attach_custom_monitor(python_context, logging_ksp_monitor_true_residual) diff --git a/gusto/preconditioners.py b/gusto/solvers/preconditioners.py similarity index 100% rename from gusto/preconditioners.py rename to gusto/solvers/preconditioners.py diff --git a/gusto/spatial_methods/__init__.py b/gusto/spatial_methods/__init__.py new file mode 100644 index 000000000..baa1e0788 --- /dev/null +++ b/gusto/spatial_methods/__init__.py @@ -0,0 +1,4 @@ +from gusto.spatial_methods.spatial_methods import * # noqa +from gusto.spatial_methods.diffusion_methods import * # noqa +from gusto.spatial_methods.transport_methods import * # noqa +from gusto.spatial_methods.limiters import * # noqa diff --git a/gusto/diffusion_methods.py b/gusto/spatial_methods/diffusion_methods.py similarity index 98% rename from gusto/diffusion_methods.py rename to gusto/spatial_methods/diffusion_methods.py index d35497964..c36fe1be3 100644 --- a/gusto/diffusion_methods.py +++ b/gusto/spatial_methods/diffusion_methods.py @@ -1,8 +1,8 @@ """Provides discretisations for diffusion terms.""" from firedrake import inner, outer, grad, avg, dx, dS_h, dS_v, dS, FacetNormal -from gusto.labels import diffusion -from gusto.spatial_methods import SpatialMethod +from gusto.core.labels import diffusion +from gusto.spatial_methods.spatial_methods import SpatialMethod __all__ = ["InteriorPenaltyDiffusion", "CGDiffusion"] diff --git a/gusto/limiters.py b/gusto/spatial_methods/limiters.py similarity index 99% rename from gusto/limiters.py rename to gusto/spatial_methods/limiters.py index ac11f1f6b..8a782a054 100644 --- a/gusto/limiters.py +++ b/gusto/spatial_methods/limiters.py @@ -8,7 +8,7 @@ from firedrake import (BrokenElement, Function, FunctionSpace, interval, FiniteElement, TensorProductElement) from firedrake.slope_limiter.vertex_based_limiter import VertexBasedLimiter -from gusto.kernels import LimitMidpoints, ClipZero +from gusto.core.kernels import LimitMidpoints, ClipZero import numpy as np diff --git a/gusto/spatial_methods.py b/gusto/spatial_methods/spatial_methods.py similarity index 98% rename from gusto/spatial_methods.py rename to gusto/spatial_methods/spatial_methods.py index 2b1751545..6ad214dcd 100644 --- a/gusto/spatial_methods.py +++ b/gusto/spatial_methods/spatial_methods.py @@ -5,7 +5,7 @@ from firedrake import split from firedrake.fml import Term, keep, drop -from gusto.labels import prognostic +from gusto.core.labels import prognostic __all__ = ['SpatialMethod'] diff --git a/gusto/transport_methods.py b/gusto/spatial_methods/transport_methods.py similarity index 99% rename from gusto/transport_methods.py rename to gusto/spatial_methods/transport_methods.py index 5420bb8ca..c2b625882 100644 --- a/gusto/transport_methods.py +++ b/gusto/spatial_methods/transport_methods.py @@ -7,10 +7,10 @@ grad, div, FacetNormal, Function, sign, avg, cross, curl, split ) from firedrake.fml import Term, keep, drop -from gusto.configuration import IntegrateByParts, TransportEquationType -from gusto.labels import prognostic, transport, transporting_velocity, ibp_label -from gusto.logging import logger -from gusto.spatial_methods import SpatialMethod +from gusto.core.configuration import IntegrateByParts, TransportEquationType +from gusto.core.labels import prognostic, transport, transporting_velocity, ibp_label +from gusto.core.logging import logger +from gusto.spatial_methods.spatial_methods import SpatialMethod __all__ = ["DefaultTransport", "DGUpwind"] diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py deleted file mode 100644 index 03a42fe81..000000000 --- a/gusto/time_discretisation.py +++ /dev/null @@ -1,2294 +0,0 @@ -u""" -Objects for discretising time derivatives. - -Time discretisation objects discretise ∂y/∂t = F(y), for variable y, time t and -operator F. -""" - -from abc import ABCMeta, abstractmethod, abstractproperty -import math -import numpy as np - -from firedrake import ( - Function, TestFunction, TestFunctions, NonlinearVariationalProblem, - NonlinearVariationalSolver, DirichletBC, split, Constant -) -from firedrake.fml import ( - replace_subject, replace_test_function, Term, all_terms, drop, keep -) -from firedrake.formmanipulation import split_form -from firedrake.utils import cached_property - -from gusto.configuration import EmbeddedDGOptions, RecoveryOptions -from gusto.labels import (time_derivative, prognostic, physics_label, - implicit, explicit) -from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual -from gusto.wrappers import * - - -__all__ = ["ForwardEuler", "BackwardEuler", "ExplicitMultistage", - "IMEXMultistage", "SSPRK3", "RK4", "Heun", "ThetaMethod", - "TrapeziumRule", "BDF2", "TR_BDF2", "Leapfrog", "AdamsMoulton", - "AdamsBashforth", "ImplicitMidpoint", "QinZhang", - "IMEX_Euler", "ARS3", "ARK2", "Trap2", "SSP3"] - - -def wrapper_apply(original_apply): - """Decorator to add steps for using a wrapper around the apply method.""" - def get_apply(self, x_out, x_in): - - if self.wrapper is not None: - - def new_apply(self, x_out, x_in): - - self.wrapper.pre_apply(x_in) - original_apply(self, self.wrapper.x_out, self.wrapper.x_in) - self.wrapper.post_apply(x_out) - - return new_apply(self, x_out, x_in) - - else: - - return original_apply(self, x_out, x_in) - - return get_apply - - -class TimeDiscretisation(object, metaclass=ABCMeta): - """Base class for time discretisation schemes.""" - - def __init__(self, domain, field_name=None, solver_parameters=None, - limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - self.domain = domain - self.field_name = field_name - self.equation = None - - self.dt = Constant(0.0) - self.dt.assign(domain.dt) - self.original_dt = Constant(0.0) - self.original_dt.assign(self.dt) - self.options = options - self.limiter = limiter - self.courant_max = None - - if options is not None: - self.wrapper_name = options.name - if self.wrapper_name == "mixed_options": - self.wrapper = MixedFSWrapper() - - for field, suboption in options.suboptions.items(): - if suboption.name == 'embedded_dg': - self.wrapper.subwrappers.update({field: EmbeddedDGWrapper(self, suboption)}) - elif suboption.name == "recovered": - self.wrapper.subwrappers.update({field: RecoveryWrapper(self, suboption)}) - elif suboption.name == "supg": - raise RuntimeError( - 'Time discretisation: suboption SUPG is currently not implemented within MixedOptions') - else: - raise RuntimeError( - f'Time discretisation: suboption wrapper {wrapper_name} not implemented') - elif self.wrapper_name == "embedded_dg": - self.wrapper = EmbeddedDGWrapper(self, options) - elif self.wrapper_name == "recovered": - self.wrapper = RecoveryWrapper(self, options) - elif self.wrapper_name == "supg": - self.wrapper = SUPGWrapper(self, options) - else: - raise RuntimeError( - f'Time discretisation: wrapper {self.wrapper_name} not implemented') - else: - self.wrapper = None - self.wrapper_name = None - - # get default solver options if none passed in - if solver_parameters is None: - self.solver_parameters = {'ksp_type': 'cg', - 'pc_type': 'bjacobi', - 'sub_pc_type': 'ilu'} - else: - self.solver_parameters = solver_parameters - - def setup(self, equation, apply_bcs=True, *active_labels): - """ - Set up the time discretisation based on the equation. - - Args: - equation (:class:`PrognosticEquation`): the model's equation. - apply_bcs (bool, optional): whether to apply the equation's boundary - conditions. Defaults to True. - *active_labels (:class:`Label`): labels indicating which terms of - the equation to include. - """ - self.equation = equation - self.residual = equation.residual - - if self.field_name is not None and hasattr(equation, "field_names"): - self.idx = equation.field_names.index(self.field_name) - self.fs = equation.spaces[self.idx] - self.residual = self.residual.label_map( - lambda t: t.get(prognostic) == self.field_name, - lambda t: Term( - split_form(t.form)[self.idx].form, - t.labels), - drop) - - else: - self.field_name = equation.field_name - self.fs = equation.function_space - self.idx = None - - bcs = equation.bcs[self.field_name] - - if len(active_labels) > 0: - self.residual = self.residual.label_map( - lambda t: any(t.has_label(time_derivative, *active_labels)), - map_if_false=drop) - - self.evaluate_source = [] - self.physics_names = [] - for t in self.residual: - if t.has_label(physics_label): - physics_name = t.get(physics_label) - if t.labels[physics_name] not in self.physics_names: - self.evaluate_source.append(t.labels[physics_name]) - self.physics_names.append(t.labels[physics_name]) - - # -------------------------------------------------------------------- # - # Set up Wrappers - # -------------------------------------------------------------------- # - - if self.wrapper is not None: - if self.wrapper_name == "mixed_options": - - self.wrapper.wrapper_spaces = equation.spaces - self.wrapper.field_names = equation.field_names - - for field, subwrapper in self.wrapper.subwrappers.items(): - - if field not in equation.field_names: - raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set") - - field_idx = equation.field_names.index(field) - subwrapper.setup(equation.spaces[field_idx]) - - # Update the function space to that needed by the wrapper - self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space - - self.wrapper.setup() - self.fs = self.wrapper.function_space - new_test_mixed = TestFunctions(self.fs) - - # Replace the original test function with one from the new - # function space defined by the subwrappers - self.residual = self.residual.label_map( - all_terms, - map_if_true=replace_test_function(new_test_mixed)) - - else: - if self.wrapper_name == "supg": - self.wrapper.setup() - else: - self.wrapper.setup(self.fs) - self.fs = self.wrapper.function_space - if self.solver_parameters is None: - self.solver_parameters = self.wrapper.solver_parameters - new_test = TestFunction(self.wrapper.test_space) - # SUPG has a special wrapper - if self.wrapper_name == "supg": - new_test = self.wrapper.test - - # Replace the original test function with the one from the wrapper - self.residual = self.residual.label_map( - all_terms, - map_if_true=replace_test_function(new_test)) - - self.residual = self.wrapper.label_terms(self.residual) - - # -------------------------------------------------------------------- # - # Make boundary conditions - # -------------------------------------------------------------------- # - - if not apply_bcs: - self.bcs = None - elif self.wrapper is not None: - # Transfer boundary conditions onto test function space - self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain) - for bc in bcs] - else: - self.bcs = bcs - - # -------------------------------------------------------------------- # - # Make the required functions - # -------------------------------------------------------------------- # - - self.x_out = Function(self.fs) - self.x1 = Function(self.fs) - - @property - def nlevels(self): - return 1 - - @abstractproperty - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x_out, old_idx=self.idx), - map_if_false=drop) - - return l.form - - @abstractproperty - def rhs(self): - """Set up the time discretisation's right hand side.""" - r = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) - - r = r.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*t) - - return r.form - - @cached_property - def solver(self): - """Set up the problem and the solver.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__ - solver = NonlinearVariationalSolver( - problem, - solver_parameters=self.solver_parameters, - options_prefix=solver_name - ) - if logger.isEnabledFor(DEBUG): - solver.snes.ksp.setMonitor(logging_ksp_monitor_true_residual) - return solver - - @abstractmethod - def apply(self, x_out, x_in): - """ - Apply the time discretisation to advance one whole time step. - - Args: - x_out (:class:`Function`): the output field to be computed. - x_in (:class:`Function`): the input field. - """ - pass - - -class IMEXMultistage(TimeDiscretisation): - """ - A class for implementing general IMEX multistage (Runge-Kutta) - methods based on two Butcher tableaus, to solve \n - - ∂y/∂t = F(y) + S(y) \n - - Where F are implicit fast terms, and S are explicit slow terms. \n - - There are three steps to move from the current solution, y^n, to the new one, y^{n+1} - - For each i = 1, s in an s stage method - we compute the intermediate solutions: \n - y_i = y^n + dt*(a_i1*F(y_1) + a_i2*F(y_2)+ ... + a_ii*F(y_i)) \n - + dt*(d_i1*S(y_1) + d_i2*S(y_2)+ ... + d_{i,i-1}*S(y_{i-1})) - - At the last stage, compute the new solution by: \n - y^{n+1} = y^n + dt*(b_1*F(y_1) + b_2*F(y_2) + .... + b_s*F(y_s)) \n - + dt*(e_1*S(y_1) + e_2*S(y_2) + .... + e_s*S(y_s)) \n - - """ - # -------------------------------------------------------------------------- - # Butcher tableaus for a s-th order - # diagonally implicit scheme (left) and explicit scheme (right): - # c_0 | a_00 0 . 0 f_0 | 0 0 . 0 - # c_1 | a_10 a_11 . 0 f_1 | d_10 0 . 0 - # . | . . . . . | . . . . - # . | . . . . . | . . . . - # c_s | a_s0 a_s1 . a_ss f_s | d_s0 d_s1 . 0 - # ------------------------- ------------------------- - # | b_1 b_2 ... b_s | b_1 b_2 ... b_s - # - # - # The corresponding square 'butcher_imp' and 'butcher_exp' matrices are: - # - # [a_00 0 0 . 0 ] [ 0 0 0 . 0 ] - # [a_10 a_11 0 . 0 ] [d_10 0 0 . 0 ] - # [a_20 a_21 a_22 . 0 ] [d_20 d_21 0 . 0 ] - # [ . . . . . ] [ . . . . . ] - # [ b_0 b_1 . b_s] [ e_0 e_1 . . e_s] - # - # -------------------------------------------------------------------------- - - def __init__(self, domain, butcher_imp, butcher_exp, field_name=None, - solver_parameters=None, limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - butcher_imp (:class:`numpy.ndarray`): A matrix containing the coefficients of - a butcher tableau defining a given implicit Runge Kutta time discretisation. - butcher_exp (:class:`numpy.ndarray`): A matrix containing the coefficients of - a butcher tableau defining a given explicit Runge Kutta time discretisation. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - super().__init__(domain, field_name=field_name, - solver_parameters=solver_parameters, - options=options) - self.butcher_imp = butcher_imp - self.butcher_exp = butcher_exp - self.nStages = int(np.shape(self.butcher_imp)[1]) - - def setup(self, equation, apply_bcs=True, *active_labels): - """ - Set up the time discretisation based on the equation. - - Args: - equation (:class:`PrognosticEquation`): the model's equation. - *active_labels (:class:`Label`): labels indicating which terms of - the equation to include. - """ - - super().setup(equation, apply_bcs, *active_labels) - - # Check all terms are labeled implicit, exlicit - for t in self.residual: - if ((not t.has_label(implicit)) and (not t.has_label(explicit)) - and (not t.has_label(time_derivative))): - raise NotImplementedError("Non time-derivative terms must be labeled as implicit or explicit") - - self.xs = [Function(self.fs) for i in range(self.nStages)] - - @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - return super(IMEXMultistage, self).lhs - - @cached_property - def rhs(self): - """Set up the discretisation's right hand side (the time derivative).""" - return super(IMEXMultistage, self).rhs - - def res(self, stage): - """Set up the discretisation's residual for a given stage.""" - # Add time derivative terms y_s - y^n for stage s - mass_form = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=drop) - residual = mass_form.label_map(all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - residual -= mass_form.label_map(all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) - # Loop through stages up to s-1 and calcualte/sum - # dt*(a_s1*F(y_1) + a_s2*F(y_2)+ ... + a_{s,s-1}*F(y_{s-1})) - # and - # dt*(d_s1*S(y_1) + d_s2*S(y_2)+ ... + d_{s,s-1}*S(y_{s-1})) - for i in range(stage): - r_exp = self.residual.label_map( - lambda t: t.has_label(explicit), - map_if_true=replace_subject(self.xs[i], old_idx=self.idx), - map_if_false=drop) - r_exp = r_exp.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=lambda t: Constant(self.butcher_exp[stage, i])*self.dt*t) - r_imp = self.residual.label_map( - lambda t: t.has_label(implicit), - map_if_true=replace_subject(self.xs[i], old_idx=self.idx), - map_if_false=drop) - r_imp = r_imp.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=lambda t: Constant(self.butcher_imp[stage, i])*self.dt*t) - residual += r_imp - residual += r_exp - # Calculate and add on dt*a_ss*F(y_s) - r_imp = self.residual.label_map( - lambda t: t.has_label(implicit), - map_if_true=replace_subject(self.x_out, old_idx=self.idx), - map_if_false=drop) - r_imp = r_imp.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=lambda t: Constant(self.butcher_imp[stage, stage])*self.dt*t) - residual += r_imp - return residual.form - - @property - def final_res(self): - """Set up the discretisation's final residual.""" - # Add time derivative terms y^{n+1} - y^n - mass_form = self.residual.label_map(lambda t: t.has_label(time_derivative), - map_if_false=drop) - residual = mass_form.label_map(all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - residual -= mass_form.label_map(all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) - # Loop through stages up to s-1 and calcualte/sum - # dt*(b_1*F(y_1) + b_2*F(y_2) + .... + b_s*F(y_s)) - # and - # dt*(e_1*S(y_1) + e_2*S(y_2) + .... + e_s*S(y_s)) - for i in range(self.nStages): - r_exp = self.residual.label_map( - lambda t: t.has_label(explicit), - map_if_true=replace_subject(self.xs[i], old_idx=self.idx), - map_if_false=drop) - r_exp = r_exp.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=lambda t: Constant(self.butcher_exp[self.nStages, i])*self.dt*t) - r_imp = self.residual.label_map( - lambda t: t.has_label(implicit), - map_if_true=replace_subject(self.xs[i], old_idx=self.idx), - map_if_false=drop) - r_imp = r_imp.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=lambda t: Constant(self.butcher_imp[self.nStages, i])*self.dt*t) - residual += r_imp - residual += r_exp - return residual.form - - @cached_property - def solvers(self): - """Set up a list of solvers for each problem at a stage.""" - solvers = [] - for stage in range(self.nStages): - # setup solver using residual defined in derived class - problem = NonlinearVariationalProblem(self.res(stage), self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage) - solvers.append(NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name)) - return solvers - - @cached_property - def final_solver(self): - """Set up a solver for the final solve to evaluate time level n+1.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.final_res, self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__ - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) - - def apply(self, x_out, x_in): - self.x1.assign(x_in) - solver_list = self.solvers - - for stage in range(self.nStages): - self.solver = solver_list[stage] - self.solver.solve() - self.xs[stage].assign(self.x_out) - - self.final_solver.solve() - x_out.assign(self.x_out) - - -class ExplicitTimeDiscretisation(TimeDiscretisation): - """Base class for explicit time discretisations.""" - - def __init__(self, domain, field_name=None, fixed_subcycles=None, - subcycle_by_courant=None, solver_parameters=None, limiter=None, - options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - fixed_subcycles (int, optional): the fixed number of sub-steps to - perform. This option cannot be specified with the - `subcycle_by_courant` argument. Defaults to None. - subcycle_by_courant (float, optional): specifying this option will - make the scheme perform adaptive sub-cycling based on the - Courant number. The specified argument is the maximum Courant - for one sub-cycle. Defaults to None, in which case adaptive - sub-cycling is not used. This option cannot be specified with the - `fixed_subcycles` argument. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - super().__init__(domain, field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - - if fixed_subcycles is not None and subcycle_by_courant is not None: - raise ValueError('Cannot specify both subcycle and subcycle_by ' - + 'arguments to a time discretisation') - self.fixed_subcycles = fixed_subcycles - self.subcycle_by_courant = subcycle_by_courant - - def setup(self, equation, apply_bcs=True, *active_labels): - """ - Set up the time discretisation based on the equation. - - Args: - equation (:class:`PrognosticEquation`): the model's equation. - apply_bcs (bool, optional): whether boundary conditions are to be - applied. Defaults to True. - *active_labels (:class:`Label`): labels indicating which terms of - the equation to include. - """ - super().setup(equation, apply_bcs, *active_labels) - - # if user has specified a number of fixed subcycles, then save this - # and rescale dt accordingly; else perform just one cycle using dt - if self.fixed_subcycles is not None: - self.dt.assign(self.dt/self.fixed_subcycles) - self.ncycles = self.fixed_subcycles - else: - self.dt = self.dt - self.ncycles = 1 - self.x0 = Function(self.fs) - self.x1 = Function(self.fs) - - @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x_out, self.idx), - map_if_false=drop) - - return l.form - - @cached_property - def solver(self): - """Set up the problem and the solver.""" - # setup linear solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs - self.rhs, self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__ - # If snes_type not specified by user, set this to ksp only to avoid outer Newton iteration - self.solver_parameters.setdefault('snes_type', 'ksponly') - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - @abstractmethod - def apply_cycle(self, x_out, x_in): - """ - Apply the time discretisation through a single sub-step. - - Args: - x_out (:class:`Function`): the output field to be computed. - x_in (:class:`Function`): the input field. - """ - pass - - @wrapper_apply - def apply(self, x_out, x_in): - """ - Apply the time discretisation to advance one whole time step. - - Args: - x_out (:class:`Function`): the output field to be computed. - x_in (:class:`Function`): the input field. - """ - # If doing adaptive subcycles, update dt and ncycles here - if self.subcycle_by_courant is not None: - self.ncycles = math.ceil(float(self.courant_max)/self.subcycle_by_courant) - self.dt.assign(self.original_dt/self.ncycles) - - self.x0.assign(x_in) - for i in range(self.ncycles): - self.apply_cycle(self.x1, self.x0) - self.x0.assign(self.x1) - x_out.assign(self.x1) - - -class ImplicitMultistage(TimeDiscretisation): - """ - A class for implementing general diagonally implicit multistage (Runge-Kutta) - methods based on its Butcher tableau. - - Unlike the explicit method, all upper diagonal a_ij elements are non-zero for implicit methods. - - There are three steps to move from the current solution, y^n, to the new one, y^{n+1} - - For each i = 1, s in an s stage method - we have the intermediate solutions: \n - y_i = y^n + dt*(a_i1*k_1 + a_i2*k_2 + ... + a_ii*k_i) \n - We compute the gradient at the intermediate location, k_i = F(y_i) \n - - At the last stage, compute the new solution by: \n - y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s) \n - - """ - # --------------------------------------------------------------------------- - # Butcher tableau for a s-th order - # diagonally implicit scheme: - # c_0 | a_00 0 . 0 - # c_1 | a_10 a_11 . 0 - # . | . . . . - # . | . . . . - # c_s | a_s0 a_s1 . a_ss - # ------------------------- - # | b_1 b_2 ... b_s - # - # - # The corresponding square 'butcher_matrix' is: - # - # [a_00 0 . 0 ] - # [a_10 a_11 . 0 ] - # [ . . . . ] - # [ b_0 b_1 . b_s] - # --------------------------------------------------------------------------- - - def __init__(self, domain, butcher_matrix, field_name=None, - solver_parameters=None, limiter=None, options=None,): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - butcher_matrix (numpy array): A matrix containing the coefficients of - a butcher tableau defining a given Runge Kutta time discretisation. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - super().__init__(domain, field_name=field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - self.butcher_matrix = butcher_matrix - self.nStages = int(np.shape(self.butcher_matrix)[1]) - - def setup(self, equation, apply_bcs=True, *active_labels): - """ - Set up the time discretisation based on the equation. - - Args: - equation (:class:`PrognosticEquation`): the model's equation. - *active_labels (:class:`Label`): labels indicating which terms of - the equation to include. - """ - - super().setup(equation, apply_bcs, *active_labels) - - self.k = [Function(self.fs) for i in range(self.nStages)] - - def lhs(self): - return super().lhs - - def rhs(self): - return super().rhs - - def solver(self, stage): - residual = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=replace_subject(self.xnph, self.idx), - ) - mass_form = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=drop) - residual += mass_form.label_map(all_terms, - replace_subject(self.x_out, self.idx)) - - problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs) - - solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage) - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - @cached_property - def solvers(self): - solvers = [] - for stage in range(self.nStages): - solvers.append(self.solver(stage)) - return solvers - - def solve_stage(self, x0, stage): - self.x1.assign(x0) - for i in range(stage): - self.x1.assign(self.x1 + self.butcher_matrix[stage, i]*self.dt*self.k[i]) - - if self.limiter is not None: - self.limiter.apply(self.x1) - - if self.idx is None and len(self.fs) > 1: - self.xnph = tuple([self.dt*self.butcher_matrix[stage, stage]*a + b for a, b in zip(split(self.x_out), split(self.x1))]) - else: - self.xnph = self.x1 + self.butcher_matrix[stage, stage]*self.dt*self.x_out - solver = self.solvers[stage] - solver.solve() - - self.k[stage].assign(self.x_out) - - def apply(self, x_out, x_in): - - for i in range(self.nStages): - self.solve_stage(x_in, i) - - x_out.assign(x_in) - for i in range(self.nStages): - x_out.assign(x_out + self.butcher_matrix[self.nStages, i]*self.dt*self.k[i]) - - if self.limiter is not None: - self.limiter.apply(x_out) - - -class ExplicitMultistage(ExplicitTimeDiscretisation): - """ - A class for implementing general explicit multistage (Runge-Kutta) - methods based on its Butcher tableau. - - A Butcher tableau is formed in the following way for a s-th order explicit scheme: \n - - All upper diagonal a_ij elements are zero for explicit methods. - - There are three steps to move from the current solution, y^n, to the new one, y^{n+1} - - For each i = 1, s in an s stage method - we have the intermediate solutions: \n - y_i = y^n + dt*(a_i1*k_1 + a_i2*k_2 + ... + a_i{i-1}*k_{i-1}) \n - We compute the gradient at the intermediate location, k_i = F(y_i) \n - - At the last stage, compute the new solution by: - y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s) \n - - """ - # --------------------------------------------------------------------------- - # Butcher tableau for a s-th order - # explicit scheme: - # c_0 | 0 0 . 0 - # c_1 | a_10 0 . 0 - # . | . . . . - # . | . . . . - # c_s | a_s0 a_s1 . 0 - # ------------------------- - # | b_1 b_2 ... b_s - # - # - # The corresponding square 'butcher_matrix' is: - # - # [a_10 0 . 0 ] - # [a_20 a_21 . 0 ] - # [ . . . . ] - # [ b_0 b_1 . b_s] - # --------------------------------------------------------------------------- - - def __init__(self, domain, butcher_matrix, field_name=None, - fixed_subcycles=None, subcycle_by_courant=None, - increment_form=True, solver_parameters=None, - limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - butcher_matrix (numpy array): A matrix containing the coefficients of - a butcher tableau defining a given Runge Kutta time discretisation. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - fixed_subcycles (int, optional): the fixed number of sub-steps to - perform. This option cannot be specified with the - `subcycle_by_courant` argument. Defaults to None. - subcycle_by_courant (float, optional): specifying this option will - make the scheme perform adaptive sub-cycling based on the - Courant number. The specified argument is the maximum Courant - for one sub-cycle. Defaults to None, in which case adaptive - sub-cycling is not used. This option cannot be specified with the - `fixed_subcycles` argument. - increment_form (bool, optional): whether to write the RK scheme in - "increment form", solving for increments rather than updated - fields. Defaults to True. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - super().__init__(domain, field_name=field_name, - fixed_subcycles=fixed_subcycles, - subcycle_by_courant=subcycle_by_courant, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - self.butcher_matrix = butcher_matrix - self.nbutcher = int(np.shape(self.butcher_matrix)[0]) - self.increment_form = increment_form - - @property - def nStages(self): - return self.nbutcher - - def setup(self, equation, apply_bcs=True, *active_labels): - """ - Set up the time discretisation based on the equation. - - Args: - equation (:class:`PrognosticEquation`): the model's equation. - *active_labels (:class:`Label`): labels indicating which terms of - the equation to include. - """ - super().setup(equation, apply_bcs, *active_labels) - - if not self.increment_form: - self.field_i = [Function(self.fs) for i in range(self.nStages+1)] - else: - self.k = [Function(self.fs) for i in range(self.nStages)] - - @cached_property - def solver(self): - if self.increment_form: - return super().solver - else: - # In this case, don't set snes_type to ksp only, as we do want the - # outer Newton iteration - solver_list = [] - - for stage in range(self.nStages): - # setup linear solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem( - self.lhs[stage].form - self.rhs[stage].form, - self.field_i[stage+1], bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__+str(stage) - solver = NonlinearVariationalSolver( - problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - solver_list.append(solver) - return solver_list - - @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - - if self.increment_form: - l = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x_out, self.idx), - map_if_false=drop) - - return l.form - - else: - lhs_list = [] - for stage in range(self.nStages): - l = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.field_i[stage+1], self.idx), - map_if_false=drop) - lhs_list.append(l) - - return lhs_list - - @cached_property - def rhs(self): - """Set up the time discretisation's right hand side.""" - - if self.increment_form: - r = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) - - r = r.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=lambda t: -1*t) - - # If there are no active labels, we may have no terms at this point - # So that we can still do xnp1 = xn, put in a zero term here - if self.increment_form and len(r.terms) == 0: - logger.warning('No terms detected for RHS of explicit problem. ' - + 'Adding a zero term to avoid failure.') - null_term = Constant(0.0)*self.residual.label_map( - lambda t: t.has_label(time_derivative), - # Drop label from this - map_if_true=lambda t: time_derivative.remove(t), - map_if_false=drop) - r += null_term - - return r.form - - else: - rhs_list = [] - - for stage in range(self.nStages): - r = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.field_i[0], old_idx=self.idx)) - - r = r.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=keep, - map_if_false=lambda t: -self.butcher_matrix[stage, 0]*self.dt*t) - - for i in range(1, stage+1): - r_i = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=replace_subject(self.field_i[i], old_idx=self.idx) - ) - - r -= self.butcher_matrix[stage, i]*self.dt*r_i - - rhs_list.append(r) - - return rhs_list - - def solve_stage(self, x0, stage): - - if self.increment_form: - self.x1.assign(x0) - - for i in range(stage): - self.x1.assign(self.x1 + self.dt*self.butcher_matrix[stage-1, i]*self.k[i]) - for evaluate in self.evaluate_source: - evaluate(self.x1, self.dt) - if self.limiter is not None: - self.limiter.apply(self.x1) - self.solver.solve() - - self.k[stage].assign(self.x_out) - - if (stage == self.nStages - 1): - self.x1.assign(x0) - for i in range(self.nStages): - self.x1.assign(self.x1 + self.dt*self.butcher_matrix[stage, i]*self.k[i]) - self.x1.assign(self.x1) - - if self.limiter is not None: - self.limiter.apply(self.x1) - - else: - # Set initial field - if stage == 0: - self.field_i[0].assign(x0) - - # Use x0 as a first guess (otherwise may not converge) - self.field_i[stage+1].assign(x0) - - # Update field_i for physics / limiters - for evaluate in self.evaluate_source: - # TODO: not implemented! Here we need to evaluate the m-th term - # in the i-th RHS with field_m - raise NotImplementedError( - 'Physics not implemented with RK schemes that do not use ' - + 'the increment form') - if self.limiter is not None: - self.limiter.apply(self.field_i[stage]) - - # Obtain field_ip1 = field_n - dt* sum_m{a_im*F[field_m]} - self.solver[stage].solve() - - if (stage == self.nStages - 1): - self.x1.assign(self.field_i[stage+1]) - if self.limiter is not None: - self.limiter.apply(self.x1) - - def apply_cycle(self, x_out, x_in): - """ - Apply the time discretisation through a single sub-step. - - Args: - x_in (:class:`Function`): the input field. - x_out (:class:`Function`): the output field to be computed. - """ - if self.limiter is not None: - self.limiter.apply(x_in) - - self.x1.assign(x_in) - - for i in range(self.nStages): - self.solve_stage(x_in, i) - x_out.assign(self.x1) - - -class ForwardEuler(ExplicitMultistage): - """ - Implements the forward Euler timestepping scheme. - - The forward Euler method for operator F is the most simple explicit scheme: \n - k0 = F[y^n] \n - y^(n+1) = y^n + dt*k0 \n - """ - def __init__(self, domain, field_name=None, fixed_subcycles=None, - subcycle_by_courant=None, increment_form=True, - solver_parameters=None, limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - fixed_subcycles (int, optional): the fixed number of sub-steps to - perform. This option cannot be specified with the - `subcycle_by_courant` argument. Defaults to None. - subcycle_by_courant (float, optional): specifying this option will - make the scheme perform adaptive sub-cycling based on the - Courant number. The specified argument is the maximum Courant - for one sub-cycle. Defaults to None, in which case adaptive - sub-cycling is not used. This option cannot be specified with the - `fixed_subcycles` argument. - increment_form (bool, optional): whether to write the RK scheme in - "increment form", solving for increments rather than updated - fields. Defaults to True. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - butcher_matrix = np.array([1.]).reshape(1, 1) - super().__init__(domain, butcher_matrix, field_name=field_name, - fixed_subcycles=fixed_subcycles, - subcycle_by_courant=subcycle_by_courant, - increment_form=increment_form, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - - -class SSPRK3(ExplicitMultistage): - u""" - Implements the 3-stage Strong-Stability-Preserving Runge-Kutta method - for solving ∂y/∂t = F(y). It can be written as: \n - - k0 = F[y^n] \n - k1 = F[y^n + dt*k1] \n - k2 = F[y^n + (1/4)*dt*(k0+k1)] \n - y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + 4*k2) \n - """ - def __init__(self, domain, field_name=None, fixed_subcycles=None, - subcycle_by_courant=None, increment_form=True, - solver_parameters=None, limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - fixed_subcycles (int, optional): the fixed number of sub-steps to - perform. This option cannot be specified with the - `subcycle_by_courant` argument. Defaults to None. - subcycle_by_courant (float, optional): specifying this option will - make the scheme perform adaptive sub-cycling based on the - Courant number. The specified argument is the maximum Courant - for one sub-cycle. Defaults to None, in which case adaptive - sub-cycling is not used. This option cannot be specified with the - `fixed_subcycles` argument. - increment_form (bool, optional): whether to write the RK scheme in - "increment form", solving for increments rather than updated - fields. Defaults to True. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - butcher_matrix = np.array([[1., 0., 0.], [1./4., 1./4., 0.], [1./6., 1./6., 2./3.]]) - - super().__init__(domain, butcher_matrix, field_name=field_name, - fixed_subcycles=fixed_subcycles, - subcycle_by_courant=subcycle_by_courant, - increment_form=increment_form, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - - -class RK4(ExplicitMultistage): - u""" - Implements the classic 4-stage Runge-Kutta method. - - The classic 4-stage Runge-Kutta method for solving ∂y/∂t = F(y). It can be - written as: \n - - k0 = F[y^n] \n - k1 = F[y^n + 1/2*dt*k1] \n - k2 = F[y^n + 1/2*dt*k2] \n - k3 = F[y^n + dt*k3] \n - y^(n+1) = y^n + (1/6) * dt * (k0 + 2*k1 + 2*k2 + k3) \n - - where superscripts indicate the time-level. \n - """ - def __init__(self, domain, field_name=None, fixed_subcycles=None, - subcycle_by_courant=None, increment_form=True, - solver_parameters=None, - limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - fixed_subcycles (int, optional): the fixed number of sub-steps to - perform. This option cannot be specified with the - `subcycle_by_courant` argument. Defaults to None. - subcycle_by_courant (float, optional): specifying this option will - make the scheme perform adaptive sub-cycling based on the - Courant number. The specified argument is the maximum Courant - for one sub-cycle. Defaults to None, in which case adaptive - sub-cycling is not used. This option cannot be specified with the - `fixed_subcycles` argument. - increment_form (bool, optional): whether to write the RK scheme in - "increment form", solving for increments rather than updated - fields. Defaults to True. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - butcher_matrix = np.array([[0.5, 0., 0., 0.], [0., 0.5, 0., 0.], [0., 0., 1., 0.], [1./6., 1./3., 1./3., 1./6.]]) - super().__init__(domain, butcher_matrix, field_name=field_name, - fixed_subcycles=fixed_subcycles, - subcycle_by_courant=subcycle_by_courant, - increment_form=increment_form, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - - -class Heun(ExplicitMultistage): - u""" - Implements Heun's method. - - The 2-stage Runge-Kutta scheme known as Heun's method,for solving - ∂y/∂t = F(y). It can be written as: \n - - y_1 = F[y^n] \n - y^(n+1) = (1/2)y^n + (1/2)F[y_1] \n - - where superscripts indicate the time-level and subscripts indicate the stage - number. - """ - def __init__(self, domain, field_name=None, fixed_subcycles=None, - subcycle_by_courant=None, increment_form=True, - solver_parameters=None, limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - fixed_subcycles (int, optional): the fixed number of sub-steps to - perform. This option cannot be specified with the - `subcycle_by_courant` argument. Defaults to None. - subcycle_by_courant (float, optional): specifying this option will - make the scheme perform adaptive sub-cycling based on the - Courant number. The specified argument is the maximum Courant - for one sub-cycle. Defaults to None, in which case adaptive - sub-cycling is not used. This option cannot be specified with the - `fixed_subcycles` argument. - increment_form (bool, optional): whether to write the RK scheme in - "increment form", solving for increments rather than updated - fields. Defaults to True. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - butcher_matrix = np.array([[1., 0.], [0.5, 0.5]]) - super().__init__(domain, butcher_matrix, field_name=field_name, - fixed_subcycles=fixed_subcycles, - subcycle_by_courant=subcycle_by_courant, - increment_form=increment_form, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - - -class BackwardEuler(TimeDiscretisation): - """ - Implements the backward Euler timestepping scheme. - - The backward Euler method for operator F is the most simple implicit scheme: \n - y^(n+1) = y^n + dt*F[y^(n+1)]. \n - """ - def __init__(self, domain, field_name=None, solver_parameters=None, - limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - fixed_subcycles (int, optional): the number of sub-steps to perform. - Defaults to None. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods. Defaults to None. - """ - if not solver_parameters: - # default solver parameters - solver_parameters = {'ksp_type': 'gmres', - 'pc_type': 'bjacobi', - 'sub_pc_type': 'ilu'} - super().__init__(domain=domain, field_name=field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - - @property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: self.dt*t) - - return l.form - - @property - def rhs(self): - """Set up the time discretisation's right hand side.""" - r = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x1, old_idx=self.idx), - map_if_false=drop) - - return r.form - - def apply(self, x_out, x_in): - """ - Apply the time discretisation to advance one whole time step. - - Args: - x_out (:class:`Function`): the output field to be computed. - x_in (:class:`Function`): the input field. - """ - for evaluate in self.evaluate_source: - evaluate(x_in, self.dt) - - if len(self.evaluate_source) > 0: - # If we have physics, use x_in as first guess - self.x_out.assign(x_in) - - self.x1.assign(x_in) - self.solver.solve() - x_out.assign(self.x_out) - - -class ThetaMethod(TimeDiscretisation): - """ - Implements the theta implicit-explicit timestepping method, which can - be thought as a generalised trapezium rule. - - The theta implicit-explicit timestepping method for operator F is written as: \n - y^(n+1) = y^n + dt*(1-theta)*F[y^n] + dt*theta*F[y^(n+1)] \n - for off-centring parameter theta. \n - """ - - def __init__(self, domain, theta, field_name=None, - solver_parameters=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - theta (float): the off-centring parameter. theta = 1 - corresponds to a backward Euler method. Defaults to None. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - - Raises: - ValueError: if theta is not provided. - """ - if (theta < 0 or theta > 1): - raise ValueError("please provide a value for theta between 0 and 1") - if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): - raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") - if not solver_parameters: - # theta method leads to asymmetric matrix, per lhs function below, - # so don't use CG - solver_parameters = {'ksp_type': 'gmres', - 'pc_type': 'bjacobi', - 'sub_pc_type': 'ilu'} - - super().__init__(domain, field_name, - solver_parameters=solver_parameters, - options=options) - - self.theta = theta - - @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: self.theta*self.dt*t) - - return l.form - - @cached_property - def rhs(self): - """Set up the time discretisation's right hand side.""" - r = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -(1-self.theta)*self.dt*t) - - return r.form - - def apply(self, x_out, x_in): - """ - Apply the time discretisation to advance one whole time step. - - Args: - x_out (:class:`Function`): the output field to be computed. - x_in (:class:`Function`): the input field. - """ - self.x1.assign(x_in) - self.solver.solve() - x_out.assign(self.x_out) - - -class TrapeziumRule(ThetaMethod): - """ - Implements the trapezium rule timestepping method, also commonly known as - Crank Nicholson. - - The trapezium rule timestepping method for operator F is written as: \n - y^(n+1) = y^n + dt/2*F[y^n] + dt/2*F[y^(n+1)]. \n - It is equivalent to the "theta" method with theta = 1/2. \n - """ - - def __init__(self, domain, field_name=None, solver_parameters=None, - options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - super().__init__(domain, 0.5, field_name, - solver_parameters=solver_parameters, - options=options) - - -class MultilevelTimeDiscretisation(TimeDiscretisation): - """Base class for multi-level timesteppers""" - - def __init__(self, domain, field_name=None, solver_parameters=None, - limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): - raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") - super().__init__(domain=domain, field_name=field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - self.initial_timesteps = 0 - - @abstractproperty - def nlevels(self): - pass - - def setup(self, equation, apply_bcs=True, *active_labels): - super().setup(equation=equation, apply_bcs=apply_bcs, *active_labels) - for n in range(self.nlevels, 1, -1): - setattr(self, "xnm%i" % (n-1), Function(self.fs)) - - -class BDF2(MultilevelTimeDiscretisation): - """ - Implements the implicit multistep BDF2 timestepping method. - - The BDF2 timestepping method for operator F is written as: \n - y^(n+1) = (4/3)*y^n - (1/3)*y^(n-1) + (2/3)*dt*F[y^(n+1)] \n - """ - - @property - def nlevels(self): - return 2 - - @property - def lhs0(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: self.dt*t) - - return l.form - - @property - def rhs0(self): - """Set up the time discretisation's right hand side for inital BDF step.""" - r = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x1, old_idx=self.idx), - map_if_false=drop) - - return r.form - - @property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: (2/3)*self.dt*t) - - return l.form - - @property - def rhs(self): - """Set up the time discretisation's right hand side for BDF2 steps.""" - xn = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x1, old_idx=self.idx), - map_if_false=drop) - xnm1 = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.xnm1, old_idx=self.idx), - map_if_false=drop) - - r = (4/3.) * xn - (1/3.) * xnm1 - - return r.form - - @property - def solver0(self): - """Set up the problem and the solver for initial BDF step.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs0-self.rhs0, self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__+"0" - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - @property - def solver(self): - """Set up the problem and the solver for BDF2 steps.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__ - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - def apply(self, x_out, *x_in): - """ - Apply the time discretisation to advance one whole time step. - - Args: - x_out (:class:`Function`): the output field to be computed. - x_in (:class:`Function`): the input field(s). - """ - if self.initial_timesteps < self.nlevels-1: - self.initial_timesteps += 1 - solver = self.solver0 - else: - solver = self.solver - - self.xnm1.assign(x_in[0]) - self.x1.assign(x_in[1]) - solver.solve() - x_out.assign(self.x_out) - - -class TR_BDF2(TimeDiscretisation): - """ - Implements the two stage implicit TR-BDF2 time stepping method, with a - trapezoidal stage (TR) followed by a second order backwards difference stage (BDF2). - - The TR-BDF2 time stepping method for operator F is written as: \n - y^(n+g) = y^n + dt*g/2*F[y^n] + dt*g/2*F[y^(n+g)] (TR stage) \n - y^(n+1) = 1/(g(2-g))*y^(n+g) - (1-g)**2/(g(2-g))*y^(n) + (1-g)/(2-g)*dt*F[y^(n+1)] (BDF2 stage) \n - for an off-centring parameter g (gamma). \n - """ - def __init__(self, domain, gamma, field_name=None, - solver_parameters=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - gamma (float): the off-centring parameter - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - if (gamma < 0. or gamma > 1.): - raise ValueError("please provide a value for gamma between 0 and 1") - if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): - raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") - if not solver_parameters: - # theta method leads to asymmetric matrix, per lhs function below, - # so don't use CG - solver_parameters = {'ksp_type': 'gmres', - 'pc_type': 'bjacobi', - 'sub_pc_type': 'ilu'} - - super().__init__(domain, field_name, - solver_parameters=solver_parameters, - options=options) - - self.gamma = gamma - - def setup(self, equation, apply_bcs=True, *active_labels): - super().setup(equation, apply_bcs, *active_labels) - self.xnpg = Function(self.fs) - self.xn = Function(self.fs) - - @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative) for the TR stage.""" - l = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.xnpg, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: 0.5*self.gamma*self.dt*t) - - return l.form - - @cached_property - def rhs(self): - """Set up the time discretisation's right hand side for the TR stage.""" - r = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.xn, old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -0.5*self.gamma*self.dt*t) - - return r.form - - @cached_property - def lhs_bdf2(self): - """Set up the discretisation's left hand side (the time derivative) for the BDF2 stage.""" - l = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: ((1.0-self.gamma)/(2.0-self.gamma))*self.dt*t) - - return l.form - - @cached_property - def rhs_bdf2(self): - """Set up the time discretisation's right hand side for the BDF2 stage.""" - xn = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.xn, old_idx=self.idx), - map_if_false=drop) - xnpg = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.xnpg, old_idx=self.idx), - map_if_false=drop) - - r = (1.0/(self.gamma*(2.0-self.gamma)))*xnpg - ((1.0-self.gamma)**2/(self.gamma*(2.0-self.gamma)))*xn - - return r.form - - @cached_property - def solver_tr(self): - """Set up the problem and the solver.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.xnpg, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__+"_tr" - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - @cached_property - def solver_bdf2(self): - """Set up the problem and the solver.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs_bdf2-self.rhs_bdf2, self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__+"_bdf2" - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - def apply(self, x_out, x_in): - """ - Apply the time discretisation to advance one whole time step. - - Args: - x_out (:class:`Function`): the output field to be computed. - x_in (:class:`Function`): the input field(s). - """ - self.xn.assign(x_in) - self.solver_tr.solve() - self.solver_bdf2.solve() - x_out.assign(self.x_out) - - -class Leapfrog(MultilevelTimeDiscretisation): - """ - Implements the multistep Leapfrog timestepping method. - - The Leapfrog timestepping method for operator F is written as: \n - y^(n+1) = y^(n-1) + 2*dt*F[y^n] \n - """ - @property - def nlevels(self): - return 2 - - @property - def rhs0(self): - """Set up the discretisation's right hand side for initial forward euler step.""" - r = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*t) - - return r.form - - @property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - return super(Leapfrog, self).lhs - - @property - def rhs(self): - """Set up the discretisation's right hand side for leapfrog steps.""" - r = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=replace_subject(self.x1, old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.xnm1, old_idx=self.idx), - map_if_false=lambda t: -2.0*self.dt*t) - - return r.form - - @property - def solver0(self): - """Set up the problem and the solver for initial forward euler step.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs0, self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__+"0" - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - @property - def solver(self): - """Set up the problem and the solver for leapfrog steps.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__ - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - def apply(self, x_out, *x_in): - """ - Apply the time discretisation to advance one whole time step. - - Args: - x_out (:class:`Function`): the output field to be computed. - x_in (:class:`Function`): the input field(s). - """ - if self.initial_timesteps < self.nlevels-1: - self.initial_timesteps += 1 - solver = self.solver0 - else: - solver = self.solver - - self.xnm1.assign(x_in[0]) - self.x1.assign(x_in[1]) - solver.solve() - x_out.assign(self.x_out) - - -class AdamsBashforth(MultilevelTimeDiscretisation): - """ - Implements the explicit multistep Adams-Bashforth timestepping - method of general order up to 5. - - The general AB timestepping method for operator F is written as: \n - y^(n+1) = y^n + dt*(b_0*F[y^(n)] + b_1*F[y^(n-1)] + b_2*F[y^(n-2)] + b_3*F[y^(n-3)] + b_4*F[y^(n-4)]) \n - """ - def __init__(self, domain, order, field_name=None, - solver_parameters=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - order (float, optional): order of scheme - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - - Raises: - ValueError: if order is not provided, or is in incorrect range. - """ - - if (order > 5 or order < 1): - raise ValueError("Adams-Bashforth of order greater than 5 not implemented") - if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): - raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") - - super().__init__(domain, field_name, - solver_parameters=solver_parameters, - options=options) - - self.order = order - - def setup(self, equation, apply_bcs=True, *active_labels): - super().setup(equation=equation, apply_bcs=apply_bcs, - *active_labels) - - self.x = [Function(self.fs) for i in range(self.nlevels)] - - if (self.order == 1): - self.b = [1.0] - elif (self.order == 2): - self.b = [-(1.0/2.0), (3.0/2.0)] - elif (self.order == 3): - self.b = [(5.0)/(12.0), -(16.0)/(12.0), (23.0)/(12.0)] - elif (self.order == 4): - self.b = [-(9.0)/(24.0), (37.0)/(24.0), -(59.0)/(24.0), (55.0)/(24.0)] - elif (self.order == 5): - self.b = [(251.0)/(720.0), -(1274.0)/(720.0), (2616.0)/(720.0), - -(2774.0)/(720.0), (2901.0)/(720.0)] - - @property - def nlevels(self): - return self.order - - @property - def rhs0(self): - """Set up the discretisation's right hand side for initial forward euler step.""" - r = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*t) - - return r.form - - @property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - return super(AdamsBashforth, self).lhs - - @property - def rhs(self): - """Set up the discretisation's right hand side for Adams Bashforth steps.""" - r = self.residual.label_map(all_terms, - map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.b[-1]*self.dt*t) - for n in range(self.nlevels-1): - rtemp = self.residual.label_map(lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=replace_subject(self.x[n], old_idx=self.idx)) - rtemp = rtemp.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*self.b[n]*t) - r += rtemp - return r.form - - @property - def solver0(self): - """Set up the problem and the solverfor initial forward euler step.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs0, self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__+"0" - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - @property - def solver(self): - """Set up the problem and the solver for Adams Bashforth steps.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__ - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - def apply(self, x_out, *x_in): - """ - Apply the time discretisation to advance one whole time step. - - Args: - x_out (:class:`Function`): the output field to be computed. - x_in (:class:`Function`): the input field(s). - """ - if self.initial_timesteps < self.nlevels-1: - self.initial_timesteps += 1 - solver = self.solver0 - else: - solver = self.solver - - for n in range(self.nlevels): - self.x[n].assign(x_in[n]) - solver.solve() - x_out.assign(self.x_out) - - -class AdamsMoulton(MultilevelTimeDiscretisation): - """ - Implements the implicit multistep Adams-Moulton - timestepping method of general order up to 5 - - The general AM timestepping method for operator F is written as \n - y^(n+1) = y^n + dt*(b_0*F[y^(n+1)] + b_1*F[y^(n)] + b_2*F[y^(n-1)] + b_3*F[y^(n-2)]) \n - """ - def __init__(self, domain, order, field_name=None, - solver_parameters=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - order (float, optional): order of scheme - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - - Raises: - ValueError: if order is not provided, or is in incorrect range. - """ - if (order > 4 or order < 1): - raise ValueError("Adams-Moulton of order greater than 5 not implemented") - if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): - raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") - if not solver_parameters: - solver_parameters = {'ksp_type': 'gmres', - 'pc_type': 'bjacobi', - 'sub_pc_type': 'ilu'} - - super().__init__(domain, field_name, - solver_parameters=solver_parameters, - options=options) - - self.order = order - - def setup(self, equation, apply_bcs=True, *active_labels): - super().setup(equation=equation, apply_bcs=apply_bcs, *active_labels) - - self.x = [Function(self.fs) for i in range(self.nlevels)] - - if (self.order == 1): - self.bl = (1.0/2.0) - self.br = [(1.0/2.0)] - elif (self.order == 2): - self.bl = (5.0/12.0) - self.br = [-(1.0/12.0), (8.0/12.0)] - elif (self.order == 3): - self.bl = (9.0/24.0) - self.br = [(1.0/24.0), -(5.0/24.0), (19.0/24.0)] - elif (self.order == 4): - self.bl = (251.0/720.0) - self.br = [-(19.0/720.0), (106.0/720.0), -(254.0/720.0), (646.0/720.0)] - - @property - def nlevels(self): - return self.order - - @property - def rhs0(self): - """Set up the discretisation's right hand side for initial trapezoidal step.""" - r = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -0.5*self.dt*t) - - return r.form - - @property - def lhs0(self): - """Set up the time discretisation's right hand side for initial trapezoidal step.""" - l = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: 0.5*self.dt*t) - return l.form - - @property - def lhs(self): - """Set up the time discretisation's right hand side for Adams Moulton steps.""" - l = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: self.bl*self.dt*t) - return l.form - - @property - def rhs(self): - """Set up the discretisation's right hand side for Adams Moulton steps.""" - r = self.residual.label_map(all_terms, - map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.br[-1]*self.dt*t) - for n in range(self.nlevels-1): - rtemp = self.residual.label_map(lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=replace_subject(self.x[n], old_idx=self.idx)) - rtemp = rtemp.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*self.br[n]*t) - r += rtemp - return r.form - - @property - def solver0(self): - """Set up the problem and the solver for initial trapezoidal step.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs0-self.rhs0, self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__+"0" - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - @property - def solver(self): - """Set up the problem and the solver for Adams Moulton steps.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__ - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - def apply(self, x_out, *x_in): - """ - Apply the time discretisation to advance one whole time step. - - Args: - x_out (:class:`Function`): the output field to be computed. - x_in (:class:`Function`): the input field(s). - """ - if self.initial_timesteps < self.nlevels-1: - self.initial_timesteps += 1 - solver = self.solver0 - else: - solver = self.solver - - for n in range(self.nlevels): - self.x[n].assign(x_in[n]) - solver.solve() - x_out.assign(self.x_out) - - -class ImplicitMidpoint(ImplicitMultistage): - u""" - Implements the Implicit Midpoint method as a 1-stage Runge Kutta method. - - The method, for solving - ∂y/∂t = F(y), can be written as: \n - - k0 = F[y^n + 0.5*dt*k0] \n - y^(n+1) = y^n + dt*k0 \n - """ - def __init__(self, domain, field_name=None, solver_parameters=None, - limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - butcher_matrix = np.array([[0.5], [1.]]) - super().__init__(domain, butcher_matrix, field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - - -class QinZhang(ImplicitMultistage): - u""" - Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method. - - The method, for solving - ∂y/∂t = F(y), can be written as: \n - - k0 = F[y^n + 0.25*dt*k0] \n - k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] \n - y^(n+1) = y^n + 0.5*dt*(k0 + k1) \n - """ - def __init__(self, domain, field_name=None, solver_parameters=None, - limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - butcher_matrix = np.array([[0.25, 0], [0.5, 0.25], [0.5, 0.5]]) - super().__init__(domain, butcher_matrix, field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - - -class IMEX_Euler(IMEXMultistage): - u""" - Implements IMEX Euler one-stage method. - - The method, for solving \n - ∂y/∂t = F(y) + S(y), can be written as: \n - - y_0 = y^n \n - y_1 = y^n + dt*F[y_1] + dt*S[y_0] \n - y^(n+1) = y^n + dt*F[y_1] + dt*S[y_0] \n - """ - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - butcher_imp = np.array([[0., 0.], [0., 1.], [0., 1.]]) - butcher_exp = np.array([[0., 0.], [1., 0.], [1., 0.]]) - super().__init__(domain, butcher_imp, butcher_exp, field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - - -class ARS3(IMEXMultistage): - u""" - Implements ARS3(2,3,3) two-stage IMEX Runge–Kutta method - from RK IMEX for HEVI (Weller et al 2013). - Where g = (3 + sqrt(3))/6. - - The method, for solving \n - ∂y/∂t = F(y) + S(y), can be written as: \n - - y_0 = y^n \n - y_1 = y^n + dt*g*F[y_1] + dt*g*S[y_0] \n - y_2 = y^n + dt*((1-2g)*F[y_1]+g*F[y_2]) \n - + dt*((g-1)*S[y_0]+2(g-1)*S[y_1]) \n - y^(n+1) = y^n + dt*(g*F[y_1]+(1-g)*F[y_2]) \n - + dt*(0.5*S[y_1]+0.5*S[y_2]) \n - """ - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - g = (3. + np.sqrt(3.))/6. - butcher_imp = np.array([[0., 0., 0.], [0., g, 0.], [0., 1-2.*g, g], [0., 0.5, 0.5]]) - butcher_exp = np.array([[0., 0., 0.], [g, 0., 0.], [g-1., 2.*(1.-g), 0.], [0., 0.5, 0.5]]) - - super().__init__(domain, butcher_imp, butcher_exp, field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - - -class ARK2(IMEXMultistage): - u""" - Implements ARK2(2,3,2) two-stage IMEX Runge–Kutta method from - RK IMEX for HEVI (Weller et al 2013). - Where g = 1 - 1/sqrt(2), a = 1/6(3 + 2sqrt(2)), d = 1/2sqrt(2). - - The method, for solving \n - ∂y/∂t = F(y) + S(y), can be written as: \n - - y_0 = y^n \n - y_1 = y^n + dt*(g*F[y_0]+g*F[y_1]) + 2*dt*g*S[y_0] \n - y_2 = y^n + dt*(d*F[y_0]+d*F[y_1]+g*F[y_2]) \n - + dt*((1-a)*S[y_0]+a*S[y_1]) \n - y^(n+1) = y^n + dt*(d*F[y_0]+d*F[y_1]+g*F[y_2]) \n - + dt*(d*S[y_0]+d*S[y_1]+g*S[y_2]) \n - """ - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - g = 1. - 1./np.sqrt(2.) - d = 1./(2.*np.sqrt(2.)) - a = 1./6.*(3. + 2.*np.sqrt(2.)) - butcher_imp = np.array([[0., 0., 0.], [g, g, 0.], [d, d, g], [d, d, g]]) - butcher_exp = np.array([[0., 0., 0.], [2.*g, 0., 0.], [1.-a, a, 0.], [d, d, g]]) - super().__init__(domain, butcher_imp, butcher_exp, field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - - -class SSP3(IMEXMultistage): - u""" - Implements SSP3(3,3,2) three-stage IMEX Runge–Kutta method from RK IMEX for HEVI (Weller et al 2013). - Where g = 1 - 1/sqrt(2) - - The method, for solving \n - ∂y/∂t = F(y) + S(y), can be written as: \n - - y_1 = y^n + dt*g*F[y_1] \n - y_2 = y^n + dt*((1-2g)*F[y_1]+g*F[y_2]) + dt*S[y_1] \n - y_3 = y^n + dt*((0.5-g)*F[y_1]+g*F[y_3]) + dt*(0.25*S[y_1]+0.25*S[y_2]) \n - y^(n+1) = y^n + dt*(1/6*F[y_1]+1/6*F[y_2]+2/3*F[y_3]) \n - + dt*(1/6*S[y_1]+1/6*S[y_2]+2/3*S[y_3]) \n - """ - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - g = 1. - (1./np.sqrt(2.)) - butcher_imp = np.array([[g, 0., 0.], [1-2.*g, g, 0.], [0.5-g, 0., g], [(1./6.), (1./6.), (2./3.)]]) - butcher_exp = np.array([[0., 0., 0.], [1., 0., 0.], [0.25, 0.25, 0.], [(1./6.), (1./6.), (2./3.)]]) - super().__init__(domain, butcher_imp, butcher_exp, field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - - -class Trap2(IMEXMultistage): - u""" - Implements Trap2(2+e,3,2) three-stage IMEX Runge–Kutta method from RK IMEX for HEVI (Weller et al 2013). - For e = 1 or 0. - - The method, for solving \n - ∂y/∂t = F(y) + S(y), can be written as: \n - - y_0 = y^n \n - y_1 = y^n + dt*e*F[y_0] + dt*S[y_0] \n - y_2 = y^n + dt*(0.5*F[y_0]+0.5*F[y_2]) + dt*(0.5*S[y_0]+0.5*S[y_1]) \n - y_3 = y^n + dt*(0.5*F[y_0]+0.5*F[y_3]) + dt*(0.5*S[y_0]+0.5*S[y_2]) \n - y^(n+1) = y^n + dt*(0.5*F[y_0]+0.5*F[y_3]) + dt*(0.5*S[y_0] + 0.5*S[y_2]) \n - """ - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): - """ - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - field_name (str, optional): name of the field to be evolved. - Defaults to None. - solver_parameters (dict, optional): dictionary of parameters to - pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. - options (:class:`AdvectionOptions`, optional): an object containing - options to either be passed to the spatial discretisation, or - to control the "wrapper" methods, such as Embedded DG or a - recovery method. Defaults to None. - """ - e = 0. - butcher_imp = np.array([[0., 0., 0., 0.], [e, 0., 0., 0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0., 0.5], [0.5, 0., 0., 0.5]]) - butcher_exp = np.array([[0., 0., 0., 0.], [1., 0., 0., 0.], [0.5, 0.5, 0., 0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0.5, 0.]]) - super().__init__(domain, butcher_imp, butcher_exp, field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) diff --git a/gusto/time_discretisation/__init__.py b/gusto/time_discretisation/__init__.py new file mode 100644 index 000000000..548cf9c40 --- /dev/null +++ b/gusto/time_discretisation/__init__.py @@ -0,0 +1,6 @@ +from gusto.time_discretisation.time_discretisation import * # noqa +from gusto.time_discretisation.explicit_runge_kutta import * # noqa +from gusto.time_discretisation.implicit_runge_kutta import * # noqa +from gusto.time_discretisation.imex_runge_kutta import * # noqa +from gusto.time_discretisation.multi_level_schemes import * # noqa +from gusto.time_discretisation.wrappers import * # noqa \ No newline at end of file diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py new file mode 100644 index 000000000..64e9545e0 --- /dev/null +++ b/gusto/time_discretisation/explicit_runge_kutta.py @@ -0,0 +1,490 @@ +"""Objects to describe explicit multi-stage (Runge-Kutta) discretisations.""" + +import numpy as np + +from firedrake import (Function, Constant, NonlinearVariationalProblem, + NonlinearVariationalSolver) +from firedrake.fml import replace_subject, all_terms, drop, keep +from firedrake.utils import cached_property + +from gusto.core.labels import time_derivative +from gusto.core.logging import logger +from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation + + +__all__ = ["ForwardEuler", "ExplicitRungeKutta", "SSPRK3", "RK4", "Heun"] + + +class ExplicitRungeKutta(ExplicitTimeDiscretisation): + """ + A class for implementing general explicit multistage (Runge-Kutta) + methods based on its Butcher tableau. + + A Butcher tableau is formed in the following way for a s-th order explicit + scheme: \n + + All upper diagonal a_ij elements are zero for explicit methods. + + There are three steps to move from the current solution, y^n, to the new + one, y^{n+1} + + For each i = 1, s in an s stage method + we have the intermediate solutions: \n + y_i = y^n + dt*(a_i1*k_1 + a_i2*k_2 + ... + a_i{i-1}*k_{i-1}) \n + We compute the gradient at the intermediate location, k_i = F(y_i) \n + + At the last stage, compute the new solution by: + y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s) \n + + """ + # --------------------------------------------------------------------------- + # Butcher tableau for a s-th order + # explicit scheme: + # c_0 | 0 0 . 0 + # c_1 | a_10 0 . 0 + # . | . . . . + # . | . . . . + # c_s | a_s0 a_s1 . 0 + # ------------------------- + # | b_1 b_2 ... b_s + # + # + # The corresponding square 'butcher_matrix' is: + # + # [a_10 0 . 0 ] + # [a_20 a_21 . 0 ] + # [ . . . . ] + # [ b_0 b_1 . b_s] + # --------------------------------------------------------------------------- + + def __init__(self, domain, butcher_matrix, field_name=None, + fixed_subcycles=None, subcycle_by_courant=None, + increment_form=True, solver_parameters=None, + limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + butcher_matrix (numpy array): A matrix containing the coefficients of + a butcher tableau defining a given Runge Kutta time discretisation. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + fixed_subcycles (int, optional): the fixed number of sub-steps to + perform. This option cannot be specified with the + `subcycle_by_courant` argument. Defaults to None. + subcycle_by_courant (float, optional): specifying this option will + make the scheme perform adaptive sub-cycling based on the + Courant number. The specified argument is the maximum Courant + for one sub-cycle. Defaults to None, in which case adaptive + sub-cycling is not used. This option cannot be specified with the + `fixed_subcycles` argument. + increment_form (bool, optional): whether to write the RK scheme in + "increment form", solving for increments rather than updated + fields. Defaults to True. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + super().__init__(domain, field_name=field_name, + fixed_subcycles=fixed_subcycles, + subcycle_by_courant=subcycle_by_courant, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + self.butcher_matrix = butcher_matrix + self.nbutcher = int(np.shape(self.butcher_matrix)[0]) + self.increment_form = increment_form + + @property + def nStages(self): + return self.nbutcher + + def setup(self, equation, apply_bcs=True, *active_labels): + """ + Set up the time discretisation based on the equation. + + Args: + equation (:class:`PrognosticEquation`): the model's equation. + *active_labels (:class:`Label`): labels indicating which terms of + the equation to include. + """ + super().setup(equation, apply_bcs, *active_labels) + + if not self.increment_form: + self.field_i = [Function(self.fs) for i in range(self.nStages+1)] + else: + self.k = [Function(self.fs) for i in range(self.nStages)] + + @cached_property + def solver(self): + if self.increment_form: + return super().solver + else: + # In this case, don't set snes_type to ksp only, as we do want the + # outer Newton iteration + solver_list = [] + + for stage in range(self.nStages): + # setup linear solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem( + self.lhs[stage].form - self.rhs[stage].form, + self.field_i[stage+1], bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__+str(stage) + solver = NonlinearVariationalSolver( + problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + solver_list.append(solver) + return solver_list + + @cached_property + def lhs(self): + """Set up the discretisation's left hand side (the time derivative).""" + + if self.increment_form: + l = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x_out, self.idx), + map_if_false=drop) + + return l.form + + else: + lhs_list = [] + for stage in range(self.nStages): + l = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.field_i[stage+1], self.idx), + map_if_false=drop) + lhs_list.append(l) + + return lhs_list + + @cached_property + def rhs(self): + """Set up the time discretisation's right hand side.""" + + if self.increment_form: + r = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x1, old_idx=self.idx)) + + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=drop, + map_if_false=lambda t: -1*t) + + # If there are no active labels, we may have no terms at this point + # So that we can still do xnp1 = xn, put in a zero term here + if self.increment_form and len(r.terms) == 0: + logger.warning('No terms detected for RHS of explicit problem. ' + + 'Adding a zero term to avoid failure.') + null_term = Constant(0.0)*self.residual.label_map( + lambda t: t.has_label(time_derivative), + # Drop label from this + map_if_true=lambda t: time_derivative.remove(t), + map_if_false=drop) + r += null_term + + return r.form + + else: + rhs_list = [] + + for stage in range(self.nStages): + r = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.field_i[0], old_idx=self.idx)) + + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=keep, + map_if_false=lambda t: -self.butcher_matrix[stage, 0]*self.dt*t) + + for i in range(1, stage+1): + r_i = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=drop, + map_if_false=replace_subject(self.field_i[i], old_idx=self.idx) + ) + + r -= self.butcher_matrix[stage, i]*self.dt*r_i + + rhs_list.append(r) + + return rhs_list + + def solve_stage(self, x0, stage): + + if self.increment_form: + self.x1.assign(x0) + + for i in range(stage): + self.x1.assign(self.x1 + self.dt*self.butcher_matrix[stage-1, i]*self.k[i]) + for evaluate in self.evaluate_source: + evaluate(self.x1, self.dt) + if self.limiter is not None: + self.limiter.apply(self.x1) + self.solver.solve() + + self.k[stage].assign(self.x_out) + + if (stage == self.nStages - 1): + self.x1.assign(x0) + for i in range(self.nStages): + self.x1.assign(self.x1 + self.dt*self.butcher_matrix[stage, i]*self.k[i]) + self.x1.assign(self.x1) + + if self.limiter is not None: + self.limiter.apply(self.x1) + + else: + # Set initial field + if stage == 0: + self.field_i[0].assign(x0) + + # Use x0 as a first guess (otherwise may not converge) + self.field_i[stage+1].assign(x0) + + # Update field_i for physics / limiters + for evaluate in self.evaluate_source: + # TODO: not implemented! Here we need to evaluate the m-th term + # in the i-th RHS with field_m + raise NotImplementedError( + 'Physics not implemented with RK schemes that do not use ' + + 'the increment form') + if self.limiter is not None: + self.limiter.apply(self.field_i[stage]) + + # Obtain field_ip1 = field_n - dt* sum_m{a_im*F[field_m]} + self.solver[stage].solve() + + if (stage == self.nStages - 1): + self.x1.assign(self.field_i[stage+1]) + if self.limiter is not None: + self.limiter.apply(self.x1) + + def apply_cycle(self, x_out, x_in): + """ + Apply the time discretisation through a single sub-step. + + Args: + x_in (:class:`Function`): the input field. + x_out (:class:`Function`): the output field to be computed. + """ + if self.limiter is not None: + self.limiter.apply(x_in) + + self.x1.assign(x_in) + + for i in range(self.nStages): + self.solve_stage(x_in, i) + x_out.assign(self.x1) + + +class ForwardEuler(ExplicitRungeKutta): + """ + Implements the forward Euler timestepping scheme. + + The forward Euler method for operator F is the most simple explicit + scheme: \n + k0 = F[y^n] \n + y^(n+1) = y^n + dt*k0 \n + """ + def __init__(self, domain, field_name=None, fixed_subcycles=None, + subcycle_by_courant=None, increment_form=True, + solver_parameters=None, limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + fixed_subcycles (int, optional): the fixed number of sub-steps to + perform. This option cannot be specified with the + `subcycle_by_courant` argument. Defaults to None. + subcycle_by_courant (float, optional): specifying this option will + make the scheme perform adaptive sub-cycling based on the + Courant number. The specified argument is the maximum Courant + for one sub-cycle. Defaults to None, in which case adaptive + sub-cycling is not used. This option cannot be specified with the + `fixed_subcycles` argument. + increment_form (bool, optional): whether to write the RK scheme in + "increment form", solving for increments rather than updated + fields. Defaults to True. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + butcher_matrix = np.array([1.]).reshape(1, 1) + super().__init__(domain, butcher_matrix, field_name=field_name, + fixed_subcycles=fixed_subcycles, + subcycle_by_courant=subcycle_by_courant, + increment_form=increment_form, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + + +class SSPRK3(ExplicitRungeKutta): + u""" + Implements the 3-stage Strong-Stability-Preserving Runge-Kutta method + for solving ∂y/∂t = F(y). It can be written as: \n + + k0 = F[y^n] \n + k1 = F[y^n + dt*k1] \n + k2 = F[y^n + (1/4)*dt*(k0+k1)] \n + y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + 4*k2) \n + """ + def __init__(self, domain, field_name=None, fixed_subcycles=None, + subcycle_by_courant=None, increment_form=True, + solver_parameters=None, limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + fixed_subcycles (int, optional): the fixed number of sub-steps to + perform. This option cannot be specified with the + `subcycle_by_courant` argument. Defaults to None. + subcycle_by_courant (float, optional): specifying this option will + make the scheme perform adaptive sub-cycling based on the + Courant number. The specified argument is the maximum Courant + for one sub-cycle. Defaults to None, in which case adaptive + sub-cycling is not used. This option cannot be specified with the + `fixed_subcycles` argument. + increment_form (bool, optional): whether to write the RK scheme in + "increment form", solving for increments rather than updated + fields. Defaults to True. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + butcher_matrix = np.array([[1., 0., 0.], [1./4., 1./4., 0.], [1./6., 1./6., 2./3.]]) + + super().__init__(domain, butcher_matrix, field_name=field_name, + fixed_subcycles=fixed_subcycles, + subcycle_by_courant=subcycle_by_courant, + increment_form=increment_form, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + + +class RK4(ExplicitRungeKutta): + u""" + Implements the classic 4-stage Runge-Kutta method. + + The classic 4-stage Runge-Kutta method for solving ∂y/∂t = F(y). It can be + written as: \n + + k0 = F[y^n] \n + k1 = F[y^n + 1/2*dt*k1] \n + k2 = F[y^n + 1/2*dt*k2] \n + k3 = F[y^n + dt*k3] \n + y^(n+1) = y^n + (1/6) * dt * (k0 + 2*k1 + 2*k2 + k3) \n + + where superscripts indicate the time-level. \n + """ + def __init__(self, domain, field_name=None, fixed_subcycles=None, + subcycle_by_courant=None, increment_form=True, + solver_parameters=None, + limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + fixed_subcycles (int, optional): the fixed number of sub-steps to + perform. This option cannot be specified with the + `subcycle_by_courant` argument. Defaults to None. + subcycle_by_courant (float, optional): specifying this option will + make the scheme perform adaptive sub-cycling based on the + Courant number. The specified argument is the maximum Courant + for one sub-cycle. Defaults to None, in which case adaptive + sub-cycling is not used. This option cannot be specified with the + `fixed_subcycles` argument. + increment_form (bool, optional): whether to write the RK scheme in + "increment form", solving for increments rather than updated + fields. Defaults to True. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + butcher_matrix = np.array([[0.5, 0., 0., 0.], [0., 0.5, 0., 0.], [0., 0., 1., 0.], [1./6., 1./3., 1./3., 1./6.]]) + super().__init__(domain, butcher_matrix, field_name=field_name, + fixed_subcycles=fixed_subcycles, + subcycle_by_courant=subcycle_by_courant, + increment_form=increment_form, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + + +class Heun(ExplicitRungeKutta): + u""" + Implements Heun's method. + + The 2-stage Runge-Kutta scheme known as Heun's method,for solving + ∂y/∂t = F(y). It can be written as: \n + + y_1 = F[y^n] \n + y^(n+1) = (1/2)y^n + (1/2)F[y_1] \n + + where superscripts indicate the time-level and subscripts indicate the stage + number. + """ + def __init__(self, domain, field_name=None, fixed_subcycles=None, + subcycle_by_courant=None, increment_form=True, + solver_parameters=None, limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + fixed_subcycles (int, optional): the fixed number of sub-steps to + perform. This option cannot be specified with the + `subcycle_by_courant` argument. Defaults to None. + subcycle_by_courant (float, optional): specifying this option will + make the scheme perform adaptive sub-cycling based on the + Courant number. The specified argument is the maximum Courant + for one sub-cycle. Defaults to None, in which case adaptive + sub-cycling is not used. This option cannot be specified with the + `fixed_subcycles` argument. + increment_form (bool, optional): whether to write the RK scheme in + "increment form", solving for increments rather than updated + fields. Defaults to True. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + butcher_matrix = np.array([[1., 0.], [0.5, 0.5]]) + super().__init__(domain, butcher_matrix, field_name=field_name, + fixed_subcycles=fixed_subcycles, + subcycle_by_courant=subcycle_by_courant, + increment_form=increment_form, + solver_parameters=solver_parameters, + limiter=limiter, options=options) diff --git a/gusto/time_discretisation/imex_runge_kutta.py b/gusto/time_discretisation/imex_runge_kutta.py new file mode 100644 index 000000000..283887e7e --- /dev/null +++ b/gusto/time_discretisation/imex_runge_kutta.py @@ -0,0 +1,417 @@ +"""Implementations of IMEX Runge-Kutta time discretisations.""" + +from firedrake import (Function, Constant, NonlinearVariationalProblem, + NonlinearVariationalSolver) +from firedrake.fml import replace_subject, all_terms, drop +from firedrake.utils import cached_property +from gusto.core.labels import time_derivative, implicit, explicit +from gusto.time_discretisation.time_discretisation import TimeDiscretisation +import numpy as np + + +__all__ = ["IMEXRungeKutta", "IMEX_Euler", "IMEX_ARS3", "IMEX_ARK2", + "IMEX_Trap2", "IMEX_SSP3"] + + +class IMEXRungeKutta(TimeDiscretisation): + """ + A class for implementing general IMEX multistage (Runge-Kutta) + methods based on two Butcher tableaus, to solve \n + + ∂y/∂t = F(y) + S(y) \n + + Where F are implicit fast terms, and S are explicit slow terms. \n + + There are three steps to move from the current solution, y^n, to the new + one, y^{n+1} \n + + For each i = 1, s in an s stage method + we compute the intermediate solutions: \n + y_i = y^n + dt*(a_i1*F(y_1) + a_i2*F(y_2)+ ... + a_ii*F(y_i)) \n + + dt*(d_i1*S(y_1) + d_i2*S(y_2)+ ... + d_{i,i-1}*S(y_{i-1})) + + At the last stage, compute the new solution by: \n + y^{n+1} = y^n + dt*(b_1*F(y_1) + b_2*F(y_2) + .... + b_s*F(y_s)) \n + + dt*(e_1*S(y_1) + e_2*S(y_2) + .... + e_s*S(y_s)) \n + + """ + # -------------------------------------------------------------------------- + # Butcher tableaus for a s-th order + # diagonally implicit scheme (left) and explicit scheme (right): + # c_0 | a_00 0 . 0 f_0 | 0 0 . 0 + # c_1 | a_10 a_11 . 0 f_1 | d_10 0 . 0 + # . | . . . . . | . . . . + # . | . . . . . | . . . . + # c_s | a_s0 a_s1 . a_ss f_s | d_s0 d_s1 . 0 + # ------------------------- ------------------------- + # | b_1 b_2 ... b_s | b_1 b_2 ... b_s + # + # + # The corresponding square 'butcher_imp' and 'butcher_exp' matrices are: + # + # [a_00 0 0 . 0 ] [ 0 0 0 . 0 ] + # [a_10 a_11 0 . 0 ] [d_10 0 0 . 0 ] + # [a_20 a_21 a_22 . 0 ] [d_20 d_21 0 . 0 ] + # [ . . . . . ] [ . . . . . ] + # [ b_0 b_1 . b_s] [ e_0 e_1 . . e_s] + # + # -------------------------------------------------------------------------- + + def __init__(self, domain, butcher_imp, butcher_exp, field_name=None, + solver_parameters=None, limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + butcher_imp (:class:`numpy.ndarray`): A matrix containing the + coefficients of a butcher tableau defining a given implicit + Runge Kutta time discretisation. + butcher_exp (:class:`numpy.ndarray`): A matrix containing the + coefficients of a butcher tableau defining a given explicit + Runge Kutta time discretisation. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + super().__init__(domain, field_name=field_name, + solver_parameters=solver_parameters, + options=options) + self.butcher_imp = butcher_imp + self.butcher_exp = butcher_exp + self.nStages = int(np.shape(self.butcher_imp)[1]) + + def setup(self, equation, apply_bcs=True, *active_labels): + """ + Set up the time discretisation based on the equation. + + Args: + equation (:class:`PrognosticEquation`): the model's equation. + *active_labels (:class:`Label`): labels indicating which terms of + the equation to include. + """ + + super().setup(equation, apply_bcs, *active_labels) + + # Check all terms are labeled implicit, exlicit + for t in self.residual: + if ((not t.has_label(implicit)) and (not t.has_label(explicit)) + and (not t.has_label(time_derivative))): + raise NotImplementedError("Non time-derivative terms must be labeled as implicit or explicit") + + self.xs = [Function(self.fs) for i in range(self.nStages)] + + @cached_property + def lhs(self): + """Set up the discretisation's left hand side (the time derivative).""" + return super(IMEXRungeKutta, self).lhs + + @cached_property + def rhs(self): + """Set up the discretisation's right hand side (the time derivative).""" + return super(IMEXRungeKutta, self).rhs + + def res(self, stage): + """Set up the discretisation's residual for a given stage.""" + # Add time derivative terms y_s - y^n for stage s + mass_form = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=drop) + residual = mass_form.label_map(all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx)) + residual -= mass_form.label_map(all_terms, + map_if_true=replace_subject(self.x1, old_idx=self.idx)) + # Loop through stages up to s-1 and calcualte/sum + # dt*(a_s1*F(y_1) + a_s2*F(y_2)+ ... + a_{s,s-1}*F(y_{s-1})) + # and + # dt*(d_s1*S(y_1) + d_s2*S(y_2)+ ... + d_{s,s-1}*S(y_{s-1})) + for i in range(stage): + r_exp = self.residual.label_map( + lambda t: t.has_label(explicit), + map_if_true=replace_subject(self.xs[i], old_idx=self.idx), + map_if_false=drop) + r_exp = r_exp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: Constant(self.butcher_exp[stage, i])*self.dt*t) + r_imp = self.residual.label_map( + lambda t: t.has_label(implicit), + map_if_true=replace_subject(self.xs[i], old_idx=self.idx), + map_if_false=drop) + r_imp = r_imp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: Constant(self.butcher_imp[stage, i])*self.dt*t) + residual += r_imp + residual += r_exp + # Calculate and add on dt*a_ss*F(y_s) + r_imp = self.residual.label_map( + lambda t: t.has_label(implicit), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop) + r_imp = r_imp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: Constant(self.butcher_imp[stage, stage])*self.dt*t) + residual += r_imp + return residual.form + + @property + def final_res(self): + """Set up the discretisation's final residual.""" + # Add time derivative terms y^{n+1} - y^n + mass_form = self.residual.label_map(lambda t: t.has_label(time_derivative), + map_if_false=drop) + residual = mass_form.label_map(all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx)) + residual -= mass_form.label_map(all_terms, + map_if_true=replace_subject(self.x1, old_idx=self.idx)) + # Loop through stages up to s-1 and calcualte/sum + # dt*(b_1*F(y_1) + b_2*F(y_2) + .... + b_s*F(y_s)) + # and + # dt*(e_1*S(y_1) + e_2*S(y_2) + .... + e_s*S(y_s)) + for i in range(self.nStages): + r_exp = self.residual.label_map( + lambda t: t.has_label(explicit), + map_if_true=replace_subject(self.xs[i], old_idx=self.idx), + map_if_false=drop) + r_exp = r_exp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: Constant(self.butcher_exp[self.nStages, i])*self.dt*t) + r_imp = self.residual.label_map( + lambda t: t.has_label(implicit), + map_if_true=replace_subject(self.xs[i], old_idx=self.idx), + map_if_false=drop) + r_imp = r_imp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: Constant(self.butcher_imp[self.nStages, i])*self.dt*t) + residual += r_imp + residual += r_exp + return residual.form + + @cached_property + def solvers(self): + """Set up a list of solvers for each problem at a stage.""" + solvers = [] + for stage in range(self.nStages): + # setup solver using residual defined in derived class + problem = NonlinearVariationalProblem(self.res(stage), self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage) + solvers.append(NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name)) + return solvers + + @cached_property + def final_solver(self): + """Set up a solver for the final solve to evaluate time level n+1.""" + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.final_res, self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__ + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) + + def apply(self, x_out, x_in): + self.x1.assign(x_in) + solver_list = self.solvers + + for stage in range(self.nStages): + self.solver = solver_list[stage] + self.solver.solve() + self.xs[stage].assign(self.x_out) + + self.final_solver.solve() + x_out.assign(self.x_out) + + +class IMEX_Euler(IMEXRungeKutta): + u""" + Implements IMEX Euler one-stage method. + + The method, for solving \n + ∂y/∂t = F(y) + S(y), can be written as: \n + + y_0 = y^n \n + y_1 = y^n + dt*F[y_1] + dt*S[y_0] \n + y^(n+1) = y^n + dt*F[y_1] + dt*S[y_0] + """ + def __init__(self, domain, field_name=None, solver_parameters=None, + limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + butcher_imp = np.array([[0., 0.], [0., 1.], [0., 1.]]) + butcher_exp = np.array([[0., 0.], [1., 0.], [1., 0.]]) + super().__init__(domain, butcher_imp, butcher_exp, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + + +class IMEX_ARS3(IMEXRungeKutta): + u""" + Implements ARS3(2,3,3) two-stage IMEX Runge–Kutta method + from RK IMEX for HEVI (Weller et al 2013). + Where g = (3 + sqrt(3))/6. + + The method, for solving \n + ∂y/∂t = F(y) + S(y), can be written as: \n + + y_0 = y^n \n + y_1 = y^n + dt*g*F[y_1] + dt*g*S[y_0] \n + y_2 = y^n + dt*((1-2g)*F[y_1]+g*F[y_2]) \n + + dt*((g-1)*S[y_0]+2(g-1)*S[y_1]) \n + y^(n+1) = y^n + dt*(g*F[y_1]+(1-g)*F[y_2]) \n + + dt*(0.5*S[y_1]+0.5*S[y_2]) + """ + def __init__(self, domain, field_name=None, solver_parameters=None, + limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + g = (3. + np.sqrt(3.))/6. + butcher_imp = np.array([[0., 0., 0.], [0., g, 0.], [0., 1-2.*g, g], [0., 0.5, 0.5]]) + butcher_exp = np.array([[0., 0., 0.], [g, 0., 0.], [g-1., 2.*(1.-g), 0.], [0., 0.5, 0.5]]) + + super().__init__(domain, butcher_imp, butcher_exp, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + + +class IMEX_ARK2(IMEXRungeKutta): + u""" + Implements ARK2(2,3,2) two-stage IMEX Runge–Kutta method from + RK IMEX for HEVI (Weller et al 2013). + Where g = 1 - 1/sqrt(2), a = 1/6(3 + 2sqrt(2)), d = 1/2sqrt(2). + + The method, for solving \n + ∂y/∂t = F(y) + S(y), can be written as: \n + + y_0 = y^n \n + y_1 = y^n + dt*(g*F[y_0]+g*F[y_1]) + 2*dt*g*S[y_0] \n + y_2 = y^n + dt*(d*F[y_0]+d*F[y_1]+g*F[y_2]) \n + + dt*((1-a)*S[y_0]+a*S[y_1]) \n + y^(n+1) = y^n + dt*(d*F[y_0]+d*F[y_1]+g*F[y_2]) \n + + dt*(d*S[y_0]+d*S[y_1]+g*S[y_2]) + """ + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + g = 1. - 1./np.sqrt(2.) + d = 1./(2.*np.sqrt(2.)) + a = 1./6.*(3. + 2.*np.sqrt(2.)) + butcher_imp = np.array([[0., 0., 0.], [g, g, 0.], [d, d, g], [d, d, g]]) + butcher_exp = np.array([[0., 0., 0.], [2.*g, 0., 0.], [1.-a, a, 0.], [d, d, g]]) + super().__init__(domain, butcher_imp, butcher_exp, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + + +class IMEX_SSP3(IMEXRungeKutta): + u""" + Implements SSP3(3,3,2) three-stage IMEX Runge–Kutta method from RK IMEX for + HEVI (Weller et al 2013). + + Let g = 1 - 1/sqrt(2). The method, for solving \n + ∂y/∂t = F(y) + S(y), can be written as: \n + + y_1 = y^n + dt*g*F[y_1] \n + y_2 = y^n + dt*((1-2g)*F[y_1]+g*F[y_2]) + dt*S[y_1] \n + y_3 = y^n + dt*((0.5-g)*F[y_1]+g*F[y_3]) + dt*(0.25*S[y_1]+0.25*S[y_2]) \n + y^(n+1) = y^n + dt*(1/6*F[y_1]+1/6*F[y_2]+2/3*F[y_3]) \n + + dt*(1/6*S[y_1]+1/6*S[y_2]+2/3*S[y_3]) + """ + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + g = 1. - (1./np.sqrt(2.)) + butcher_imp = np.array([[g, 0., 0.], [1-2.*g, g, 0.], [0.5-g, 0., g], [(1./6.), (1./6.), (2./3.)]]) + butcher_exp = np.array([[0., 0., 0.], [1., 0., 0.], [0.25, 0.25, 0.], [(1./6.), (1./6.), (2./3.)]]) + super().__init__(domain, butcher_imp, butcher_exp, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + + +class IMEX_Trap2(IMEXRungeKutta): + u""" + Implements Trap2(2+e,3,2) three-stage IMEX Runge–Kutta method from RK IMEX for HEVI (Weller et al 2013). + For e = 1 or 0. + + The method, for solving \n + ∂y/∂t = F(y) + S(y), can be written as: \n + + y_0 = y^n \n + y_1 = y^n + dt*e*F[y_0] + dt*S[y_0] \n + y_2 = y^n + dt*(0.5*F[y_0]+0.5*F[y_2]) + dt*(0.5*S[y_0]+0.5*S[y_1]) \n + y_3 = y^n + dt*(0.5*F[y_0]+0.5*F[y_3]) + dt*(0.5*S[y_0]+0.5*S[y_2]) \n + y^(n+1) = y^n + dt*(0.5*F[y_0]+0.5*F[y_3]) + dt*(0.5*S[y_0] + 0.5*S[y_2]) \n + """ + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + e = 0. + butcher_imp = np.array([[0., 0., 0., 0.], [e, 0., 0., 0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0., 0.5], [0.5, 0., 0., 0.5]]) + butcher_exp = np.array([[0., 0., 0., 0.], [1., 0., 0., 0.], [0.5, 0.5, 0., 0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0.5, 0.]]) + super().__init__(domain, butcher_imp, butcher_exp, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) diff --git a/gusto/time_discretisation/implicit_runge_kutta.py b/gusto/time_discretisation/implicit_runge_kutta.py new file mode 100644 index 000000000..4184dd3e5 --- /dev/null +++ b/gusto/time_discretisation/implicit_runge_kutta.py @@ -0,0 +1,222 @@ +"""Objects to describe implicit multi-stage (Runge-Kutta) discretisations.""" + +import numpy as np + +from firedrake import (Function, split, NonlinearVariationalProblem, + NonlinearVariationalSolver) +from firedrake.fml import replace_subject, all_terms, drop +from firedrake.utils import cached_property + +from gusto.core.labels import time_derivative +from gusto.time_discretisation.time_discretisation import TimeDiscretisation + + +__all__ = ["ImplicitRungeKutta", "ImplicitMidpoint", "QinZhang"] + + +class ImplicitRungeKutta(TimeDiscretisation): + """ + A class for implementing general diagonally implicit multistage (Runge-Kutta) + methods based on its Butcher tableau. + + Unlike the explicit method, all upper diagonal a_ij elements are non-zero + for implicit methods. + + There are three steps to move from the current solution, y^n, to the new + one, y^{n+1} + + For each i = 1, s in an s stage method + we have the intermediate solutions: \n + y_i = y^n + dt*(a_i1*k_1 + a_i2*k_2 + ... + a_ii*k_i) \n + We compute the gradient at the intermediate location, k_i = F(y_i) \n + + At the last stage, compute the new solution by: \n + y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s) + """ + # --------------------------------------------------------------------------- + # Butcher tableau for a s-th order + # diagonally implicit scheme: + # c_0 | a_00 0 . 0 + # c_1 | a_10 a_11 . 0 + # . | . . . . + # . | . . . . + # c_s | a_s0 a_s1 . a_ss + # ------------------------- + # | b_1 b_2 ... b_s + # + # + # The corresponding square 'butcher_matrix' is: + # + # [a_00 0 . 0 ] + # [a_10 a_11 . 0 ] + # [ . . . . ] + # [ b_0 b_1 . b_s] + # --------------------------------------------------------------------------- + + def __init__(self, domain, butcher_matrix, field_name=None, + solver_parameters=None, limiter=None, options=None,): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + butcher_matrix (numpy array): A matrix containing the coefficients + of a butcher tableau defining a given Runge Kutta time + discretisation. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + super().__init__(domain, field_name=field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + self.butcher_matrix = butcher_matrix + self.nStages = int(np.shape(self.butcher_matrix)[1]) + + def setup(self, equation, apply_bcs=True, *active_labels): + """ + Set up the time discretisation based on the equation. + + Args: + equation (:class:`PrognosticEquation`): the model's equation. + *active_labels (:class:`Label`): labels indicating which terms of + the equation to include. + """ + + super().setup(equation, apply_bcs, *active_labels) + + self.k = [Function(self.fs) for i in range(self.nStages)] + + def lhs(self): + return super().lhs + + def rhs(self): + return super().rhs + + def solver(self, stage): + residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=drop, + map_if_false=replace_subject(self.xnph, self.idx), + ) + mass_form = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=drop) + residual += mass_form.label_map(all_terms, + replace_subject(self.x_out, self.idx)) + + problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs) + + solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage) + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + @cached_property + def solvers(self): + solvers = [] + for stage in range(self.nStages): + solvers.append(self.solver(stage)) + return solvers + + def solve_stage(self, x0, stage): + self.x1.assign(x0) + for i in range(stage): + self.x1.assign(self.x1 + self.butcher_matrix[stage, i]*self.dt*self.k[i]) + + if self.limiter is not None: + self.limiter.apply(self.x1) + + if self.idx is None and len(self.fs) > 1: + self.xnph = tuple([self.dt*self.butcher_matrix[stage, stage]*a + b + for a, b in zip(split(self.x_out), split(self.x1))]) + else: + self.xnph = self.x1 + self.butcher_matrix[stage, stage]*self.dt*self.x_out + solver = self.solvers[stage] + solver.solve() + + self.k[stage].assign(self.x_out) + + def apply(self, x_out, x_in): + + for i in range(self.nStages): + self.solve_stage(x_in, i) + + x_out.assign(x_in) + for i in range(self.nStages): + x_out.assign(x_out + self.butcher_matrix[self.nStages, i]*self.dt*self.k[i]) + + if self.limiter is not None: + self.limiter.apply(x_out) + + +class ImplicitMidpoint(ImplicitRungeKutta): + u""" + Implements the Implicit Midpoint method as a 1-stage Runge Kutta method. + + The method, for solving + ∂y/∂t = F(y), can be written as: \n + + k0 = F[y^n + 0.5*dt*k0] \n + y^(n+1) = y^n + dt*k0 \n + """ + def __init__(self, domain, field_name=None, solver_parameters=None, + limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + butcher_matrix = np.array([[0.5], [1.]]) + super().__init__(domain, butcher_matrix, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + + +class QinZhang(ImplicitRungeKutta): + u""" + Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method. + + The method, for solving + ∂y/∂t = F(y), can be written as: \n + + k0 = F[y^n + 0.25*dt*k0] \n + k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] \n + y^(n+1) = y^n + 0.5*dt*(k0 + k1) \n + """ + def __init__(self, domain, field_name=None, solver_parameters=None, + limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + butcher_matrix = np.array([[0.25, 0], [0.5, 0.25], [0.5, 0.5]]) + super().__init__(domain, butcher_matrix, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) diff --git a/gusto/time_discretisation/multi_level_schemes.py b/gusto/time_discretisation/multi_level_schemes.py new file mode 100644 index 000000000..47e3b9652 --- /dev/null +++ b/gusto/time_discretisation/multi_level_schemes.py @@ -0,0 +1,511 @@ +"""Implements time discretisations with multiple time levels.""" + +from abc import abstractproperty + +from firedrake import (Function, NonlinearVariationalProblem, + NonlinearVariationalSolver) +from firedrake.fml import replace_subject, all_terms, drop +from gusto.core.configuration import EmbeddedDGOptions, RecoveryOptions +from gusto.core.labels import time_derivative +from gusto.time_discretisation.time_discretisation import TimeDiscretisation + + +__all__ = ["BDF2", "Leapfrog", "AdamsMoulton", "AdamsBashforth"] + + +class MultilevelTimeDiscretisation(TimeDiscretisation): + """Base class for multi-level timesteppers""" + + def __init__(self, domain, field_name=None, solver_parameters=None, + limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): + raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") + super().__init__(domain=domain, field_name=field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + self.initial_timesteps = 0 + + @abstractproperty + def nlevels(self): + pass + + def setup(self, equation, apply_bcs=True, *active_labels): + super().setup(equation=equation, apply_bcs=apply_bcs, *active_labels) + for n in range(self.nlevels, 1, -1): + setattr(self, "xnm%i" % (n-1), Function(self.fs)) + + +class BDF2(MultilevelTimeDiscretisation): + """ + Implements the implicit multistep BDF2 timestepping method. + + The BDF2 timestepping method for operator F is written as: \n + y^(n+1) = (4/3)*y^n - (1/3)*y^(n-1) + (2/3)*dt*F[y^(n+1)] \n + """ + + @property + def nlevels(self): + return 2 + + @property + def lhs0(self): + """Set up the discretisation's left hand side (the time derivative).""" + l = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx)) + l = l.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: self.dt*t) + + return l.form + + @property + def rhs0(self): + """Set up the time discretisation's right hand side for inital BDF step.""" + r = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x1, old_idx=self.idx), + map_if_false=drop) + + return r.form + + @property + def lhs(self): + """Set up the discretisation's left hand side (the time derivative).""" + l = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx)) + l = l.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: (2/3)*self.dt*t) + + return l.form + + @property + def rhs(self): + """Set up the time discretisation's right hand side for BDF2 steps.""" + xn = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x1, old_idx=self.idx), + map_if_false=drop) + xnm1 = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.xnm1, old_idx=self.idx), + map_if_false=drop) + + r = (4/3.) * xn - (1/3.) * xnm1 + + return r.form + + @property + def solver0(self): + """Set up the problem and the solver for initial BDF step.""" + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.lhs0-self.rhs0, self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__+"0" + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + @property + def solver(self): + """Set up the problem and the solver for BDF2 steps.""" + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__ + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + def apply(self, x_out, *x_in): + """ + Apply the time discretisation to advance one whole time step. + + Args: + x_out (:class:`Function`): the output field to be computed. + x_in (:class:`Function`): the input field(s). + """ + if self.initial_timesteps < self.nlevels-1: + self.initial_timesteps += 1 + solver = self.solver0 + else: + solver = self.solver + + self.xnm1.assign(x_in[0]) + self.x1.assign(x_in[1]) + solver.solve() + x_out.assign(self.x_out) + + +class Leapfrog(MultilevelTimeDiscretisation): + """ + Implements the multistep Leapfrog timestepping method. + + The Leapfrog timestepping method for operator F is written as: \n + y^(n+1) = y^(n-1) + 2*dt*F[y^n] + """ + @property + def nlevels(self): + return 2 + + @property + def rhs0(self): + """Set up the discretisation's right hand side for initial forward euler step.""" + r = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x1, old_idx=self.idx)) + r = r.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt*t) + + return r.form + + @property + def lhs(self): + """Set up the discretisation's left hand side (the time derivative).""" + return super(Leapfrog, self).lhs + + @property + def rhs(self): + """Set up the discretisation's right hand side for leapfrog steps.""" + r = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=replace_subject(self.x1, old_idx=self.idx)) + r = r.label_map(lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.xnm1, old_idx=self.idx), + map_if_false=lambda t: -2.0*self.dt*t) + + return r.form + + @property + def solver0(self): + """Set up the problem and the solver for initial forward euler step.""" + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.lhs-self.rhs0, self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__+"0" + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + @property + def solver(self): + """Set up the problem and the solver for leapfrog steps.""" + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__ + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + def apply(self, x_out, *x_in): + """ + Apply the time discretisation to advance one whole time step. + + Args: + x_out (:class:`Function`): the output field to be computed. + x_in (:class:`Function`): the input field(s). + """ + if self.initial_timesteps < self.nlevels-1: + self.initial_timesteps += 1 + solver = self.solver0 + else: + solver = self.solver + + self.xnm1.assign(x_in[0]) + self.x1.assign(x_in[1]) + solver.solve() + x_out.assign(self.x_out) + + +class AdamsBashforth(MultilevelTimeDiscretisation): + """ + Implements the explicit multistep Adams-Bashforth timestepping + method of general order up to 5. + + The general AB timestepping method for operator F is written as: \n + y^(n+1) = y^n + dt*(b_0*F[y^(n)] + b_1*F[y^(n-1)] + b_2*F[y^(n-2)] + b_3*F[y^(n-3)] + b_4*F[y^(n-4)]) + """ + def __init__(self, domain, order, field_name=None, + solver_parameters=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + order (float, optional): order of scheme + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + + Raises: + ValueError: if order is not provided, or is in incorrect range. + """ + + if (order > 5 or order < 1): + raise ValueError("Adams-Bashforth of order greater than 5 not implemented") + if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): + raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") + + super().__init__(domain, field_name, + solver_parameters=solver_parameters, + options=options) + + self.order = order + + def setup(self, equation, apply_bcs=True, *active_labels): + super().setup(equation=equation, apply_bcs=apply_bcs, + *active_labels) + + self.x = [Function(self.fs) for i in range(self.nlevels)] + + if (self.order == 1): + self.b = [1.0] + elif (self.order == 2): + self.b = [-(1.0/2.0), (3.0/2.0)] + elif (self.order == 3): + self.b = [(5.0)/(12.0), -(16.0)/(12.0), (23.0)/(12.0)] + elif (self.order == 4): + self.b = [-(9.0)/(24.0), (37.0)/(24.0), -(59.0)/(24.0), (55.0)/(24.0)] + elif (self.order == 5): + self.b = [(251.0)/(720.0), -(1274.0)/(720.0), (2616.0)/(720.0), + -(2774.0)/(720.0), (2901.0)/(720.0)] + + @property + def nlevels(self): + return self.order + + @property + def rhs0(self): + """Set up the discretisation's right hand side for initial forward euler step.""" + r = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) + r = r.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt*t) + + return r.form + + @property + def lhs(self): + """Set up the discretisation's left hand side (the time derivative).""" + return super(AdamsBashforth, self).lhs + + @property + def rhs(self): + """Set up the discretisation's right hand side for Adams Bashforth steps.""" + r = self.residual.label_map(all_terms, + map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) + r = r.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.b[-1]*self.dt*t) + for n in range(self.nlevels-1): + rtemp = self.residual.label_map(lambda t: t.has_label(time_derivative), + map_if_true=drop, + map_if_false=replace_subject(self.x[n], old_idx=self.idx)) + rtemp = rtemp.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt*self.b[n]*t) + r += rtemp + return r.form + + @property + def solver0(self): + """Set up the problem and the solverfor initial forward euler step.""" + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.lhs-self.rhs0, self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__+"0" + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + @property + def solver(self): + """Set up the problem and the solver for Adams Bashforth steps.""" + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__ + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + def apply(self, x_out, *x_in): + """ + Apply the time discretisation to advance one whole time step. + + Args: + x_out (:class:`Function`): the output field to be computed. + x_in (:class:`Function`): the input field(s). + """ + if self.initial_timesteps < self.nlevels-1: + self.initial_timesteps += 1 + solver = self.solver0 + else: + solver = self.solver + + for n in range(self.nlevels): + self.x[n].assign(x_in[n]) + solver.solve() + x_out.assign(self.x_out) + + +class AdamsMoulton(MultilevelTimeDiscretisation): + """ + Implements the implicit multistep Adams-Moulton + timestepping method of general order up to 5 + + The general AM timestepping method for operator F is written as \n + y^(n+1) = y^n + dt*(b_0*F[y^(n+1)] + b_1*F[y^(n)] + b_2*F[y^(n-1)] + b_3*F[y^(n-2)]) + """ + def __init__(self, domain, order, field_name=None, + solver_parameters=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + order (float, optional): order of scheme + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + + Raises: + ValueError: if order is not provided, or is in incorrect range. + """ + if (order > 4 or order < 1): + raise ValueError("Adams-Moulton of order greater than 5 not implemented") + if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): + raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") + if not solver_parameters: + solver_parameters = {'ksp_type': 'gmres', + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu'} + + super().__init__(domain, field_name, + solver_parameters=solver_parameters, + options=options) + + self.order = order + + def setup(self, equation, apply_bcs=True, *active_labels): + super().setup(equation=equation, apply_bcs=apply_bcs, *active_labels) + + self.x = [Function(self.fs) for i in range(self.nlevels)] + + if (self.order == 1): + self.bl = (1.0/2.0) + self.br = [(1.0/2.0)] + elif (self.order == 2): + self.bl = (5.0/12.0) + self.br = [-(1.0/12.0), (8.0/12.0)] + elif (self.order == 3): + self.bl = (9.0/24.0) + self.br = [(1.0/24.0), -(5.0/24.0), (19.0/24.0)] + elif (self.order == 4): + self.bl = (251.0/720.0) + self.br = [-(19.0/720.0), (106.0/720.0), -(254.0/720.0), (646.0/720.0)] + + @property + def nlevels(self): + return self.order + + @property + def rhs0(self): + """ + Set up the discretisation's right hand side for initial trapezoidal + step. + """ + r = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) + r = r.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -0.5*self.dt*t) + + return r.form + + @property + def lhs0(self): + """ + Set up the time discretisation's right hand side for initial + trapezoidal step. + """ + l = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx)) + l = l.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: 0.5*self.dt*t) + return l.form + + @property + def lhs(self): + """Set up the time discretisation's right hand side for Adams Moulton steps.""" + l = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx)) + l = l.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: self.bl*self.dt*t) + return l.form + + @property + def rhs(self): + """Set up the discretisation's right hand side for Adams Moulton steps.""" + r = self.residual.label_map(all_terms, + map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) + r = r.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.br[-1]*self.dt*t) + for n in range(self.nlevels-1): + rtemp = self.residual.label_map(lambda t: t.has_label(time_derivative), + map_if_true=drop, + map_if_false=replace_subject(self.x[n], old_idx=self.idx)) + rtemp = rtemp.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt*self.br[n]*t) + r += rtemp + return r.form + + @property + def solver0(self): + """Set up the problem and the solver for initial trapezoidal step.""" + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.lhs0-self.rhs0, self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__+"0" + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + @property + def solver(self): + """Set up the problem and the solver for Adams Moulton steps.""" + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__ + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + def apply(self, x_out, *x_in): + """ + Apply the time discretisation to advance one whole time step. + + Args: + x_out (:class:`Function`): the output field to be computed. + x_in (:class:`Function`): the input field(s). + """ + if self.initial_timesteps < self.nlevels-1: + self.initial_timesteps += 1 + solver = self.solver0 + else: + solver = self.solver + + for n in range(self.nlevels): + self.x[n].assign(x_in[n]) + solver.solve() + x_out.assign(self.x_out) diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py new file mode 100644 index 000000000..d1e27c425 --- /dev/null +++ b/gusto/time_discretisation/time_discretisation.py @@ -0,0 +1,723 @@ +u""" +Objects for discretising time derivatives. + +Time discretisation objects discretise ∂y/∂t = F(y), for variable y, time t and +operator F. +""" + +from abc import ABCMeta, abstractmethod, abstractproperty +import math + +from firedrake import (Function, TestFunction, TestFunctions, DirichletBC, + Constant, NonlinearVariationalProblem, + NonlinearVariationalSolver) +from firedrake.fml import (replace_subject, replace_test_function, Term, + all_terms, drop) +from firedrake.formmanipulation import split_form +from firedrake.utils import cached_property + +from gusto.core.configuration import EmbeddedDGOptions, RecoveryOptions +from gusto.core.labels import time_derivative, prognostic, physics_label +from gusto.core.logging import logger, DEBUG, logging_ksp_monitor_true_residual +from gusto.time_discretisation.wrappers import * + + +__all__ = ["TimeDiscretisation", "ExplicitTimeDiscretisation", "BackwardEuler", + "ThetaMethod", "TrapeziumRule", "TR_BDF2"] + + +def wrapper_apply(original_apply): + """Decorator to add steps for using a wrapper around the apply method.""" + def get_apply(self, x_out, x_in): + + if self.wrapper is not None: + + def new_apply(self, x_out, x_in): + + self.wrapper.pre_apply(x_in) + original_apply(self, self.wrapper.x_out, self.wrapper.x_in) + self.wrapper.post_apply(x_out) + + return new_apply(self, x_out, x_in) + + else: + + return original_apply(self, x_out, x_in) + + return get_apply + + +class TimeDiscretisation(object, metaclass=ABCMeta): + """Base class for time discretisation schemes.""" + + def __init__(self, domain, field_name=None, solver_parameters=None, + limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + self.domain = domain + self.field_name = field_name + self.equation = None + + self.dt = Constant(0.0) + self.dt.assign(domain.dt) + self.original_dt = Constant(0.0) + self.original_dt.assign(self.dt) + self.options = options + self.limiter = limiter + self.courant_max = None + + if options is not None: + self.wrapper_name = options.name + if self.wrapper_name == "mixed_options": + self.wrapper = MixedFSWrapper() + + for field, suboption in options.suboptions.items(): + if suboption.name == 'embedded_dg': + self.wrapper.subwrappers.update({field: EmbeddedDGWrapper(self, suboption)}) + elif suboption.name == "recovered": + self.wrapper.subwrappers.update({field: RecoveryWrapper(self, suboption)}) + elif suboption.name == "supg": + raise RuntimeError( + 'Time discretisation: suboption SUPG is currently not implemented within MixedOptions') + else: + raise RuntimeError( + f'Time discretisation: suboption wrapper {wrapper_name} not implemented') + elif self.wrapper_name == "embedded_dg": + self.wrapper = EmbeddedDGWrapper(self, options) + elif self.wrapper_name == "recovered": + self.wrapper = RecoveryWrapper(self, options) + elif self.wrapper_name == "supg": + self.wrapper = SUPGWrapper(self, options) + else: + raise RuntimeError( + f'Time discretisation: wrapper {self.wrapper_name} not implemented') + else: + self.wrapper = None + self.wrapper_name = None + + # get default solver options if none passed in + if solver_parameters is None: + self.solver_parameters = {'ksp_type': 'cg', + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu'} + else: + self.solver_parameters = solver_parameters + + def setup(self, equation, apply_bcs=True, *active_labels): + """ + Set up the time discretisation based on the equation. + + Args: + equation (:class:`PrognosticEquation`): the model's equation. + apply_bcs (bool, optional): whether to apply the equation's boundary + conditions. Defaults to True. + *active_labels (:class:`Label`): labels indicating which terms of + the equation to include. + """ + self.equation = equation + self.residual = equation.residual + + if self.field_name is not None and hasattr(equation, "field_names"): + self.idx = equation.field_names.index(self.field_name) + self.fs = equation.spaces[self.idx] + self.residual = self.residual.label_map( + lambda t: t.get(prognostic) == self.field_name, + lambda t: Term( + split_form(t.form)[self.idx].form, + t.labels), + drop) + + else: + self.field_name = equation.field_name + self.fs = equation.function_space + self.idx = None + + bcs = equation.bcs[self.field_name] + + if len(active_labels) > 0: + self.residual = self.residual.label_map( + lambda t: any(t.has_label(time_derivative, *active_labels)), + map_if_false=drop) + + self.evaluate_source = [] + self.physics_names = [] + for t in self.residual: + if t.has_label(physics_label): + physics_name = t.get(physics_label) + if t.labels[physics_name] not in self.physics_names: + self.evaluate_source.append(t.labels[physics_name]) + self.physics_names.append(t.labels[physics_name]) + + # -------------------------------------------------------------------- # + # Set up Wrappers + # -------------------------------------------------------------------- # + + if self.wrapper is not None: + if self.wrapper_name == "mixed_options": + + self.wrapper.wrapper_spaces = equation.spaces + self.wrapper.field_names = equation.field_names + + for field, subwrapper in self.wrapper.subwrappers.items(): + + if field not in equation.field_names: + raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set") + + field_idx = equation.field_names.index(field) + subwrapper.setup(equation.spaces[field_idx]) + + # Update the function space to that needed by the wrapper + self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space + + self.wrapper.setup() + self.fs = self.wrapper.function_space + new_test_mixed = TestFunctions(self.fs) + + # Replace the original test function with one from the new + # function space defined by the subwrappers + self.residual = self.residual.label_map( + all_terms, + map_if_true=replace_test_function(new_test_mixed)) + + else: + if self.wrapper_name == "supg": + self.wrapper.setup() + else: + self.wrapper.setup(self.fs) + self.fs = self.wrapper.function_space + if self.solver_parameters is None: + self.solver_parameters = self.wrapper.solver_parameters + new_test = TestFunction(self.wrapper.test_space) + # SUPG has a special wrapper + if self.wrapper_name == "supg": + new_test = self.wrapper.test + + # Replace the original test function with the one from the wrapper + self.residual = self.residual.label_map( + all_terms, + map_if_true=replace_test_function(new_test)) + + self.residual = self.wrapper.label_terms(self.residual) + + # -------------------------------------------------------------------- # + # Make boundary conditions + # -------------------------------------------------------------------- # + + if not apply_bcs: + self.bcs = None + elif self.wrapper is not None: + # Transfer boundary conditions onto test function space + self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain) + for bc in bcs] + else: + self.bcs = bcs + + # -------------------------------------------------------------------- # + # Make the required functions + # -------------------------------------------------------------------- # + + self.x_out = Function(self.fs) + self.x1 = Function(self.fs) + + @property + def nlevels(self): + return 1 + + @abstractproperty + def lhs(self): + """Set up the discretisation's left hand side (the time derivative).""" + l = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop) + + return l.form + + @abstractproperty + def rhs(self): + """Set up the time discretisation's right hand side.""" + r = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x1, old_idx=self.idx)) + + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt*t) + + return r.form + + @cached_property + def solver(self): + """Set up the problem and the solver.""" + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__ + solver = NonlinearVariationalSolver( + problem, + solver_parameters=self.solver_parameters, + options_prefix=solver_name + ) + if logger.isEnabledFor(DEBUG): + solver.snes.ksp.setMonitor(logging_ksp_monitor_true_residual) + return solver + + @abstractmethod + def apply(self, x_out, x_in): + """ + Apply the time discretisation to advance one whole time step. + + Args: + x_out (:class:`Function`): the output field to be computed. + x_in (:class:`Function`): the input field. + """ + pass + + +class ExplicitTimeDiscretisation(TimeDiscretisation): + """Base class for explicit time discretisations.""" + + def __init__(self, domain, field_name=None, fixed_subcycles=None, + subcycle_by_courant=None, solver_parameters=None, limiter=None, + options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + fixed_subcycles (int, optional): the fixed number of sub-steps to + perform. This option cannot be specified with the + `subcycle_by_courant` argument. Defaults to None. + subcycle_by_courant (float, optional): specifying this option will + make the scheme perform adaptive sub-cycling based on the + Courant number. The specified argument is the maximum Courant + for one sub-cycle. Defaults to None, in which case adaptive + sub-cycling is not used. This option cannot be specified with the + `fixed_subcycles` argument. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + super().__init__(domain, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + + if fixed_subcycles is not None and subcycle_by_courant is not None: + raise ValueError('Cannot specify both subcycle and subcycle_by ' + + 'arguments to a time discretisation') + self.fixed_subcycles = fixed_subcycles + self.subcycle_by_courant = subcycle_by_courant + + def setup(self, equation, apply_bcs=True, *active_labels): + """ + Set up the time discretisation based on the equation. + + Args: + equation (:class:`PrognosticEquation`): the model's equation. + apply_bcs (bool, optional): whether boundary conditions are to be + applied. Defaults to True. + *active_labels (:class:`Label`): labels indicating which terms of + the equation to include. + """ + super().setup(equation, apply_bcs, *active_labels) + + # if user has specified a number of fixed subcycles, then save this + # and rescale dt accordingly; else perform just one cycle using dt + if self.fixed_subcycles is not None: + self.dt.assign(self.dt/self.fixed_subcycles) + self.ncycles = self.fixed_subcycles + else: + self.dt = self.dt + self.ncycles = 1 + self.x0 = Function(self.fs) + self.x1 = Function(self.fs) + + @cached_property + def lhs(self): + """Set up the discretisation's left hand side (the time derivative).""" + l = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x_out, self.idx), + map_if_false=drop) + + return l.form + + @cached_property + def solver(self): + """Set up the problem and the solver.""" + # setup linear solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.lhs - self.rhs, self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__ + # If snes_type not specified by user, set this to ksp only to avoid outer Newton iteration + self.solver_parameters.setdefault('snes_type', 'ksponly') + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + @abstractmethod + def apply_cycle(self, x_out, x_in): + """ + Apply the time discretisation through a single sub-step. + + Args: + x_out (:class:`Function`): the output field to be computed. + x_in (:class:`Function`): the input field. + """ + pass + + @wrapper_apply + def apply(self, x_out, x_in): + """ + Apply the time discretisation to advance one whole time step. + + Args: + x_out (:class:`Function`): the output field to be computed. + x_in (:class:`Function`): the input field. + """ + # If doing adaptive subcycles, update dt and ncycles here + if self.subcycle_by_courant is not None: + self.ncycles = math.ceil(float(self.courant_max)/self.subcycle_by_courant) + self.dt.assign(self.original_dt/self.ncycles) + + self.x0.assign(x_in) + for i in range(self.ncycles): + self.apply_cycle(self.x1, self.x0) + self.x0.assign(self.x1) + x_out.assign(self.x1) + + +class BackwardEuler(TimeDiscretisation): + """ + Implements the backward Euler timestepping scheme. + + The backward Euler method for operator F is the most simple implicit scheme: \n + y^(n+1) = y^n + dt*F[y^(n+1)]. \n + """ + def __init__(self, domain, field_name=None, solver_parameters=None, + limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + fixed_subcycles (int, optional): the number of sub-steps to perform. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods. Defaults to None. + """ + if not solver_parameters: + # default solver parameters + solver_parameters = {'ksp_type': 'gmres', + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu'} + super().__init__(domain=domain, field_name=field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + + @property + def lhs(self): + """Set up the discretisation's left hand side (the time derivative).""" + l = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx)) + l = l.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: self.dt*t) + + return l.form + + @property + def rhs(self): + """Set up the time discretisation's right hand side.""" + r = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x1, old_idx=self.idx), + map_if_false=drop) + + return r.form + + def apply(self, x_out, x_in): + """ + Apply the time discretisation to advance one whole time step. + + Args: + x_out (:class:`Function`): the output field to be computed. + x_in (:class:`Function`): the input field. + """ + for evaluate in self.evaluate_source: + evaluate(x_in, self.dt) + + if len(self.evaluate_source) > 0: + # If we have physics, use x_in as first guess + self.x_out.assign(x_in) + + self.x1.assign(x_in) + self.solver.solve() + x_out.assign(self.x_out) + + +class ThetaMethod(TimeDiscretisation): + """ + Implements the theta implicit-explicit timestepping method, which can + be thought as a generalised trapezium rule. + + The theta implicit-explicit timestepping method for operator F is written + as: \n + y^(n+1) = y^n + dt*(1-theta)*F[y^n] + dt*theta*F[y^(n+1)] \n + for off-centring parameter theta. + """ + + def __init__(self, domain, theta, field_name=None, + solver_parameters=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + theta (float): the off-centring parameter. theta = 1 + corresponds to a backward Euler method. Defaults to None. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + + Raises: + ValueError: if theta is not provided. + """ + if (theta < 0 or theta > 1): + raise ValueError("please provide a value for theta between 0 and 1") + if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): + raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") + if not solver_parameters: + # theta method leads to asymmetric matrix, per lhs function below, + # so don't use CG + solver_parameters = {'ksp_type': 'gmres', + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu'} + + super().__init__(domain, field_name, + solver_parameters=solver_parameters, + options=options) + + self.theta = theta + + @cached_property + def lhs(self): + """Set up the discretisation's left hand side (the time derivative).""" + l = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx)) + l = l.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: self.theta*self.dt*t) + + return l.form + + @cached_property + def rhs(self): + """Set up the time discretisation's right hand side.""" + r = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x1, old_idx=self.idx)) + r = r.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -(1-self.theta)*self.dt*t) + + return r.form + + def apply(self, x_out, x_in): + """ + Apply the time discretisation to advance one whole time step. + + Args: + x_out (:class:`Function`): the output field to be computed. + x_in (:class:`Function`): the input field. + """ + self.x1.assign(x_in) + self.solver.solve() + x_out.assign(self.x_out) + + +class TrapeziumRule(ThetaMethod): + """ + Implements the trapezium rule timestepping method, also commonly known as + Crank Nicholson. + + The trapezium rule timestepping method for operator F is written as: \n + y^(n+1) = y^n + dt/2*F[y^n] + dt/2*F[y^(n+1)]. \n + It is equivalent to the "theta" method with theta = 1/2. \n + """ + + def __init__(self, domain, field_name=None, solver_parameters=None, + options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + super().__init__(domain, 0.5, field_name, + solver_parameters=solver_parameters, + options=options) + + +# TODO: this should be implemented as an ImplicitRK +class TR_BDF2(TimeDiscretisation): + """ + Implements the two stage implicit TR-BDF2 time stepping method, with a + trapezoidal stage (TR) followed by a second order backwards difference stage + (BDF2). + + The TR-BDF2 time stepping method for operator F is written as: \n + y^(n+g) = y^n + dt*g/2*F[y^n] + dt*g/2*F[y^(n+g)] (TR stage) \n + y^(n+1) = 1/(g(2-g))*y^(n+g) - (1-g)**2/(g(2-g))*y^(n) + (1-g)/(2-g)*dt*F[y^(n+1)] (BDF2 stage) \n + for an off-centring parameter g (gamma). + """ + def __init__(self, domain, gamma, field_name=None, + solver_parameters=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + gamma (float): the off-centring parameter + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + if (gamma < 0. or gamma > 1.): + raise ValueError("please provide a value for gamma between 0 and 1") + if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): + raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") + if not solver_parameters: + # theta method leads to asymmetric matrix, per lhs function below, + # so don't use CG + solver_parameters = {'ksp_type': 'gmres', + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu'} + + super().__init__(domain, field_name, + solver_parameters=solver_parameters, + options=options) + + self.gamma = gamma + + def setup(self, equation, apply_bcs=True, *active_labels): + super().setup(equation, apply_bcs, *active_labels) + self.xnpg = Function(self.fs) + self.xn = Function(self.fs) + + @cached_property + def lhs(self): + """Set up the discretisation's left hand side (the time derivative) for the TR stage.""" + l = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.xnpg, old_idx=self.idx)) + l = l.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: 0.5*self.gamma*self.dt*t) + + return l.form + + @cached_property + def rhs(self): + """Set up the time discretisation's right hand side for the TR stage.""" + r = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.xn, old_idx=self.idx)) + r = r.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -0.5*self.gamma*self.dt*t) + + return r.form + + @cached_property + def lhs_bdf2(self): + """Set up the discretisation's left hand side (the time derivative) for + the BDF2 stage.""" + l = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx)) + l = l.label_map(lambda t: t.has_label(time_derivative), + map_if_false=lambda t: ((1.0-self.gamma)/(2.0-self.gamma))*self.dt*t) + + return l.form + + @cached_property + def rhs_bdf2(self): + """Set up the time discretisation's right hand side for the BDF2 stage.""" + xn = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.xn, old_idx=self.idx), + map_if_false=drop) + xnpg = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.xnpg, old_idx=self.idx), + map_if_false=drop) + + r = (1.0/(self.gamma*(2.0-self.gamma)))*xnpg - ((1.0-self.gamma)**2/(self.gamma*(2.0-self.gamma)))*xn + + return r.form + + @cached_property + def solver_tr(self): + """Set up the problem and the solver.""" + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.xnpg, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__+"_tr" + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + @cached_property + def solver_bdf2(self): + """Set up the problem and the solver.""" + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.lhs_bdf2-self.rhs_bdf2, self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__+"_bdf2" + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + def apply(self, x_out, x_in): + """ + Apply the time discretisation to advance one whole time step. + + Args: + x_out (:class:`Function`): the output field to be computed. + x_in (:class:`Function`): the input field(s). + """ + self.xn.assign(x_in) + self.solver_tr.solve() + self.solver_bdf2.solve() + x_out.assign(self.x_out) diff --git a/gusto/wrappers.py b/gusto/time_discretisation/wrappers.py similarity index 99% rename from gusto/wrappers.py rename to gusto/time_discretisation/wrappers.py index 3ca9c4347..8b1a5c4f1 100644 --- a/gusto/wrappers.py +++ b/gusto/time_discretisation/wrappers.py @@ -10,9 +10,9 @@ VectorElement, Constant, as_ufl, dot, grad, TestFunction, MixedFunctionSpace ) from firedrake.fml import Term -from gusto.configuration import EmbeddedDGOptions, RecoveryOptions, SUPGOptions +from gusto.core.configuration import EmbeddedDGOptions, RecoveryOptions, SUPGOptions from gusto.recovery import Recoverer, ReversibleRecoverer -from gusto.labels import transporting_velocity +from gusto.core.labels import transporting_velocity import ufl __all__ = ["EmbeddedDGWrapper", "RecoveryWrapper", "SUPGWrapper", "MixedFSWrapper"] diff --git a/gusto/timeloop.py b/gusto/timeloop.py deleted file mode 100644 index 717b9002e..000000000 --- a/gusto/timeloop.py +++ /dev/null @@ -1,869 +0,0 @@ -"""Classes for controlling the timestepping loop.""" - -from abc import ABCMeta, abstractmethod, abstractproperty -from firedrake import Function, Projector, split, Constant -from firedrake.fml import drop, Label, Term -from pyop2.profiling import timed_stage -from gusto.equations import PrognosticEquationSet -from gusto.fields import TimeLevelFields, StateFields -from gusto.forcing import Forcing -from gusto.labels import ( - transport, diffusion, time_derivative, linearisation, prognostic, - physics_label, transporting_velocity -) -from gusto.linear_solvers import LinearTimesteppingSolver -from gusto.logging import logger -from gusto.time_discretisation import ExplicitTimeDiscretisation -from gusto.transport_methods import TransportMethod -import ufl - -__all__ = ["Timestepper", "SplitPhysicsTimestepper", "SplitPrescribedTransport", - "SemiImplicitQuasiNewton", "PrescribedTransport"] - - -class BaseTimestepper(object, metaclass=ABCMeta): - """Base class for timesteppers.""" - - def __init__(self, equation, io): - """ - Args: - equation (:class:`PrognosticEquation`): the prognostic equation. - io (:class:`IO`): the model's object for controlling input/output. - """ - - self.equation = equation - self.io = io - self.dt = self.equation.domain.dt - self.t = self.equation.domain.t - self.reference_profiles_initialised = False - - self.setup_fields() - self.setup_scheme() - - self.io.log_parameters(equation) - - @abstractproperty - def transporting_velocity(self): - return NotImplementedError - - @abstractmethod - def setup_fields(self): - """Set up required fields. Must be implemented in child classes""" - pass - - @abstractmethod - def setup_scheme(self): - """Set up required scheme(s). Must be implemented in child classes""" - pass - - @abstractmethod - def timestep(self): - """Defines the timestep. Must be implemented in child classes""" - return NotImplementedError - - def set_initial_timesteps(self, num_steps): - """Sets the number of initial time steps for a multi-level scheme.""" - can_set = (hasattr(self, 'scheme') - and hasattr(self.scheme, 'initial_timesteps') - and num_steps is not None) - if can_set: - self.scheme.initial_timesteps = num_steps - - def get_initial_timesteps(self): - """Gets the number of initial time steps from a multi-level scheme.""" - can_get = (hasattr(self, 'scheme') - and hasattr(self.scheme, 'initial_timesteps')) - # Return None if this is not applicable - return self.scheme.initial_timesteps if can_get else None - - def setup_equation(self, equation): - """ - Sets up the spatial methods for an equation, by the setting the - forms used for transport/diffusion in the equation. - - Args: - equation (:class:`PrognosticEquation`): the equation that the - transport method is to be applied to. - """ - - # For now, we only have methods for transport and diffusion - for term_label in [transport, diffusion]: - # ---------------------------------------------------------------- # - # Check that appropriate methods have been provided - # ---------------------------------------------------------------- # - # Extract all terms corresponding to this type of term - residual = equation.residual.label_map( - lambda t: t.has_label(term_label), map_if_false=drop - ) - variables = [t.get(prognostic) for t in residual.terms] - methods = list(filter(lambda t: t.term_label == term_label, - self.spatial_methods)) - method_variables = [method.variable for method in methods] - for variable in variables: - if variable not in method_variables: - message = f'Variable {variable} has a {term_label.label} ' \ - + 'term but no method for this has been specified. ' \ - + 'Using default form for this term' - logger.warning(message) - - # -------------------------------------------------------------------- # - # Check that appropriate methods have been provided - # -------------------------------------------------------------------- # - # Replace forms in equation - if self.spatial_methods is not None: - for method in self.spatial_methods: - method.replace_form(equation) - - def setup_transporting_velocity(self, scheme): - """ - Set up the time discretisation by replacing the transporting velocity - used by the appropriate one for this time loop. - - Args: - scheme (:class:`TimeDiscretisation`): the time discretisation whose - transport term should be replaced with the transport term of - this discretisation. - """ - if self.transporting_velocity == "prognostic" and "u" in self.fields._field_names: - # Use the prognostic wind variable as the transporting velocity - u_idx = self.equation.field_names.index('u') - uadv = split(self.equation.X)[u_idx] - else: - uadv = self.transporting_velocity - - scheme.residual = scheme.residual.label_map( - lambda t: t.has_label(transporting_velocity), - map_if_true=lambda t: - Term(ufl.replace(t.form, {t.get(transporting_velocity): uadv}), t.labels) - ) - - scheme.residual = transporting_velocity.update_value(scheme.residual, uadv) - - def log_timestep(self): - """ - Logs the start of a time step. - """ - logger.info('') - logger.info('='*40) - logger.info(f'at start of timestep {self.step}, t={float(self.t)}, dt={float(self.dt)}') - - def run(self, t, tmax, pick_up=False): - """ - Runs the model for the specified time, from t to tmax - - Args: - t (float): the start time of the run - tmax (float): the end time of the run - pick_up: (bool): specify whether to pick_up from a previous run - """ - - # Set up diagnostics, which may set up some fields necessary to pick up - self.io.setup_diagnostics(self.fields) - self.io.setup_log_courant(self.fields) - if self.equation.domain.mesh.extruded: - self.io.setup_log_courant(self.fields, component='horizontal') - self.io.setup_log_courant(self.fields, component='vertical') - if self.transporting_velocity != "prognostic": - self.io.setup_log_courant(self.fields, name='transporting_velocity', - expression=self.transporting_velocity) - - if pick_up: - # Pick up fields, and return other info to be picked up - t, reference_profiles, self.step, initial_timesteps = self.io.pick_up_from_checkpoint(self.fields) - self.set_reference_profiles(reference_profiles) - self.set_initial_timesteps(initial_timesteps) - else: - self.step = 1 - - # Set up dump, which may also include an initial dump - with timed_stage("Dump output"): - self.io.setup_dump(self.fields, t, pick_up) - - self.t.assign(t) - - # Time loop - while float(self.t) < tmax - 0.5*float(self.dt): - self.log_timestep() - - self.x.update() - - self.io.log_courant(self.fields) - if self.equation.domain.mesh.extruded: - self.io.log_courant(self.fields, component='horizontal', message='horizontal') - self.io.log_courant(self.fields, component='vertical', message='vertical') - - self.timestep() - - self.t.assign(self.t + self.dt) - self.step += 1 - - with timed_stage("Dump output"): - self.io.dump(self.fields, float(self.t), self.step, self.get_initial_timesteps()) - - if self.io.output.checkpoint and self.io.output.checkpoint_method == 'dumbcheckpoint': - self.io.chkpt.close() - - logger.info(f'TIMELOOP complete. t={float(self.t)}, tmax={tmax}') - - def set_reference_profiles(self, reference_profiles): - """ - Initialise the model's reference profiles. - - reference_profiles (list): an iterable of pairs: (field_name, expr), - where 'field_name' is the string giving the name of the reference - profile field expr is the :class:`ufl.Expr` whose value is used to - set the reference field. - """ - for field_name, profile in reference_profiles: - if field_name+'_bar' in self.fields: - # For reference profiles already added to state, allow - # interpolation from expressions - ref = self.fields(field_name+'_bar') - elif isinstance(profile, Function): - # Need to add reference profile to state so profile must be - # a Function - ref = self.fields(field_name+'_bar', space=profile.function_space(), - pick_up=True, dump=False, field_type='reference') - else: - raise ValueError(f'When initialising reference profile {field_name}' - + ' the passed profile must be a Function') - # if field name is not prognostic we need to add it - ref.interpolate(profile) - # Assign profile to X_ref belonging to equation - if isinstance(self.equation, PrognosticEquationSet): - if field_name in self.equation.field_names: - idx = self.equation.field_names.index(field_name) - X_ref = self.equation.X_ref.subfunctions[idx] - X_ref.assign(ref) - else: - # reference profile of a diagnostic - # warn user in case they made a typo - logger.warning(f'Setting reference profile for diagnostic {field_name}') - # Don't need to do anything else as value in field container has already been set - self.reference_profiles_initialised = True - - -class Timestepper(BaseTimestepper): - """ - Implements a timeloop by applying a scheme to a prognostic equation. - """ - - def __init__(self, equation, scheme, io, spatial_methods=None, - physics_parametrisations=None): - """ - Args: - equation (:class:`PrognosticEquation`): the prognostic equation - scheme (:class:`TimeDiscretisation`): the scheme to use to timestep - the prognostic equation - io (:class:`IO`): the model's object for controlling input/output. - spatial_methods (iter, optional): a list of objects describing the - methods to use for discretising transport or diffusion terms - for each transported/diffused variable. Defaults to None, - in which case the terms follow the original discretisation in - the equation. - physics_parametrisations: (iter, optional): an iterable of - :class:`PhysicsParametrisation` objects that describe physical - parametrisations to be included to add to the equation. They can - only be used when the time discretisation `scheme` is explicit. - Defaults to None. - """ - self.scheme = scheme - if spatial_methods is not None: - self.spatial_methods = spatial_methods - else: - self.spatial_methods = [] - - if physics_parametrisations is not None: - self.physics_parametrisations = physics_parametrisations - if len(self.physics_parametrisations) > 1: - assert isinstance(scheme, ExplicitTimeDiscretisation), \ - ('Physics parametrisations can only be used with the ' - + 'basic TimeStepper when the time discretisation is ' - + 'explicit. If you want to use an implicit scheme, the ' - + 'SplitPhysicsTimestepper is more appropriate.') - else: - self.physics_parametrisations = [] - - super().__init__(equation=equation, io=io) - - @property - def transporting_velocity(self): - return "prognostic" - - def setup_fields(self): - self.x = TimeLevelFields(self.equation, self.scheme.nlevels) - self.fields = StateFields(self.x, self.equation.prescribed_fields, - *self.io.output.dumplist) - - def setup_scheme(self): - self.setup_equation(self.equation) - self.scheme.setup(self.equation) - self.setup_transporting_velocity(self.scheme) - if self.io.output.log_courant: - self.scheme.courant_max = self.io.courant_max - - def timestep(self): - """ - Implement the timestep - """ - xnp1 = self.x.np1 - name = self.equation.field_name - x_in = [x(name) for x in self.x.previous[-self.scheme.nlevels:]] - - self.scheme.apply(xnp1(name), *x_in) - - -class SplitPhysicsTimestepper(Timestepper): - """ - Implements a timeloop by applying schemes separately to the physics and - dynamics. This 'splits' the physics from the dynamics and allows a different - scheme to be applied to the physics terms than the prognostic equation. - """ - - def __init__(self, equation, scheme, io, spatial_methods=None, - physics_schemes=None): - """ - Args: - equation (:class:`PrognosticEquation`): the prognostic equation - scheme (:class:`TimeDiscretisation`): the scheme to use to timestep - the prognostic equation - io (:class:`IO`): the model's object for controlling input/output. - spatial_methods (iter,optional): a list of objects describing the - methods to use for discretising transport or diffusion terms - for each transported/diffused variable. Defaults to None, - in which case the terms follow the original discretisation in - the equation. - physics_schemes: (list, optional): a list of tuples of the form - (:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`), - pairing physics parametrisations and timestepping schemes to use - for each. Timestepping schemes for physics must be explicit. - Defaults to None. - """ - - # As we handle physics differently to the Timestepper, these are not - # passed to the super __init__ - super().__init__(equation, scheme, io, spatial_methods=spatial_methods) - - if physics_schemes is not None: - self.physics_schemes = physics_schemes - else: - self.physics_schemes = [] - - for parametrisation, phys_scheme in self.physics_schemes: - # check that the supplied schemes for physics are valid - if hasattr(parametrisation, "explicit_only") and parametrisation.explicit_only: - assert isinstance(phys_scheme, ExplicitTimeDiscretisation), \ - ("Only explicit time discretisations can be used with " - + f"physics scheme {parametrisation.label.label}") - apply_bcs = False - phys_scheme.setup(equation, apply_bcs, parametrisation.label) - - @property - def transporting_velocity(self): - return "prognostic" - - def setup_scheme(self): - self.setup_equation(self.equation) - # Go through and label all non-physics terms with a "dynamics" label - dynamics = Label('dynamics') - self.equation.label_terms(lambda t: not any(t.has_label(time_derivative, physics_label)), dynamics) - apply_bcs = True - self.scheme.setup(self.equation, apply_bcs, dynamics) - self.setup_transporting_velocity(self.scheme) - if self.io.output.log_courant: - self.scheme.courant_max = self.io.courant_max - - def timestep(self): - - super().timestep() - - with timed_stage("Physics"): - for _, scheme in self.physics_schemes: - scheme.apply(self.x.np1(scheme.field_name), self.x.np1(scheme.field_name)) - - -class SplitPrescribedTransport(Timestepper): - """ - Implements a timeloop where the physics terms are solved separately from - the dynamics, like with SplitPhysicsTimestepper, but here we define - a prescribed transporting velocity, as opposed to having the - velocity as a prognostic variable. - """ - - def __init__(self, equation, scheme, io, spatial_methods=None, - physics_schemes=None, - prescribed_transporting_velocity=None): - """ - Args: - equation (:class:`PrognosticEquation`): the prognostic equation - scheme (:class:`TimeDiscretisation`): the scheme to use to timestep - the prognostic equation - io (:class:`IO`): the model's object for controlling input/output. - spatial_methods (iter,optional): a list of objects describing the - methods to use for discretising transport or diffusion terms - for each transported/diffused variable. Defaults to None, - in which case the terms follow the original discretisation in - the equation. - physics_schemes: (list, optional): a list of tuples of the form - (:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`), - pairing physics parametrisations and timestepping schemes to use - for each. Timestepping schemes for physics can be explicit - or implicit. Defaults to None. - prescribed_transporting_velocity: (field, optional): A known - velocity field that is used for the transport of tracers. - This can be made time-varying by defining a python function - that uses time as an argument. - Defaults to None. - """ - - # As we handle physics differently to the Timestepper, these are not - # passed to the super __init__ - super().__init__(equation, scheme, io, spatial_methods=spatial_methods) - - if physics_schemes is not None: - self.physics_schemes = physics_schemes - else: - self.physics_schemes = [] - - for parametrisation, phys_scheme in self.physics_schemes: - # check that the supplied schemes for physics are valid - if hasattr(parametrisation, "explicit_only") and parametrisation.explicit_only: - assert isinstance(phys_scheme, ExplicitTimeDiscretisation), \ - ("Only explicit time discretisations can be used with " - + f"physics scheme {parametrisation.label.label}") - apply_bcs = False - phys_scheme.setup(equation, apply_bcs, parametrisation.label) - - if prescribed_transporting_velocity is not None: - self.velocity_projection = Projector( - prescribed_transporting_velocity(self.t), - self.fields('u')) - else: - self.velocity_projection = None - - @property - def transporting_velocity(self): - return self.fields('u') - - def setup_scheme(self): - self.setup_equation(self.equation) - # Go through and label all non-physics terms with a "dynamics" label - dynamics = Label('dynamics') - self.equation.label_terms(lambda t: not any(t.has_label(time_derivative, physics_label)), dynamics) - apply_bcs = True - self.scheme.setup(self.equation, apply_bcs, dynamics) - self.setup_transporting_velocity(self.scheme) - if self.io.output.log_courant: - self.scheme.courant_max = self.io.courant_max - - def timestep(self): - - if self.velocity_projection is not None: - self.velocity_projection.project() - - super().timestep() - - with timed_stage("Physics"): - for _, scheme in self.physics_schemes: - scheme.apply(self.x.np1(scheme.field_name), self.x.np1(scheme.field_name)) - - -class SemiImplicitQuasiNewton(BaseTimestepper): - """ - Implements a semi-implicit quasi-Newton discretisation, - with Strang splitting and auxiliary semi-Lagrangian transport. - - The timestep consists of an outer loop applying the transport and an - inner loop to perform the quasi-Newton interations for the fast-wave - terms. - """ - - def __init__(self, equation_set, io, transport_schemes, spatial_methods, - auxiliary_equations_and_schemes=None, linear_solver=None, - diffusion_schemes=None, physics_schemes=None, - slow_physics_schemes=None, fast_physics_schemes=None, - alpha=Constant(0.5), off_centred_u=False, - num_outer=2, num_inner=2): - - """ - Args: - equation_set (:class:`PrognosticEquationSet`): the prognostic - equation set to be solved - io (:class:`IO`): the model's object for controlling input/output. - transport_schemes: iterable of ``(field_name, scheme)`` pairs - indicating the name of the field (str) to transport, and the - :class:`TimeDiscretisation` to use - spatial_methods (iter): a list of objects describing the spatial - discretisations of transport or diffusion terms to be used. - auxiliary_equations_and_schemes: iterable of ``(equation, scheme)`` - pairs indicating any additional equations to be solved and the - scheme to use to solve them. Defaults to None. - linear_solver (:class:`TimesteppingSolver`, optional): the object - to use for the linear solve. Defaults to None. - diffusion_schemes (iter, optional): iterable of pairs of the form - ``(field_name, scheme)`` indicating the fields to diffuse, and - the :class:`~.TimeDiscretisation` to use. Defaults to None. - physics_schemes: (list, optional): a list of tuples of the form - (:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`), - pairing physics parametrisations and timestepping schemes to use - for each. Timestepping schemes for physics must be explicit. - These schemes are all evaluated at the end of the time step. - Defaults to None. - slow_physics_schemes: (list, optional): a list of tuples of the form - (:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`). - These schemes are all evaluated at the start of the time step. - Defaults to None. - fast_physics_schemes: (list, optional): a list of tuples of the form - (:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`). - These schemes are evaluated within the outer loop. Defaults to - None. - alpha (`ufl.Constant`, optional): the semi-implicit off-centering - parameter. A value of 1 corresponds to fully implicit, while 0 - corresponds to fully explicit. Defaults to Constant(0.5). - off_centred_u (bool, optional): option to offcentre the transporting - velocity. Defaults to False, in which case transporting velocity - is centred. If True offcentring uses value of alpha. - num_outer (int, optional): number of outer iterations in the semi- - implicit algorithm. The outer loop includes transport and any - fast physics schemes. Defaults to 2. Note that default used by - the Met Office's ENDGame and GungHo models is 2. - num_inner (int, optional): number of inner iterations in the semi- - implicit algorithm. The inner loop includes the evaluation of - implicit forcing (pressure gradient and Coriolis) terms, and the - linear solve. Defaults to 2. Note that default used by the Met - Office's ENDGame and GungHo models is 2. - """ - - self.num_outer = num_outer - self.num_inner = num_inner - self.alpha = alpha - - # default is to not offcentre transporting velocity but if it - # is offcentred then use the same value as alpha - self.alpha_u = Constant(alpha) if off_centred_u else Constant(0.5) - - self.spatial_methods = spatial_methods - - if physics_schemes is not None: - self.final_physics_schemes = physics_schemes - else: - self.final_physics_schemes = [] - if slow_physics_schemes is not None: - self.slow_physics_schemes = slow_physics_schemes - else: - self.slow_physics_schemes = [] - if fast_physics_schemes is not None: - self.fast_physics_schemes = fast_physics_schemes - else: - self.fast_physics_schemes = [] - self.all_physics_schemes = (self.slow_physics_schemes - + self.fast_physics_schemes - + self.final_physics_schemes) - - for parametrisation, scheme in self.all_physics_schemes: - assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" - if hasattr(parametrisation, "explicit_only") and parametrisation.explicit_only: - assert isinstance(scheme, ExplicitTimeDiscretisation), \ - ("Only explicit time discretisations can be used with " - + f"physics scheme {parametrisation.label.label}") - - self.active_transport = [] - for scheme in transport_schemes: - assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" - assert scheme.field_name in equation_set.field_names - self.active_transport.append((scheme.field_name, scheme)) - # Check that there is a corresponding transport method - method_found = False - for method in spatial_methods: - if scheme.field_name == method.variable and method.term_label == transport: - method_found = True - assert method_found, f'No transport method found for variable {scheme.field_name}' - - self.diffusion_schemes = [] - if diffusion_schemes is not None: - for scheme in diffusion_schemes: - assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" - assert scheme.field_name in equation_set.field_names - self.diffusion_schemes.append((scheme.field_name, scheme)) - # Check that there is a corresponding transport method - method_found = False - for method in spatial_methods: - if scheme.field_name == method.variable and method.term_label == diffusion: - method_found = True - assert method_found, f'No diffusion method found for variable {scheme.field_name}' - - if auxiliary_equations_and_schemes is not None: - for eqn, scheme in auxiliary_equations_and_schemes: - assert not hasattr(eqn, "field_names"), 'Cannot use auxiliary schemes with multiple fields' - self.auxiliary_schemes = [ - (eqn.field_name, scheme) - for eqn, scheme in auxiliary_equations_and_schemes] - - else: - auxiliary_equations_and_schemes = [] - self.auxiliary_schemes = [] - self.auxiliary_equations_and_schemes = auxiliary_equations_and_schemes - - super().__init__(equation_set, io) - - for aux_eqn, aux_scheme in self.auxiliary_equations_and_schemes: - self.setup_equation(aux_eqn) - aux_scheme.setup(aux_eqn) - self.setup_transporting_velocity(aux_scheme) - - self.tracers_to_copy = [] - for name in equation_set.field_names: - # Extract time derivative for that prognostic - mass_form = equation_set.residual.label_map( - lambda t: (t.has_label(time_derivative) and t.get(prognostic) == name), - map_if_false=drop) - # Copy over field if the time derivative term has no linearisation - if not mass_form.terms[0].has_label(linearisation): - self.tracers_to_copy.append(name) - - self.field_name = equation_set.field_name - W = equation_set.function_space - self.xrhs = Function(W) - self.xrhs_phys = Function(W) - self.dy = Function(W) - if linear_solver is None: - self.linear_solver = LinearTimesteppingSolver(equation_set, self.alpha) - else: - self.linear_solver = linear_solver - self.forcing = Forcing(equation_set, self.alpha) - self.bcs = equation_set.bcs - - def _apply_bcs(self): - """ - Set the zero boundary conditions in the velocity. - """ - unp1 = self.x.np1("u") - - for bc in self.bcs['u']: - bc.apply(unp1) - - @property - def transporting_velocity(self): - """Computes ubar=(1-alpha)*un + alpha*unp1""" - xn = self.x.n - xnp1 = self.x.np1 - # computes ubar from un and unp1 - return xn('u') + self.alpha_u*(xnp1('u')-xn('u')) - - def setup_fields(self): - """Sets up time levels n, star, p and np1""" - self.x = TimeLevelFields(self.equation, 1) - self.x.add_fields(self.equation, levels=("star", "p", "after_slow", "after_fast")) - for aux_eqn, _ in self.auxiliary_equations_and_schemes: - self.x.add_fields(aux_eqn) - # Prescribed fields for auxiliary eqns should come from prognostics of - # other equations, so only the prescribed fields of the main equation - # need passing to StateFields - self.fields = StateFields(self.x, self.equation.prescribed_fields, - *self.io.output.dumplist) - - def setup_scheme(self): - """Sets up transport, diffusion and physics schemes""" - # TODO: apply_bcs should be False for advection but this means - # tests with KGOs fail - apply_bcs = True - self.setup_equation(self.equation) - for _, scheme in self.active_transport: - scheme.setup(self.equation, apply_bcs, transport) - self.setup_transporting_velocity(scheme) - if self.io.output.log_courant: - scheme.courant_max = self.io.courant_max - - apply_bcs = True - for _, scheme in self.diffusion_schemes: - scheme.setup(self.equation, apply_bcs, diffusion) - for parametrisation, scheme in self.all_physics_schemes: - apply_bcs = True - scheme.setup(self.equation, apply_bcs, parametrisation.label) - - def copy_active_tracers(self, x_in, x_out): - """ - Copies active tracers from one set of fields to another, if those fields - are not included in the linear solver. This is determined by whether the - time derivative term for that tracer has a linearisation. - - Args: - x_in: The input set of fields - x_out: The output set of fields - """ - - for name in self.tracers_to_copy: - x_out(name).assign(x_in(name)) - - def timestep(self): - """Defines the timestep""" - xn = self.x.n - xnp1 = self.x.np1 - xstar = self.x.star - xp = self.x.p - x_after_slow = self.x.after_slow - x_after_fast = self.x.after_fast - xrhs = self.xrhs - xrhs_phys = self.xrhs_phys - dy = self.dy - - x_after_slow(self.field_name).assign(xn(self.field_name)) - if len(self.slow_physics_schemes) > 0: - with timed_stage("Slow physics"): - logger.info('SIQN: Slow physics') - for _, scheme in self.slow_physics_schemes: - scheme.apply(x_after_slow(scheme.field_name), x_after_slow(scheme.field_name)) - - with timed_stage("Apply forcing terms"): - logger.info('SIQN: Explicit forcing') - # Put explicit forcing into xstar - self.forcing.apply(x_after_slow, xn, xstar(self.field_name), "explicit") - - # set xp here so that variables that are not transported have - # the correct values - xp(self.field_name).assign(xstar(self.field_name)) - - for outer in range(self.num_outer): - - with timed_stage("Transport"): - self.io.log_courant(self.fields, 'transporting_velocity', - message=f'transporting velocity, outer iteration {outer}') - for name, scheme in self.active_transport: - logger.info(f'SIQN: Transport {outer}: {name}') - # transports a field from xstar and puts result in xp - scheme.apply(xp(name), xstar(name)) - - x_after_fast(self.field_name).assign(xp(self.field_name)) - if len(self.fast_physics_schemes) > 0: - with timed_stage("Fast physics"): - logger.info(f'SIQN: Fast physics {outer}') - for _, scheme in self.fast_physics_schemes: - scheme.apply(x_after_fast(scheme.field_name), x_after_fast(scheme.field_name)) - - xrhs.assign(0.) # xrhs is the residual which goes in the linear solve - xrhs_phys.assign(x_after_fast(self.field_name) - xp(self.field_name)) - - for inner in range(self.num_inner): - - # TODO: this is where to update the reference state - - with timed_stage("Apply forcing terms"): - logger.info(f'SIQN: Implicit forcing {(outer, inner)}') - self.forcing.apply(xp, xnp1, xrhs, "implicit") - - xrhs -= xnp1(self.field_name) - xrhs += xrhs_phys - - with timed_stage("Implicit solve"): - logger.info(f'SIQN: Mixed solve {(outer, inner)}') - self.linear_solver.solve(xrhs, dy) # solves linear system and places result in dy - - xnp1X = xnp1(self.field_name) - xnp1X += dy - - # Update xnp1 values for active tracers not included in the linear solve - self.copy_active_tracers(x_after_fast, xnp1) - - self._apply_bcs() - - for name, scheme in self.auxiliary_schemes: - # transports a field from xn and puts result in xnp1 - scheme.apply(xnp1(name), xn(name)) - - with timed_stage("Diffusion"): - for name, scheme in self.diffusion_schemes: - scheme.apply(xnp1(name), xnp1(name)) - - if len(self.final_physics_schemes) > 0: - with timed_stage("Final Physics"): - for _, scheme in self.final_physics_schemes: - scheme.apply(xnp1(scheme.field_name), xnp1(scheme.field_name)) - - def run(self, t, tmax, pick_up=False): - """ - Runs the model for the specified time, from t to tmax. - - Args: - t (float): the start time of the run - tmax (float): the end time of the run - pick_up: (bool): specify whether to pick_up from a previous run - """ - - if not pick_up: - assert self.reference_profiles_initialised, \ - 'Reference profiles for must be initialised to use Semi-Implicit Timestepper' - - super().run(t, tmax, pick_up=pick_up) - - -class PrescribedTransport(Timestepper): - """ - Implements a timeloop with a prescibed transporting velocity. - """ - def __init__(self, equation, scheme, io, transport_method, - physics_parametrisations=None, - prescribed_transporting_velocity=None): - """ - Args: - equation (:class:`PrognosticEquation`): the prognostic equation - scheme (:class:`TimeDiscretisation`): the scheme to use to timestep - the prognostic equation. - transport_method (:class:`TransportMethod`): describes the method - used for discretising the transport term. - io (:class:`IO`): the model's object for controlling input/output. - physics_schemes: (list, optional): a list of :class:`Physics` and - :class:`TimeDiscretisation` options describing physical - parametrisations and timestepping schemes to use for each. - Timestepping schemes for physics must be explicit. Defaults to - None. - physics_parametrisations: (iter, optional): an iterable of - :class:`PhysicsParametrisation` objects that describe physical - parametrisations to be included to add to the equation. They can - only be used when the time discretisation `scheme` is explicit. - Defaults to None. - """ - - if isinstance(transport_method, TransportMethod): - transport_methods = [transport_method] - else: - # Assume an iterable has been provided - transport_methods = transport_method - - super().__init__(equation, scheme, io, spatial_methods=transport_methods, - physics_parametrisations=physics_parametrisations) - - if prescribed_transporting_velocity is not None: - self.velocity_projection = Projector( - prescribed_transporting_velocity(self.t), - self.fields('u')) - else: - self.velocity_projection = None - - @property - def transporting_velocity(self): - return self.fields('u') - - def setup_fields(self): - self.x = TimeLevelFields(self.equation, self.scheme.nlevels) - self.fields = StateFields(self.x, self.equation.prescribed_fields, - *self.io.output.dumplist) - - def run(self, t, tmax, pick_up=False): - """ - Runs the model for the specified time, from t to tmax - Args: - t (float): the start time of the run - tmax (float): the end time of the run - pick_up: (bool): specify whether to pick_up from a previous run - """ - # It's best to have evaluated the velocity before we start - if self.velocity_projection is not None: - self.velocity_projection.project() - - super().run(t, tmax, pick_up=pick_up) - - def timestep(self): - if self.velocity_projection is not None: - self.velocity_projection.project() - - super().timestep() diff --git a/gusto/timestepping/__init__.py b/gusto/timestepping/__init__.py new file mode 100644 index 000000000..cb216fa33 --- /dev/null +++ b/gusto/timestepping/__init__.py @@ -0,0 +1,3 @@ +from gusto.timestepping.timestepper import * # noqa +from gusto.timestepping.split_timestepper import * # noqa +from gusto.timestepping.semi_implicit_quasi_newton import * # noqa \ No newline at end of file diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py new file mode 100644 index 000000000..bb933ed8e --- /dev/null +++ b/gusto/timestepping/semi_implicit_quasi_newton.py @@ -0,0 +1,475 @@ +""" +The Semi-Implicit Quasi-Newton timestepper used by the Met Office's ENDGame +and GungHo dynamical cores. +""" + +from firedrake import (Function, Constant, TrialFunctions, DirichletBC, + LinearVariationalProblem, LinearVariationalSolver) +from firedrake.fml import drop, replace_subject +from pyop2.profiling import timed_stage +from gusto.core import TimeLevelFields, StateFields +from gusto.core.labels import (transport, diffusion, time_derivative, + linearisation, prognostic, hydrostatic, + physics_label, sponge, incompressible) +from gusto.solvers import LinearTimesteppingSolver +from gusto.core.logging import logger, DEBUG, logging_ksp_monitor_true_residual +from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation +from gusto.timestepping.timestepper import BaseTimestepper + + +__all__ = ["SemiImplicitQuasiNewton"] + + +class SemiImplicitQuasiNewton(BaseTimestepper): + """ + Implements a semi-implicit quasi-Newton discretisation, + with Strang splitting and auxiliary semi-Lagrangian transport. + + The timestep consists of an outer loop applying the transport and an + inner loop to perform the quasi-Newton interations for the fast-wave + terms. + """ + + def __init__(self, equation_set, io, transport_schemes, spatial_methods, + auxiliary_equations_and_schemes=None, linear_solver=None, + diffusion_schemes=None, physics_schemes=None, + slow_physics_schemes=None, fast_physics_schemes=None, + alpha=Constant(0.5), off_centred_u=False, + num_outer=2, num_inner=2): + + """ + Args: + equation_set (:class:`PrognosticEquationSet`): the prognostic + equation set to be solved + io (:class:`IO`): the model's object for controlling input/output. + transport_schemes: iterable of ``(field_name, scheme)`` pairs + indicating the name of the field (str) to transport, and the + :class:`TimeDiscretisation` to use + spatial_methods (iter): a list of objects describing the spatial + discretisations of transport or diffusion terms to be used. + auxiliary_equations_and_schemes: iterable of ``(equation, scheme)`` + pairs indicating any additional equations to be solved and the + scheme to use to solve them. Defaults to None. + linear_solver (:class:`TimesteppingSolver`, optional): the object + to use for the linear solve. Defaults to None. + diffusion_schemes (iter, optional): iterable of pairs of the form + ``(field_name, scheme)`` indicating the fields to diffuse, and + the :class:`~.TimeDiscretisation` to use. Defaults to None. + physics_schemes: (list, optional): a list of tuples of the form + (:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`), + pairing physics parametrisations and timestepping schemes to use + for each. Timestepping schemes for physics must be explicit. + These schemes are all evaluated at the end of the time step. + Defaults to None. + slow_physics_schemes: (list, optional): a list of tuples of the form + (:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`). + These schemes are all evaluated at the start of the time step. + Defaults to None. + fast_physics_schemes: (list, optional): a list of tuples of the form + (:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`). + These schemes are evaluated within the outer loop. Defaults to + None. + alpha (`ufl.Constant`, optional): the semi-implicit off-centering + parameter. A value of 1 corresponds to fully implicit, while 0 + corresponds to fully explicit. Defaults to Constant(0.5). + off_centred_u (bool, optional): option to offcentre the transporting + velocity. Defaults to False, in which case transporting velocity + is centred. If True offcentring uses value of alpha. + num_outer (int, optional): number of outer iterations in the semi- + implicit algorithm. The outer loop includes transport and any + fast physics schemes. Defaults to 2. Note that default used by + the Met Office's ENDGame and GungHo models is 2. + num_inner (int, optional): number of inner iterations in the semi- + implicit algorithm. The inner loop includes the evaluation of + implicit forcing (pressure gradient and Coriolis) terms, and the + linear solve. Defaults to 2. Note that default used by the Met + Office's ENDGame and GungHo models is 2. + """ + + self.num_outer = num_outer + self.num_inner = num_inner + self.alpha = alpha + + # default is to not offcentre transporting velocity but if it + # is offcentred then use the same value as alpha + self.alpha_u = Constant(alpha) if off_centred_u else Constant(0.5) + + self.spatial_methods = spatial_methods + + if physics_schemes is not None: + self.final_physics_schemes = physics_schemes + else: + self.final_physics_schemes = [] + if slow_physics_schemes is not None: + self.slow_physics_schemes = slow_physics_schemes + else: + self.slow_physics_schemes = [] + if fast_physics_schemes is not None: + self.fast_physics_schemes = fast_physics_schemes + else: + self.fast_physics_schemes = [] + self.all_physics_schemes = (self.slow_physics_schemes + + self.fast_physics_schemes + + self.final_physics_schemes) + + for parametrisation, scheme in self.all_physics_schemes: + assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" + if hasattr(parametrisation, "explicit_only") and parametrisation.explicit_only: + assert isinstance(scheme, ExplicitTimeDiscretisation), \ + ("Only explicit time discretisations can be used with " + + f"physics scheme {parametrisation.label.label}") + + self.active_transport = [] + for scheme in transport_schemes: + assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" + assert scheme.field_name in equation_set.field_names + self.active_transport.append((scheme.field_name, scheme)) + # Check that there is a corresponding transport method + method_found = False + for method in spatial_methods: + if scheme.field_name == method.variable and method.term_label == transport: + method_found = True + assert method_found, f'No transport method found for variable {scheme.field_name}' + + self.diffusion_schemes = [] + if diffusion_schemes is not None: + for scheme in diffusion_schemes: + assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" + assert scheme.field_name in equation_set.field_names + self.diffusion_schemes.append((scheme.field_name, scheme)) + # Check that there is a corresponding transport method + method_found = False + for method in spatial_methods: + if scheme.field_name == method.variable and method.term_label == diffusion: + method_found = True + assert method_found, f'No diffusion method found for variable {scheme.field_name}' + + if auxiliary_equations_and_schemes is not None: + for eqn, scheme in auxiliary_equations_and_schemes: + assert not hasattr(eqn, "field_names"), 'Cannot use auxiliary schemes with multiple fields' + self.auxiliary_schemes = [ + (eqn.field_name, scheme) + for eqn, scheme in auxiliary_equations_and_schemes] + + else: + auxiliary_equations_and_schemes = [] + self.auxiliary_schemes = [] + self.auxiliary_equations_and_schemes = auxiliary_equations_and_schemes + + super().__init__(equation_set, io) + + for aux_eqn, aux_scheme in self.auxiliary_equations_and_schemes: + self.setup_equation(aux_eqn) + aux_scheme.setup(aux_eqn) + self.setup_transporting_velocity(aux_scheme) + + self.tracers_to_copy = [] + for name in equation_set.field_names: + # Extract time derivative for that prognostic + mass_form = equation_set.residual.label_map( + lambda t: (t.has_label(time_derivative) and t.get(prognostic) == name), + map_if_false=drop) + # Copy over field if the time derivative term has no linearisation + if not mass_form.terms[0].has_label(linearisation): + self.tracers_to_copy.append(name) + + self.field_name = equation_set.field_name + W = equation_set.function_space + self.xrhs = Function(W) + self.xrhs_phys = Function(W) + self.dy = Function(W) + if linear_solver is None: + self.linear_solver = LinearTimesteppingSolver(equation_set, self.alpha) + else: + self.linear_solver = linear_solver + self.forcing = Forcing(equation_set, self.alpha) + self.bcs = equation_set.bcs + + def _apply_bcs(self): + """ + Set the zero boundary conditions in the velocity. + """ + unp1 = self.x.np1("u") + + for bc in self.bcs['u']: + bc.apply(unp1) + + @property + def transporting_velocity(self): + """Computes ubar=(1-alpha)*un + alpha*unp1""" + xn = self.x.n + xnp1 = self.x.np1 + # computes ubar from un and unp1 + return xn('u') + self.alpha_u*(xnp1('u')-xn('u')) + + def setup_fields(self): + """Sets up time levels n, star, p and np1""" + self.x = TimeLevelFields(self.equation, 1) + self.x.add_fields(self.equation, levels=("star", "p", "after_slow", "after_fast")) + for aux_eqn, _ in self.auxiliary_equations_and_schemes: + self.x.add_fields(aux_eqn) + # Prescribed fields for auxiliary eqns should come from prognostics of + # other equations, so only the prescribed fields of the main equation + # need passing to StateFields + self.fields = StateFields(self.x, self.equation.prescribed_fields, + *self.io.output.dumplist) + + def setup_scheme(self): + """Sets up transport, diffusion and physics schemes""" + # TODO: apply_bcs should be False for advection but this means + # tests with KGOs fail + apply_bcs = True + self.setup_equation(self.equation) + for _, scheme in self.active_transport: + scheme.setup(self.equation, apply_bcs, transport) + self.setup_transporting_velocity(scheme) + if self.io.output.log_courant: + scheme.courant_max = self.io.courant_max + + apply_bcs = True + for _, scheme in self.diffusion_schemes: + scheme.setup(self.equation, apply_bcs, diffusion) + for parametrisation, scheme in self.all_physics_schemes: + apply_bcs = True + scheme.setup(self.equation, apply_bcs, parametrisation.label) + + def copy_active_tracers(self, x_in, x_out): + """ + Copies active tracers from one set of fields to another, if those fields + are not included in the linear solver. This is determined by whether the + time derivative term for that tracer has a linearisation. + + Args: + x_in: The input set of fields + x_out: The output set of fields + """ + + for name in self.tracers_to_copy: + x_out(name).assign(x_in(name)) + + def timestep(self): + """Defines the timestep""" + xn = self.x.n + xnp1 = self.x.np1 + xstar = self.x.star + xp = self.x.p + x_after_slow = self.x.after_slow + x_after_fast = self.x.after_fast + xrhs = self.xrhs + xrhs_phys = self.xrhs_phys + dy = self.dy + + x_after_slow(self.field_name).assign(xn(self.field_name)) + if len(self.slow_physics_schemes) > 0: + with timed_stage("Slow physics"): + logger.info('SIQN: Slow physics') + for _, scheme in self.slow_physics_schemes: + scheme.apply(x_after_slow(scheme.field_name), x_after_slow(scheme.field_name)) + + with timed_stage("Apply forcing terms"): + logger.info('SIQN: Explicit forcing') + # Put explicit forcing into xstar + self.forcing.apply(x_after_slow, xn, xstar(self.field_name), "explicit") + + # set xp here so that variables that are not transported have + # the correct values + xp(self.field_name).assign(xstar(self.field_name)) + + for outer in range(self.num_outer): + + with timed_stage("Transport"): + self.io.log_courant(self.fields, 'transporting_velocity', + message=f'transporting velocity, outer iteration {outer}') + for name, scheme in self.active_transport: + logger.info(f'SIQN: Transport {outer}: {name}') + # transports a field from xstar and puts result in xp + scheme.apply(xp(name), xstar(name)) + + x_after_fast(self.field_name).assign(xp(self.field_name)) + if len(self.fast_physics_schemes) > 0: + with timed_stage("Fast physics"): + logger.info(f'SIQN: Fast physics {outer}') + for _, scheme in self.fast_physics_schemes: + scheme.apply(x_after_fast(scheme.field_name), x_after_fast(scheme.field_name)) + + xrhs.assign(0.) # xrhs is the residual which goes in the linear solve + xrhs_phys.assign(x_after_fast(self.field_name) - xp(self.field_name)) + + for inner in range(self.num_inner): + + # TODO: this is where to update the reference state + + with timed_stage("Apply forcing terms"): + logger.info(f'SIQN: Implicit forcing {(outer, inner)}') + self.forcing.apply(xp, xnp1, xrhs, "implicit") + + xrhs -= xnp1(self.field_name) + xrhs += xrhs_phys + + with timed_stage("Implicit solve"): + logger.info(f'SIQN: Mixed solve {(outer, inner)}') + self.linear_solver.solve(xrhs, dy) # solves linear system and places result in dy + + xnp1X = xnp1(self.field_name) + xnp1X += dy + + # Update xnp1 values for active tracers not included in the linear solve + self.copy_active_tracers(x_after_fast, xnp1) + + self._apply_bcs() + + for name, scheme in self.auxiliary_schemes: + # transports a field from xn and puts result in xnp1 + scheme.apply(xnp1(name), xn(name)) + + with timed_stage("Diffusion"): + for name, scheme in self.diffusion_schemes: + scheme.apply(xnp1(name), xnp1(name)) + + if len(self.final_physics_schemes) > 0: + with timed_stage("Final Physics"): + for _, scheme in self.final_physics_schemes: + scheme.apply(xnp1(scheme.field_name), xnp1(scheme.field_name)) + + def run(self, t, tmax, pick_up=False): + """ + Runs the model for the specified time, from t to tmax. + + Args: + t (float): the start time of the run + tmax (float): the end time of the run + pick_up: (bool): specify whether to pick_up from a previous run + """ + + if not pick_up: + assert self.reference_profiles_initialised, \ + 'Reference profiles for must be initialised to use Semi-Implicit Timestepper' + + super().run(t, tmax, pick_up=pick_up) + + +class Forcing(object): + """ + Discretises forcing terms for the Semi-Implicit Quasi-Newton timestepper. + + This class describes the evaluation of forcing terms, for instance the + gravitational force, the Coriolis force or the pressure gradient force. + These are terms that can simply be evaluated, generally as part of some + semi-implicit time discretisation. + """ + + def __init__(self, equation, alpha): + """ + Args: + equation (:class:`PrognosticEquationSet`): the prognostic equations + containing the forcing terms. + alpha (:class:`Constant`): semi-implicit off-centering factor. An + alpha of 0 corresponds to fully explicit, while a factor of 1 + corresponds to fully implicit. + """ + + self.field_name = equation.field_name + implicit_terms = [incompressible, sponge] + dt = equation.domain.dt + + W = equation.function_space + self.x0 = Function(W) + self.xF = Function(W) + + # set up boundary conditions on the u subspace of W + bcs = [DirichletBC(W.sub(0), bc.function_arg, bc.sub_domain) for bc in equation.bcs['u']] + + # drop terms relating to transport, diffusion and physics + residual = equation.residual.label_map( + 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 + trials = TrialFunctions(W) + a = residual.label_map(lambda t: t.has_label(time_derivative), + replace_subject(trials), + map_if_false=drop) + + # the explicit forms are multiplied by (1-alpha) and moved to the rhs + L_explicit = -(1-alpha)*dt*residual.label_map( + lambda t: + any(t.has_label(time_derivative, hydrostatic, *implicit_terms, + return_tuple=True)), + drop, + replace_subject(self.x0)) + + # the implicit forms are multiplied by alpha and moved to the rhs + L_implicit = -alpha*dt*residual.label_map( + lambda t: + any(t.has_label( + time_derivative, hydrostatic, *implicit_terms, + return_tuple=True)), + drop, + replace_subject(self.x0)) + + # now add the terms that are always fully implicit + L_implicit -= dt*residual.label_map( + lambda t: any(t.has_label(*implicit_terms, return_tuple=True)), + replace_subject(self.x0), + drop) + + # the hydrostatic equations require some additional forms: + if any([t.has_label(hydrostatic) for t in residual]): + L_explicit += residual.label_map( + lambda t: t.has_label(hydrostatic), + replace_subject(self.x0), + drop) + + L_implicit -= residual.label_map( + lambda t: t.has_label(hydrostatic), + replace_subject(self.x0), + drop) + + # now we can set up the explicit and implicit problems + explicit_forcing_problem = LinearVariationalProblem( + a.form, L_explicit.form, self.xF, bcs=bcs + ) + + implicit_forcing_problem = LinearVariationalProblem( + a.form, L_implicit.form, self.xF, bcs=bcs + ) + + self.solvers = {} + self.solvers["explicit"] = LinearVariationalSolver( + explicit_forcing_problem, + options_prefix="ExplicitForcingSolver" + ) + self.solvers["implicit"] = LinearVariationalSolver( + implicit_forcing_problem, + options_prefix="ImplicitForcingSolver" + ) + + if logger.isEnabledFor(DEBUG): + self.solvers["explicit"].snes.ksp.setMonitor(logging_ksp_monitor_true_residual) + self.solvers["implicit"].snes.ksp.setMonitor(logging_ksp_monitor_true_residual) + + def apply(self, x_in, x_nl, x_out, label): + """ + Applies the discretisation for a forcing term F(x). + + This takes x_in and x_nl and computes F(x_nl), and updates x_out to \n + x_out = x_in + scale*F(x_nl) \n + where 'scale' is the appropriate semi-implicit factor. + + Args: + x_in (:class:`FieldCreator`): the field to be incremented. + x_nl (:class:`FieldCreator`): the field which the forcing term is + evaluated on. + x_out (:class:`FieldCreator`): the output field to be updated. + label (str): denotes which forcing to apply. Should be 'explicit' or + 'implicit'. TODO: there should be a check on this. Or this + should be an actual label. + """ + + self.x0.assign(x_nl(self.field_name)) + + self.solvers[label].solve() # places forcing in self.xF + + x_out.assign(x_in(self.field_name)) + x_out += self.xF diff --git a/gusto/timestepping/split_timestepper.py b/gusto/timestepping/split_timestepper.py new file mode 100644 index 000000000..b38935987 --- /dev/null +++ b/gusto/timestepping/split_timestepper.py @@ -0,0 +1,165 @@ +"""Split timestepping methods for generically solving terms separately.""" + +from firedrake import Projector +from firedrake.fml import Label +from pyop2.profiling import timed_stage +from gusto.core.labels import time_derivative, physics_label +from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation +from gusto.timestepping.timestepper import Timestepper + +__all__ = ["SplitPhysicsTimestepper", "SplitPrescribedTransport"] + + +class SplitPhysicsTimestepper(Timestepper): + """ + Implements a timeloop by applying schemes separately to the physics and + dynamics. This 'splits' the physics from the dynamics and allows a different + scheme to be applied to the physics terms than the prognostic equation. + """ + + def __init__(self, equation, scheme, io, spatial_methods=None, + physics_schemes=None): + """ + Args: + equation (:class:`PrognosticEquation`): the prognostic equation + scheme (:class:`TimeDiscretisation`): the scheme to use to timestep + the prognostic equation + io (:class:`IO`): the model's object for controlling input/output. + spatial_methods (iter,optional): a list of objects describing the + methods to use for discretising transport or diffusion terms + for each transported/diffused variable. Defaults to None, + in which case the terms follow the original discretisation in + the equation. + physics_schemes: (list, optional): a list of tuples of the form + (:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`), + pairing physics parametrisations and timestepping schemes to use + for each. Timestepping schemes for physics must be explicit. + Defaults to None. + """ + + # As we handle physics differently to the Timestepper, these are not + # passed to the super __init__ + super().__init__(equation, scheme, io, spatial_methods=spatial_methods) + + if physics_schemes is not None: + self.physics_schemes = physics_schemes + else: + self.physics_schemes = [] + + for parametrisation, phys_scheme in self.physics_schemes: + # check that the supplied schemes for physics are valid + if hasattr(parametrisation, "explicit_only") and parametrisation.explicit_only: + assert isinstance(phys_scheme, ExplicitTimeDiscretisation), \ + ("Only explicit time discretisations can be used with " + + f"physics scheme {parametrisation.label.label}") + apply_bcs = False + phys_scheme.setup(equation, apply_bcs, parametrisation.label) + + @property + def transporting_velocity(self): + return "prognostic" + + def setup_scheme(self): + self.setup_equation(self.equation) + # Go through and label all non-physics terms with a "dynamics" label + dynamics = Label('dynamics') + self.equation.label_terms(lambda t: not any(t.has_label(time_derivative, physics_label)), dynamics) + apply_bcs = True + self.scheme.setup(self.equation, apply_bcs, dynamics) + self.setup_transporting_velocity(self.scheme) + if self.io.output.log_courant: + self.scheme.courant_max = self.io.courant_max + + def timestep(self): + + super().timestep() + + with timed_stage("Physics"): + for _, scheme in self.physics_schemes: + scheme.apply(self.x.np1(scheme.field_name), self.x.np1(scheme.field_name)) + + +class SplitPrescribedTransport(Timestepper): + """ + Implements a timeloop where the physics terms are solved separately from + the dynamics, like with SplitPhysicsTimestepper, but here we define + a prescribed transporting velocity, as opposed to having the + velocity as a prognostic variable. + """ + + def __init__(self, equation, scheme, io, spatial_methods=None, + physics_schemes=None, + prescribed_transporting_velocity=None): + """ + Args: + equation (:class:`PrognosticEquation`): the prognostic equation + scheme (:class:`TimeDiscretisation`): the scheme to use to timestep + the prognostic equation + io (:class:`IO`): the model's object for controlling input/output. + spatial_methods (iter,optional): a list of objects describing the + methods to use for discretising transport or diffusion terms + for each transported/diffused variable. Defaults to None, + in which case the terms follow the original discretisation in + the equation. + physics_schemes: (list, optional): a list of tuples of the form + (:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`), + pairing physics parametrisations and timestepping schemes to use + for each. Timestepping schemes for physics can be explicit + or implicit. Defaults to None. + prescribed_transporting_velocity: (field, optional): A known + velocity field that is used for the transport of tracers. + This can be made time-varying by defining a python function + that uses time as an argument. + Defaults to None. + """ + + # As we handle physics differently to the Timestepper, these are not + # passed to the super __init__ + super().__init__(equation, scheme, io, spatial_methods=spatial_methods) + + if physics_schemes is not None: + self.physics_schemes = physics_schemes + else: + self.physics_schemes = [] + + for parametrisation, phys_scheme in self.physics_schemes: + # check that the supplied schemes for physics are valid + if hasattr(parametrisation, "explicit_only") and parametrisation.explicit_only: + assert isinstance(phys_scheme, ExplicitTimeDiscretisation), \ + ("Only explicit time discretisations can be used with " + + f"physics scheme {parametrisation.label.label}") + apply_bcs = False + phys_scheme.setup(equation, apply_bcs, parametrisation.label) + + if prescribed_transporting_velocity is not None: + self.velocity_projection = Projector( + prescribed_transporting_velocity(self.t), + self.fields('u')) + else: + self.velocity_projection = None + + @property + def transporting_velocity(self): + return self.fields('u') + + def setup_scheme(self): + self.setup_equation(self.equation) + # Go through and label all non-physics terms with a "dynamics" label + dynamics = Label('dynamics') + self.equation.label_terms(lambda t: not any(t.has_label(time_derivative, physics_label)), dynamics) + apply_bcs = True + self.scheme.setup(self.equation, apply_bcs, dynamics) + self.setup_transporting_velocity(self.scheme) + if self.io.output.log_courant: + self.scheme.courant_max = self.io.courant_max + + def timestep(self): + + if self.velocity_projection is not None: + self.velocity_projection.project() + + super().timestep() + + with timed_stage("Physics"): + for _, scheme in self.physics_schemes: + scheme.apply(self.x.np1(scheme.field_name), self.x.np1(scheme.field_name)) diff --git a/gusto/timestepping/timestepper.py b/gusto/timestepping/timestepper.py new file mode 100644 index 000000000..e8e1a1d38 --- /dev/null +++ b/gusto/timestepping/timestepper.py @@ -0,0 +1,380 @@ +"""Defines the basic timestepper objects.""" + +from abc import ABCMeta, abstractmethod, abstractproperty +from firedrake import Function, Projector, split +from firedrake.fml import drop, Term +from pyop2.profiling import timed_stage +from gusto.equations import PrognosticEquationSet +from gusto.core import TimeLevelFields, StateFields +from gusto.core.labels import transport, diffusion, prognostic, transporting_velocity +from gusto.core.logging import logger +from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation +from gusto.spatial_methods.transport_methods import TransportMethod +import ufl + +__all__ = ["BaseTimestepper", "Timestepper", "PrescribedTransport"] + + +class BaseTimestepper(object, metaclass=ABCMeta): + """Base class for timesteppers.""" + + def __init__(self, equation, io): + """ + Args: + equation (:class:`PrognosticEquation`): the prognostic equation. + io (:class:`IO`): the model's object for controlling input/output. + """ + + self.equation = equation + self.io = io + self.dt = self.equation.domain.dt + self.t = self.equation.domain.t + self.reference_profiles_initialised = False + + self.setup_fields() + self.setup_scheme() + + self.io.log_parameters(equation) + + @abstractproperty + def transporting_velocity(self): + return NotImplementedError + + @abstractmethod + def setup_fields(self): + """Set up required fields. Must be implemented in child classes""" + pass + + @abstractmethod + def setup_scheme(self): + """Set up required scheme(s). Must be implemented in child classes""" + pass + + @abstractmethod + def timestep(self): + """Defines the timestep. Must be implemented in child classes""" + return NotImplementedError + + def set_initial_timesteps(self, num_steps): + """Sets the number of initial time steps for a multi-level scheme.""" + can_set = (hasattr(self, 'scheme') + and hasattr(self.scheme, 'initial_timesteps') + and num_steps is not None) + if can_set: + self.scheme.initial_timesteps = num_steps + + def get_initial_timesteps(self): + """Gets the number of initial time steps from a multi-level scheme.""" + can_get = (hasattr(self, 'scheme') + and hasattr(self.scheme, 'initial_timesteps')) + # Return None if this is not applicable + return self.scheme.initial_timesteps if can_get else None + + def setup_equation(self, equation): + """ + Sets up the spatial methods for an equation, by the setting the + forms used for transport/diffusion in the equation. + + Args: + equation (:class:`PrognosticEquation`): the equation that the + transport method is to be applied to. + """ + + # For now, we only have methods for transport and diffusion + for term_label in [transport, diffusion]: + # ---------------------------------------------------------------- # + # Check that appropriate methods have been provided + # ---------------------------------------------------------------- # + # Extract all terms corresponding to this type of term + residual = equation.residual.label_map( + lambda t: t.has_label(term_label), map_if_false=drop + ) + variables = [t.get(prognostic) for t in residual.terms] + methods = list(filter(lambda t: t.term_label == term_label, + self.spatial_methods)) + method_variables = [method.variable for method in methods] + for variable in variables: + if variable not in method_variables: + message = f'Variable {variable} has a {term_label.label} ' \ + + 'term but no method for this has been specified. ' \ + + 'Using default form for this term' + logger.warning(message) + + # -------------------------------------------------------------------- # + # Check that appropriate methods have been provided + # -------------------------------------------------------------------- # + # Replace forms in equation + if self.spatial_methods is not None: + for method in self.spatial_methods: + method.replace_form(equation) + + def setup_transporting_velocity(self, scheme): + """ + Set up the time discretisation by replacing the transporting velocity + used by the appropriate one for this time loop. + + Args: + scheme (:class:`TimeDiscretisation`): the time discretisation whose + transport term should be replaced with the transport term of + this discretisation. + """ + if self.transporting_velocity == "prognostic" and "u" in self.fields._field_names: + # Use the prognostic wind variable as the transporting velocity + u_idx = self.equation.field_names.index('u') + uadv = split(self.equation.X)[u_idx] + else: + uadv = self.transporting_velocity + + scheme.residual = scheme.residual.label_map( + lambda t: t.has_label(transporting_velocity), + map_if_true=lambda t: + Term(ufl.replace(t.form, {t.get(transporting_velocity): uadv}), t.labels) + ) + + scheme.residual = transporting_velocity.update_value(scheme.residual, uadv) + + def log_timestep(self): + """ + Logs the start of a time step. + """ + logger.info('') + logger.info('='*40) + logger.info(f'at start of timestep {self.step}, t={float(self.t)}, dt={float(self.dt)}') + + def run(self, t, tmax, pick_up=False): + """ + Runs the model for the specified time, from t to tmax + + Args: + t (float): the start time of the run + tmax (float): the end time of the run + pick_up: (bool): specify whether to pick_up from a previous run + """ + + # Set up diagnostics, which may set up some fields necessary to pick up + self.io.setup_diagnostics(self.fields) + self.io.setup_log_courant(self.fields) + if self.equation.domain.mesh.extruded: + self.io.setup_log_courant(self.fields, component='horizontal') + self.io.setup_log_courant(self.fields, component='vertical') + if self.transporting_velocity != "prognostic": + self.io.setup_log_courant(self.fields, name='transporting_velocity', + expression=self.transporting_velocity) + + if pick_up: + # Pick up fields, and return other info to be picked up + t, reference_profiles, self.step, initial_timesteps = self.io.pick_up_from_checkpoint(self.fields) + self.set_reference_profiles(reference_profiles) + self.set_initial_timesteps(initial_timesteps) + else: + self.step = 1 + + # Set up dump, which may also include an initial dump + with timed_stage("Dump output"): + self.io.setup_dump(self.fields, t, pick_up) + + self.t.assign(t) + + # Time loop + while float(self.t) < tmax - 0.5*float(self.dt): + self.log_timestep() + + self.x.update() + + self.io.log_courant(self.fields) + if self.equation.domain.mesh.extruded: + self.io.log_courant(self.fields, component='horizontal', message='horizontal') + self.io.log_courant(self.fields, component='vertical', message='vertical') + + self.timestep() + + self.t.assign(self.t + self.dt) + self.step += 1 + + with timed_stage("Dump output"): + self.io.dump(self.fields, float(self.t), self.step, self.get_initial_timesteps()) + + if self.io.output.checkpoint and self.io.output.checkpoint_method == 'dumbcheckpoint': + self.io.chkpt.close() + + logger.info(f'TIMELOOP complete. t={float(self.t)}, tmax={tmax}') + + def set_reference_profiles(self, reference_profiles): + """ + Initialise the model's reference profiles. + + reference_profiles (list): an iterable of pairs: (field_name, expr), + where 'field_name' is the string giving the name of the reference + profile field expr is the :class:`ufl.Expr` whose value is used to + set the reference field. + """ + for field_name, profile in reference_profiles: + if field_name+'_bar' in self.fields: + # For reference profiles already added to state, allow + # interpolation from expressions + ref = self.fields(field_name+'_bar') + elif isinstance(profile, Function): + # Need to add reference profile to state so profile must be + # a Function + ref = self.fields(field_name+'_bar', space=profile.function_space(), + pick_up=True, dump=False, field_type='reference') + else: + raise ValueError(f'When initialising reference profile {field_name}' + + ' the passed profile must be a Function') + # if field name is not prognostic we need to add it + ref.interpolate(profile) + # Assign profile to X_ref belonging to equation + if isinstance(self.equation, PrognosticEquationSet): + if field_name in self.equation.field_names: + idx = self.equation.field_names.index(field_name) + X_ref = self.equation.X_ref.subfunctions[idx] + X_ref.assign(ref) + else: + # reference profile of a diagnostic + # warn user in case they made a typo + logger.warning(f'Setting reference profile for diagnostic {field_name}') + # Don't need to do anything else as value in field container has already been set + self.reference_profiles_initialised = True + + +class Timestepper(BaseTimestepper): + """ + Implements a timeloop by applying a scheme to a prognostic equation. + """ + + def __init__(self, equation, scheme, io, spatial_methods=None, + physics_parametrisations=None): + """ + Args: + equation (:class:`PrognosticEquation`): the prognostic equation + scheme (:class:`TimeDiscretisation`): the scheme to use to timestep + the prognostic equation + io (:class:`IO`): the model's object for controlling input/output. + spatial_methods (iter, optional): a list of objects describing the + methods to use for discretising transport or diffusion terms + for each transported/diffused variable. Defaults to None, + in which case the terms follow the original discretisation in + the equation. + physics_parametrisations: (iter, optional): an iterable of + :class:`PhysicsParametrisation` objects that describe physical + parametrisations to be included to add to the equation. They can + only be used when the time discretisation `scheme` is explicit. + Defaults to None. + """ + self.scheme = scheme + if spatial_methods is not None: + self.spatial_methods = spatial_methods + else: + self.spatial_methods = [] + + if physics_parametrisations is not None: + self.physics_parametrisations = physics_parametrisations + if len(self.physics_parametrisations) > 1: + assert isinstance(scheme, ExplicitTimeDiscretisation), \ + ('Physics parametrisations can only be used with the ' + + 'basic TimeStepper when the time discretisation is ' + + 'explicit. If you want to use an implicit scheme, the ' + + 'SplitPhysicsTimestepper is more appropriate.') + else: + self.physics_parametrisations = [] + + super().__init__(equation=equation, io=io) + + @property + def transporting_velocity(self): + return "prognostic" + + def setup_fields(self): + self.x = TimeLevelFields(self.equation, self.scheme.nlevels) + self.fields = StateFields(self.x, self.equation.prescribed_fields, + *self.io.output.dumplist) + + def setup_scheme(self): + self.setup_equation(self.equation) + self.scheme.setup(self.equation) + self.setup_transporting_velocity(self.scheme) + if self.io.output.log_courant: + self.scheme.courant_max = self.io.courant_max + + def timestep(self): + """ + Implement the timestep + """ + xnp1 = self.x.np1 + name = self.equation.field_name + x_in = [x(name) for x in self.x.previous[-self.scheme.nlevels:]] + + self.scheme.apply(xnp1(name), *x_in) + + +class PrescribedTransport(Timestepper): + """ + Implements a timeloop with a prescibed transporting velocity. + """ + def __init__(self, equation, scheme, io, transport_method, + physics_parametrisations=None, + prescribed_transporting_velocity=None): + """ + Args: + equation (:class:`PrognosticEquation`): the prognostic equation + scheme (:class:`TimeDiscretisation`): the scheme to use to timestep + the prognostic equation. + transport_method (:class:`TransportMethod`): describes the method + used for discretising the transport term. + io (:class:`IO`): the model's object for controlling input/output. + physics_schemes: (list, optional): a list of :class:`Physics` and + :class:`TimeDiscretisation` options describing physical + parametrisations and timestepping schemes to use for each. + Timestepping schemes for physics must be explicit. Defaults to + None. + physics_parametrisations: (iter, optional): an iterable of + :class:`PhysicsParametrisation` objects that describe physical + parametrisations to be included to add to the equation. They can + only be used when the time discretisation `scheme` is explicit. + Defaults to None. + """ + + if isinstance(transport_method, TransportMethod): + transport_methods = [transport_method] + else: + # Assume an iterable has been provided + transport_methods = transport_method + + super().__init__(equation, scheme, io, spatial_methods=transport_methods, + physics_parametrisations=physics_parametrisations) + + if prescribed_transporting_velocity is not None: + self.velocity_projection = Projector( + prescribed_transporting_velocity(self.t), + self.fields('u')) + else: + self.velocity_projection = None + + @property + def transporting_velocity(self): + return self.fields('u') + + def setup_fields(self): + self.x = TimeLevelFields(self.equation, self.scheme.nlevels) + self.fields = StateFields(self.x, self.equation.prescribed_fields, + *self.io.output.dumplist) + + def run(self, t, tmax, pick_up=False): + """ + Runs the model for the specified time, from t to tmax + Args: + t (float): the start time of the run + tmax (float): the end time of the run + pick_up: (bool): specify whether to pick_up from a previous run + """ + # It's best to have evaluated the velocity before we start + if self.velocity_projection is not None: + self.velocity_projection.project() + + super().run(t, tmax, pick_up=pick_up) + + def timestep(self): + if self.velocity_projection is not None: + self.velocity_projection.project() + + super().timestep() diff --git a/integration-tests/equations/test_moist_compressible.py b/integration-tests/equations/test_moist_compressible.py index 44d53acf2..f2909f75e 100644 --- a/integration-tests/equations/test_moist_compressible.py +++ b/integration-tests/equations/test_moist_compressible.py @@ -5,7 +5,7 @@ from os.path import join, abspath, dirname from gusto import * -import gusto.thermodynamics as tde +import gusto.equations.thermodynamics as tde from firedrake import (SpatialCoordinate, PeriodicIntervalMesh, exp, sqrt, ExtrudedMesh, as_vector) import numpy as np diff --git a/integration-tests/equations/test_sw_fplane.py b/integration-tests/equations/test_sw_fplane.py index 65a237183..8e0779c2b 100644 --- a/integration-tests/equations/test_sw_fplane.py +++ b/integration-tests/equations/test_sw_fplane.py @@ -6,7 +6,7 @@ from os.path import join, abspath, dirname from gusto import * from firedrake import (PeriodicSquareMesh, SpatialCoordinate, Function, - cos, pi, as_vector) + cos, pi, as_vector, sin) import numpy as np diff --git a/integration-tests/model/test_IMEX.py b/integration-tests/model/test_IMEX.py index d4ce82c22..aab8b8514 100644 --- a/integration-tests/model/test_IMEX.py +++ b/integration-tests/model/test_IMEX.py @@ -22,13 +22,13 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup): eqn.label_terms(lambda t: t.has_label(transport), explicit) if scheme == "ssp3": - transport_scheme = SSP3(domain) + transport_scheme = IMEX_SSP3(domain) elif scheme == "ark2": - transport_scheme = ARK2(domain) + transport_scheme = IMEX_ARK2(domain) elif scheme == "ars3": - transport_scheme = ARS3(domain) + transport_scheme = IMEX_ARS3(domain) elif scheme == "trap2": - transport_scheme = Trap2(domain) + transport_scheme = IMEX_Trap2(domain) elif scheme == "euler": transport_scheme = IMEX_Euler(domain) diff --git a/integration-tests/physics/test_boundary_layer_mixing.py b/integration-tests/physics/test_boundary_layer_mixing.py index a72d0b3c1..cc76a7371 100644 --- a/integration-tests/physics/test_boundary_layer_mixing.py +++ b/integration-tests/physics/test_boundary_layer_mixing.py @@ -3,10 +3,10 @@ """ from gusto import * -from gusto.labels import physics_label +from gusto.core.labels import physics_label from firedrake import (VectorFunctionSpace, PeriodicIntervalMesh, as_vector, exp, SpatialCoordinate, ExtrudedMesh, Function) -from firedrake.fml import identity +from firedrake.fml import identity, drop import pytest diff --git a/integration-tests/physics/test_saturation_adjustment.py b/integration-tests/physics/test_saturation_adjustment.py index 3671114ea..74e6e129d 100644 --- a/integration-tests/physics/test_saturation_adjustment.py +++ b/integration-tests/physics/test_saturation_adjustment.py @@ -5,12 +5,12 @@ """ from os import path +import gusto.equations.thermodynamics as td from gusto import * -import gusto.thermodynamics as td from firedrake import (norm, Constant, PeriodicIntervalMesh, SpatialCoordinate, ExtrudedMesh, Function, sqrt, conditional) -from firedrake.fml import identity +from firedrake.fml import identity, drop from netCDF4 import Dataset import pytest diff --git a/integration-tests/physics/test_static_adjustment.py b/integration-tests/physics/test_static_adjustment.py index 2264d6675..0803d9b2c 100644 --- a/integration-tests/physics/test_static_adjustment.py +++ b/integration-tests/physics/test_static_adjustment.py @@ -5,10 +5,10 @@ """ from gusto import * -from gusto.labels import physics_label +from gusto.core.labels import physics_label from firedrake import (Constant, PeriodicIntervalMesh, SpatialCoordinate, ExtrudedMesh, Function) -from firedrake.fml import identity +from firedrake.fml import identity, drop import pytest diff --git a/integration-tests/physics/test_suppress_vertical_wind.py b/integration-tests/physics/test_suppress_vertical_wind.py index cab6efc76..8765b3985 100644 --- a/integration-tests/physics/test_suppress_vertical_wind.py +++ b/integration-tests/physics/test_suppress_vertical_wind.py @@ -4,10 +4,10 @@ """ from gusto import * -from gusto.labels import physics_label +from gusto.core.labels import physics_label from firedrake import (Constant, PeriodicIntervalMesh, as_vector, sin, norm, - SpatialCoordinate, ExtrudedMesh, Function, dot) -from firedrake.fml import identity + SpatialCoordinate, ExtrudedMesh, Function, dot, pi) +from firedrake.fml import identity, drop def run_suppress_vertical_wind(dirname): diff --git a/integration-tests/physics/test_surface_fluxes.py b/integration-tests/physics/test_surface_fluxes.py index fd4f92838..35d8e9b89 100644 --- a/integration-tests/physics/test_surface_fluxes.py +++ b/integration-tests/physics/test_surface_fluxes.py @@ -5,11 +5,11 @@ """ from gusto import * -import gusto.thermodynamics as td -from gusto.labels import physics_label +import gusto.equations.thermodynamics as td +from gusto.core.labels import physics_label from firedrake import (norm, Constant, PeriodicIntervalMesh, as_vector, SpatialCoordinate, ExtrudedMesh, Function, conditional) -from firedrake.fml import identity +from firedrake.fml import identity, drop import pytest diff --git a/integration-tests/physics/test_sw_saturation_adjustment.py b/integration-tests/physics/test_sw_saturation_adjustment.py index a758744c4..3bec4bb0f 100644 --- a/integration-tests/physics/test_sw_saturation_adjustment.py +++ b/integration-tests/physics/test_sw_saturation_adjustment.py @@ -12,7 +12,7 @@ from os import path from gusto import * from firedrake import (IcosahedralSphereMesh, acos, sin, cos, Constant, norm, - max_value, min_value) + max_value, min_value, pi, conditional) from netCDF4 import Dataset import pytest diff --git a/integration-tests/physics/test_wind_drag.py b/integration-tests/physics/test_wind_drag.py index c6aec9bc4..cc70ef356 100644 --- a/integration-tests/physics/test_wind_drag.py +++ b/integration-tests/physics/test_wind_drag.py @@ -3,11 +3,11 @@ """ from gusto import * -import gusto.thermodynamics as td -from gusto.labels import physics_label +import gusto.equations.thermodynamics as td +from gusto.core.labels import physics_label from firedrake import (norm, Constant, PeriodicIntervalMesh, as_vector, dot, SpatialCoordinate, ExtrudedMesh, Function, conditional) -from firedrake.fml import identity +from firedrake.fml import identity, drop import pytest diff --git a/server.py b/server.py deleted file mode 100644 index 81848fb8a..000000000 --- a/server.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Launch a livereload server serving up the html documention. Watch the -sphinx source directory for changes and rebuild the html documentation. Watch -the pyop2 package directory for changes and rebuild the API documentation. - -Requires livereload_ (or falls back to SimpleHTTPServer) :: - - pip install git+https://github.com/lepture/python-livereload - -.. _livereload: https://github.com/lepture/python-livereload""" - -try: - from livereload import Server - - server = Server() - server.watch('source', 'make buildhtml') - server.watch('../firedrake', 'make apidoc') - server.serve(root='build/html', open_url=True) -except ImportError: - import SimpleHTTPServer - import SocketServer - - Handler = SimpleHTTPServer.SimpleHTTPRequestHandler - httpd = SocketServer.TCPServer(("build/html", 8000), Handler) - httpd.serve_forever() diff --git a/setup.py b/setup.py index 3918a6b99..8f7239293 100644 --- a/setup.py +++ b/setup.py @@ -1,10 +1,20 @@ #!/usr/bin/env python -from distutils.core import setup +from setuptools import setup setup(name="gusto", version="1.0", description="Toolkit for compatible finite element dynamical cores", author="The Gusto Team", url="http://www.firedrakeproject.org/gusto/", - packages=["gusto", "gusto.recovery"]) + packages=["gusto", + "gusto.core", + "gusto.diagnostics", + "gusto.equations", + "gusto.initialisation", + "gusto.physics", + "gusto.recovery", + "gusto.solvers", + "gusto.spatial_methods", + "gusto.time_discretisation", + "gusto.timestepping"]) diff --git a/unit-tests/test_distances.py b/unit-tests/test_distances.py index e7895314a..4186be27f 100644 --- a/unit-tests/test_distances.py +++ b/unit-tests/test_distances.py @@ -4,7 +4,7 @@ import importlib import numpy as np import firedrake as fd -from gusto.coord_transforms import * +from gusto.core.coord_transforms import * import pytest tol = 1e-12 diff --git a/unit-tests/test_rotated_coords.py b/unit-tests/test_rotated_coords.py index aa01b7f47..5f03503ec 100644 --- a/unit-tests/test_rotated_coords.py +++ b/unit-tests/test_rotated_coords.py @@ -2,7 +2,7 @@ Test the formulae for rotating spherical coordinates. """ import numpy as np -from gusto.coord_transforms import * +from gusto.core.coord_transforms import * tol = 1e-12 diff --git a/unit-tests/test_rotated_vectors.py b/unit-tests/test_rotated_vectors.py index e326f0938..580a5e5f9 100644 --- a/unit-tests/test_rotated_vectors.py +++ b/unit-tests/test_rotated_vectors.py @@ -2,7 +2,7 @@ Test the formulae for rotating spherical vectors. """ import numpy as np -from gusto.coord_transforms import * +from gusto.core.coord_transforms import * tol = 1e-12 diff --git a/unit-tests/test_spherical_coord_transforms.py b/unit-tests/test_spherical_coord_transforms.py index 486c950e2..e7546d9a1 100644 --- a/unit-tests/test_spherical_coord_transforms.py +++ b/unit-tests/test_spherical_coord_transforms.py @@ -4,7 +4,7 @@ import importlib import numpy as np import firedrake as fd -from gusto.coord_transforms import * +from gusto.core.coord_transforms import * import pytest tol = 1e-12 diff --git a/unit-tests/test_spherical_rotations.py b/unit-tests/test_spherical_rotations.py index 07063d457..e235ea24c 100644 --- a/unit-tests/test_spherical_rotations.py +++ b/unit-tests/test_spherical_rotations.py @@ -4,7 +4,7 @@ import importlib import numpy as np import firedrake as fd -from gusto.coord_transforms import * +from gusto.core.coord_transforms import * import pytest tol = 1e-12 diff --git a/unit-tests/test_spherical_vector_transforms.py b/unit-tests/test_spherical_vector_transforms.py index b309db1fa..d78c0e429 100644 --- a/unit-tests/test_spherical_vector_transforms.py +++ b/unit-tests/test_spherical_vector_transforms.py @@ -6,7 +6,7 @@ import importlib import numpy as np import firedrake as fd -from gusto.coord_transforms import * +from gusto.core.coord_transforms import * import pytest tol = 1e-12