From abdaf635ec88bb2de229742c287a840429ae56a1 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Tue, 30 Aug 2022 16:01:01 +0100 Subject: [PATCH 01/12] make new files for domain and io but so far just copy state --- gusto/__init__.py | 2 + gusto/domain.py | 583 ++++++++++++++++++++++++++++++++++++++ gusto/io.py | 695 ++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 1280 insertions(+) create mode 100644 gusto/domain.py create mode 100644 gusto/io.py diff --git a/gusto/__init__.py b/gusto/__init__.py index c85012ef0..cc74c9d6b 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -1,11 +1,13 @@ from gusto.active_tracers import * # noqa from gusto.configuration import * # noqa +from gusto.domain import * # noqa from gusto.diagnostics import * # noqa from gusto.diffusion import * # noqa from gusto.equations import * # noqa from gusto.fml 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 diff --git a/gusto/domain.py b/gusto/domain.py new file mode 100644 index 000000000..bf3a04b48 --- /dev/null +++ b/gusto/domain.py @@ -0,0 +1,583 @@ +from os import path, makedirs +import itertools +from netCDF4 import Dataset +import sys +import time +from gusto.diagnostics import Diagnostics, Perturbation, SteadyStateError +from firedrake import (FiniteElement, TensorProductElement, HDiv, + FunctionSpace, VectorFunctionSpace, + interval, Function, Mesh, functionspaceimpl, + File, SpatialCoordinate, sqrt, Constant, inner, + op2, DumbCheckpoint, FILE_CREATE, FILE_READ, interpolate, + CellNormal, cross, as_vector) +import numpy as np +from gusto.configuration import logger, set_log_handler + +__all__ = ["Domain"] + + +class SpaceCreator(object): + + def __init__(self, mesh): + self.mesh = mesh + self.extruded_mesh = hasattr(mesh, "_base_mesh") + self._initialised_base_spaces = False + + def __call__(self, name, family=None, degree=None, V=None): + try: + return getattr(self, name) + except AttributeError: + if V is not None: + value = V + elif name == "HDiv" and family in ["BDM", "RT", "CG", "RTCF"]: + value = self.build_hdiv_space(family, degree) + elif name == "theta": + value = self.build_theta_space(degree) + elif name == "DG1_equispaced": + value = self.build_dg_space(1, variant='equispaced') + elif family == "DG": + value = self.build_dg_space(degree) + elif family == "CG": + value = self.build_cg_space(degree) + else: + raise ValueError(f'State has no space corresponding to {name}') + setattr(self, name, value) + return value + + def build_compatible_spaces(self, family, degree): + if self.extruded_mesh and not self._initialised_base_spaces: + self.build_base_spaces(family, degree) + Vu = self.build_hdiv_space(family, degree) + setattr(self, "HDiv", Vu) + Vdg = self.build_dg_space(degree) + setattr(self, "DG", Vdg) + Vth = self.build_theta_space(degree) + setattr(self, "theta", Vth) + return Vu, Vdg, Vth + else: + Vu = self.build_hdiv_space(family, degree) + setattr(self, "HDiv", Vu) + Vdg = self.build_dg_space(degree) + setattr(self, "DG", Vdg) + return Vu, Vdg + + def build_base_spaces(self, family, degree): + + cell = self.mesh._base_mesh.ufl_cell().cellname() + + # horizontal base spaces + self.S1 = FiniteElement(family, cell, degree+1) + self.S2 = FiniteElement("DG", cell, degree) + + # vertical base spaces + self.T0 = FiniteElement("CG", interval, degree+1) + self.T1 = FiniteElement("DG", interval, degree) + + self._initialised_base_spaces = True + + def build_hdiv_space(self, family, degree): + if self.extruded_mesh: + if not self._initialised_base_spaces: + self.build_base_spaces(family, degree) + Vh_elt = HDiv(TensorProductElement(self.S1, self.T1)) + Vt_elt = TensorProductElement(self.S2, self.T0) + Vv_elt = HDiv(Vt_elt) + V_elt = Vh_elt + Vv_elt + else: + cell = self.mesh.ufl_cell().cellname() + V_elt = FiniteElement(family, cell, degree+1) + return FunctionSpace(self.mesh, V_elt, name='HDiv') + + def build_dg_space(self, degree, variant=None): + if self.extruded_mesh: + if not self._initialised_base_spaces or self.T1.degree() != degree or self.T1.variant() != variant: + cell = self.mesh._base_mesh.ufl_cell().cellname() + S2 = FiniteElement("DG", cell, degree, variant=variant) + T1 = FiniteElement("DG", interval, degree, variant=variant) + else: + S2 = self.S2 + T1 = self.T1 + V_elt = TensorProductElement(S2, T1) + else: + cell = self.mesh.ufl_cell().cellname() + V_elt = FiniteElement("DG", cell, degree, variant=variant) + name = f'DG{degree}_equispaced' if variant == 'equispaced' else f'DG{degree}' + return FunctionSpace(self.mesh, V_elt, name=name) + + def build_theta_space(self, degree): + assert self.extruded_mesh + if not self._initialised_base_spaces: + cell = self.mesh._base_mesh.ufl_cell().cellname() + self.S2 = FiniteElement("DG", cell, degree) + self.T0 = FiniteElement("CG", interval, degree+1) + V_elt = TensorProductElement(self.S2, self.T0) + return FunctionSpace(self.mesh, V_elt, name='Vtheta') + + def build_cg_space(self, degree): + return FunctionSpace(self.mesh, "CG", degree, name=f'CG{degree}') + + +class FieldCreator(object): + + def __init__(self, equations): + self.fields = [] + for eqn in equations: + subfield_names = eqn.field_names if hasattr(eqn, "field_names") else None + self.add_field(eqn.field_name, eqn.function_space, subfield_names) + + def add_field(self, name, space, subfield_names=None): + value = Function(space, name=name) + setattr(self, name, value) + self.fields.append(value) + + if len(space) > 1: + assert len(space) == len(subfield_names) + for field_name, field in zip(subfield_names, value.split()): + setattr(self, field_name, field) + field.rename(field_name) + self.fields.append(field) + + def __call__(self, name): + return getattr(self, name) + + def __iter__(self): + return iter(self.fields) + + +class StateFields(FieldCreator): + + def __init__(self, *fields_to_dump): + self.fields = [] + self.output_specified = len(fields_to_dump) > 0 + self.to_dump = set((fields_to_dump)) + self.to_pickup = set(()) + + def __call__(self, name, space=None, subfield_names=None, dump=True, + pickup=False): + try: + return getattr(self, name) + except AttributeError: + self.add_field(name, space, subfield_names) + if dump: + if subfield_names is not None: + self.to_dump.update(subfield_names) + else: + self.to_dump.add(name) + if pickup: + if subfield_names is not None: + self.to_pickup.update(subfield_names) + else: + self.to_pickup.add(name) + return getattr(self, name) + + +class Domain(object): + """ + An object holding the mesh used by the model and its function spaces. + + The :class:`Domain` holds the mesh and builds the compatible finite element + spaces upon that mesh. + + Args: + mesh (:class:`Mesh`): The mesh to use. + family (str): The family of HDiv finite elements to use. With the + degree, this defines the de Rham complex. + degree (int, optional): The polynomial degree used by the de Rham + complex. Specifically this is the degree of the discontinuous space + at the bottom of the complex. Defaults to None, but if None is + provided then then horizontal degree (and possibly vertical degree) + must be specified. + horizontal_degree (int, optional): The polynomial degree used by the + horizontal component of a Tensor Product de Rham complex. Can only + be used with extruded meshes. Defaults to None, and + vertical_degree (int, optional): + + """ + + def __init__(self, mesh, family, degree=None, + horizontal_degree=None, vertical_degree=None): + + # Check degree arguments do not conflict with one another + if degree is not None and (horizontal_degree is not None + or vertical_degree is not None): + raise ValueError( + 'If degree argument is provided, cannot also provide ' \ + + 'horizontal_degree or vertical_degree argument') + + + if horizontal_degree is None: + if degree is None: + raise ValueError('The degree or horizontal_degree must be specified') + # Set horizontal degree to degree + horizontal_degree = degree + + # For extruded meshes, set vertical degree + if mesh.extruded: + if vertical_degree is None: + if degree is None: + raise ValueError('The degree or vertical_degree must be specified for an extruded mesh') + # Set vertical degree to degree + vertical_degree = degree + elif vertical_degree is not None: + raise ValueError('Cannot provide a vertical_degree for a non-extruded mesh') + + self.mesh = mesh + self.spaces = SpaceCreator(mesh) + + if self.output.dumplist is None: + + self.output.dumplist = [] + + self.fields = StateFields(*self.output.dumplist) + + self.dumpdir = None + self.dumpfile = None + self.to_pickup = None + + # figure out if we're on a sphere + try: + self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3 and mesh._base_mesh.topological_dimension() == 2) + except AttributeError: + self.on_sphere = (mesh.geometric_dimension() == 3 and mesh.topological_dimension() == 2) + + # build the vertical normal and define perp for 2d geometries + dim = mesh.topological_dimension() + if self.on_sphere: + x = SpatialCoordinate(mesh) + R = sqrt(inner(x, x)) + self.k = interpolate(x/R, mesh.coordinates.function_space()) + if dim == 2: + outward_normals = CellNormal(mesh) + self.perp = lambda u: cross(outward_normals, u) + else: + kvec = [0.0]*dim + kvec[dim-1] = 1.0 + self.k = Constant(kvec) + if dim == 2: + self.perp = lambda u: as_vector([-u[1], u[0]]) + + # setup logger + logger.setLevel(output.log_level) + set_log_handler(mesh.comm) + if parameters is not None: + logger.info("Physical parameters that take non-default values:") + logger.info(", ".join("%s: %s" % (k, float(v)) for (k, v) in vars(parameters).items())) + + # Constant to hold current time + self.t = Constant(0.0) + if type(dt) is Constant: + self.dt = dt + elif type(dt) in (float, int): + self.dt = Constant(dt) + else: + raise TypeError(f'dt must be a Constant, float or int, not {type(dt)}') + + def setup_diagnostics(self): + """ + Add special case diagnostic fields + """ + for name in self.output.perturbation_fields: + f = Perturbation(name) + self.diagnostic_fields.append(f) + + for name in self.output.steady_state_error_fields: + f = SteadyStateError(self, name) + self.diagnostic_fields.append(f) + + fields = set([f.name() for f in self.fields]) + field_deps = [(d, sorted(set(d.required_fields).difference(fields),)) for d in self.diagnostic_fields] + schedule = topo_sort(field_deps) + self.diagnostic_fields = schedule + for diagnostic in self.diagnostic_fields: + diagnostic.setup(self) + self.diagnostics.register(diagnostic.name) + + def setup_dump(self, t, tmax, pickup=False): + """ + Setup dump files + Check for existence of directory so as not to overwrite + output files + Setup checkpoint file + + :arg tmax: model stop time + :arg pickup: recover state from the checkpointing file if true, + otherwise dump and checkpoint to disk. (default is False). + """ + + if any([self.output.dump_vtus, self.output.dumplist_latlon, + self.output.dump_diagnostics, self.output.point_data, + self.output.checkpoint and not pickup]): + # setup output directory and check that it does not already exist + self.dumpdir = path.join("results", self.output.dirname) + running_tests = '--running-tests' in sys.argv or "pytest" in self.output.dirname + if self.mesh.comm.rank == 0: + if not running_tests and path.exists(self.dumpdir) and not pickup: + raise IOError("results directory '%s' already exists" + % self.dumpdir) + else: + if not running_tests: + makedirs(self.dumpdir) + + if self.output.dump_vtus: + + # setup pvd output file + outfile = path.join(self.dumpdir, "field_output.pvd") + self.dumpfile = File( + outfile, project_output=self.output.project_fields, + comm=self.mesh.comm) + + # make list of fields to dump + self.to_dump = [f for f in self.fields if f.name() in self.fields.to_dump] + + # make dump counter + self.dumpcount = itertools.count() + + # if there are fields to be dumped in latlon coordinates, + # setup the latlon coordinate mesh and make output file + if len(self.output.dumplist_latlon) > 0: + mesh_ll = get_latlon_mesh(self.mesh) + outfile_ll = path.join(self.dumpdir, "field_output_latlon.pvd") + self.dumpfile_ll = File(outfile_ll, + project_output=self.output.project_fields, + comm=self.mesh.comm) + + # make functions on latlon mesh, as specified by dumplist_latlon + self.to_dump_latlon = [] + for name in self.output.dumplist_latlon: + f = self.fields(name) + field = Function( + functionspaceimpl.WithGeometry.create( + f.function_space(), mesh_ll), + val=f.topological, name=name+'_ll') + self.to_dump_latlon.append(field) + + # we create new netcdf files to write to, unless pickup=True, in + # which case we just need the filenames + if self.output.dump_diagnostics: + diagnostics_filename = self.dumpdir+"/diagnostics.nc" + self.diagnostic_output = DiagnosticsOutput(diagnostics_filename, + self.diagnostics, + self.output.dirname, + self.mesh.comm, + create=not pickup) + + if len(self.output.point_data) > 0: + # set up point data output + pointdata_filename = self.dumpdir+"/point_data.nc" + ndt = int(tmax/float(self.dt)) + self.pointdata_output = PointDataOutput(pointdata_filename, ndt, + self.output.point_data, + self.output.dirname, + self.fields, + self.mesh.comm, + self.output.tolerance, + create=not pickup) + + # make point data dump counter + self.pddumpcount = itertools.count() + + # set frequency of point data output - defaults to + # dumpfreq if not set by user + if self.output.pddumpfreq is None: + self.output.pddumpfreq = self.output.dumpfreq + + # if we want to checkpoint and are not picking up from a previous + # checkpoint file, setup the checkpointing + if self.output.checkpoint: + if not pickup: + self.chkpt = DumbCheckpoint(path.join(self.dumpdir, "chkpt"), + mode=FILE_CREATE) + # make list of fields to pickup (this doesn't include + # diagnostic fields) + self.to_pickup = [f for f in self.fields if f.name() in self.fields.to_pickup] + + # if we want to checkpoint then make a checkpoint counter + if self.output.checkpoint: + self.chkptcount = itertools.count() + + # dump initial fields + self.dump(t) + + def pickup_from_checkpoint(self): + """ + :arg t: the current model time (default is zero). + """ + # TODO: this duplicates some code from setup_dump. Can this be avoided? + # It is because we don't know if we are picking up or setting dump first + if self.to_pickup is None: + self.to_pickup = [f for f in self.fields if f.name() in self.fields.to_pickup] + # Set dumpdir if has not been done already + if self.dumpdir is None: + self.dumpdir = path.join("results", self.output.dirname) + + if self.output.checkpoint: + # Open the checkpointing file for writing + if self.output.checkpoint_pickup_filename is not None: + chkfile = self.output.checkpoint_pickup_filename + else: + chkfile = path.join(self.dumpdir, "chkpt") + with DumbCheckpoint(chkfile, mode=FILE_READ) as chk: + # Recover all the fields from the checkpoint + for field in self.to_pickup: + chk.load(field) + t = chk.read_attribute("/", "time") + # Setup new checkpoint + self.chkpt = DumbCheckpoint(path.join(self.dumpdir, "chkpt"), mode=FILE_CREATE) + else: + raise ValueError("Must set checkpoint True if pickup") + + return t + + def dump(self, t): + """ + Dump output + """ + output = self.output + + # Diagnostics: + # Compute diagnostic fields + for field in self.diagnostic_fields: + field(self) + + if output.dump_diagnostics: + # Output diagnostic data + self.diagnostic_output.dump(self, t) + + if len(output.point_data) > 0 and (next(self.pddumpcount) % output.pddumpfreq) == 0: + # Output pointwise data + self.pointdata_output.dump(self.fields, t) + + # Dump all the fields to the checkpointing file (backup version) + if output.checkpoint and (next(self.chkptcount) % output.chkptfreq) == 0: + for field in self.to_pickup: + self.chkpt.store(field) + self.chkpt.write_attribute("/", "time", t) + + if output.dump_vtus and (next(self.dumpcount) % output.dumpfreq) == 0: + # dump fields + self.dumpfile.write(*self.to_dump) + + # dump fields on latlon mesh + if len(output.dumplist_latlon) > 0: + self.dumpfile_ll.write(*self.to_dump_latlon) + + def initialise(self, initial_conditions): + """ + Initialise state variables + + :arg initial_conditions: An iterable of pairs (field_name, pointwise_value) + """ + for name, ic in initial_conditions: + f_init = getattr(self.fields, name) + f_init.assign(ic) + f_init.rename(name) + + def set_reference_profiles(self, reference_profiles): + """ + Initialise reference profiles + + :arg reference_profiles: An iterable of pairs (field_name, interpolatory_value) + """ + for name, profile in reference_profiles: + if name+'bar' in self.fields: + # For reference profiles already added to state, allow + # interpolation from expressions + ref = self.fields(name+'bar') + elif isinstance(profile, Function): + # Need to add reference profile to state so profile must be + # a Function + ref = self.fields(name+'bar', space=profile.function_space(), dump=False) + else: + raise ValueError(f'When initialising reference profile {name}' + + ' the passed profile must be a Function') + ref.interpolate(profile) + + +def get_latlon_mesh(mesh): + coords_orig = mesh.coordinates + coords_fs = coords_orig.function_space() + + if coords_fs.extruded: + cell = mesh._base_mesh.ufl_cell().cellname() + DG1_hori_elt = FiniteElement("DG", cell, 1, variant="equispaced") + DG1_vert_elt = FiniteElement("DG", interval, 1, variant="equispaced") + DG1_elt = TensorProductElement(DG1_hori_elt, DG1_vert_elt) + else: + cell = mesh.ufl_cell().cellname() + DG1_elt = FiniteElement("DG", cell, 1, variant="equispaced") + vec_DG1 = VectorFunctionSpace(mesh, DG1_elt) + coords_dg = Function(vec_DG1).interpolate(coords_orig) + coords_latlon = Function(vec_DG1) + shapes = {"nDOFs": vec_DG1.finat_element.space_dimension(), 'dim': 3} + + radius = np.min(np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2)) + # lat-lon 'x' = atan2(y, x) + coords_latlon.dat.data[:, 0] = np.arctan2(coords_dg.dat.data[:, 1], coords_dg.dat.data[:, 0]) + # lat-lon 'y' = asin(z/sqrt(x^2 + y^2 + z^2)) + coords_latlon.dat.data[:, 1] = np.arcsin(coords_dg.dat.data[:, 2]/np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2)) + # our vertical coordinate is radius - the minimum radius + coords_latlon.dat.data[:, 2] = np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2) - radius + +# We need to ensure that all points in a cell are on the same side of the branch cut in longitude coords +# This kernel amends the longitude coords so that all longitudes in one cell are close together + kernel = op2.Kernel(""" +#define PI 3.141592653589793 +#define TWO_PI 6.283185307179586 +void splat_coords(double *coords) {{ + double max_diff = 0.0; + double diff = 0.0; + + for (int i=0; i<{nDOFs}; i++) {{ + for (int j=0; j<{nDOFs}; j++) {{ + diff = coords[i*{dim}] - coords[j*{dim}]; + if (fabs(diff) > max_diff) {{ + max_diff = diff; + }} + }} + }} + + if (max_diff > PI) {{ + for (int i=0; i<{nDOFs}; i++) {{ + if (coords[i*{dim}] < 0) {{ + coords[i*{dim}] += TWO_PI; + }} + }} + }} +}} +""".format(**shapes), "splat_coords") + + op2.par_loop(kernel, coords_latlon.cell_set, + coords_latlon.dat(op2.RW, coords_latlon.cell_node_map())) + return Mesh(coords_latlon) + + +def topo_sort(field_deps): + name2field = dict((f.name, f) for f, _ in field_deps) + # map node: (input_deps, output_deps) + graph = dict((f.name, (list(deps), [])) for f, deps in field_deps) + roots = [] + for f, input_deps in field_deps: + if len(input_deps) == 0: + # No dependencies, candidate for evaluation + roots.append(f.name) + for d in input_deps: + # add f as output dependency + graph[d][1].append(f.name) + + schedule = [] + while roots: + n = roots.pop() + schedule.append(n) + output_deps = list(graph[n][1]) + for m in output_deps: + # Remove edge + graph[m][0].remove(n) + graph[n][1].remove(m) + # If m now as no input deps, candidate for evaluation + if len(graph[m][0]) == 0: + roots.append(m) + if any(len(i) for i, _ in graph.values()): + cycle = "\n".join("%s -> %s" % (f, i) for f, (i, _) in graph.items() + if f not in schedule) + raise RuntimeError("Field dependencies have a cycle:\n\n%s" % cycle) + return list(map(name2field.__getitem__, schedule)) diff --git a/gusto/io.py b/gusto/io.py new file mode 100644 index 000000000..25e26282d --- /dev/null +++ b/gusto/io.py @@ -0,0 +1,695 @@ +from os import path, makedirs +import itertools +from netCDF4 import Dataset +import sys +import time +from gusto.diagnostics import Diagnostics, Perturbation, SteadyStateError +from firedrake import (FiniteElement, TensorProductElement, HDiv, + FunctionSpace, VectorFunctionSpace, + interval, Function, Mesh, functionspaceimpl, + File, SpatialCoordinate, sqrt, Constant, inner, + op2, DumbCheckpoint, FILE_CREATE, FILE_READ, interpolate, + CellNormal, cross, as_vector) +import numpy as np +from gusto.configuration import logger, set_log_handler + +__all__ = ["IO"] + + +class SpaceCreator(object): + + def __init__(self, mesh): + self.mesh = mesh + self.extruded_mesh = hasattr(mesh, "_base_mesh") + self._initialised_base_spaces = False + + def __call__(self, name, family=None, degree=None, V=None): + try: + return getattr(self, name) + except AttributeError: + if V is not None: + value = V + elif name == "HDiv" and family in ["BDM", "RT", "CG", "RTCF"]: + value = self.build_hdiv_space(family, degree) + elif name == "theta": + value = self.build_theta_space(degree) + elif name == "DG1_equispaced": + value = self.build_dg_space(1, variant='equispaced') + elif family == "DG": + value = self.build_dg_space(degree) + elif family == "CG": + value = self.build_cg_space(degree) + else: + raise ValueError(f'State has no space corresponding to {name}') + setattr(self, name, value) + return value + + def build_compatible_spaces(self, family, degree): + if self.extruded_mesh and not self._initialised_base_spaces: + self.build_base_spaces(family, degree) + Vu = self.build_hdiv_space(family, degree) + setattr(self, "HDiv", Vu) + Vdg = self.build_dg_space(degree) + setattr(self, "DG", Vdg) + Vth = self.build_theta_space(degree) + setattr(self, "theta", Vth) + return Vu, Vdg, Vth + else: + Vu = self.build_hdiv_space(family, degree) + setattr(self, "HDiv", Vu) + Vdg = self.build_dg_space(degree) + setattr(self, "DG", Vdg) + return Vu, Vdg + + def build_base_spaces(self, family, degree): + + cell = self.mesh._base_mesh.ufl_cell().cellname() + + # horizontal base spaces + self.S1 = FiniteElement(family, cell, degree+1) + self.S2 = FiniteElement("DG", cell, degree) + + # vertical base spaces + self.T0 = FiniteElement("CG", interval, degree+1) + self.T1 = FiniteElement("DG", interval, degree) + + self._initialised_base_spaces = True + + def build_hdiv_space(self, family, degree): + if self.extruded_mesh: + if not self._initialised_base_spaces: + self.build_base_spaces(family, degree) + Vh_elt = HDiv(TensorProductElement(self.S1, self.T1)) + Vt_elt = TensorProductElement(self.S2, self.T0) + Vv_elt = HDiv(Vt_elt) + V_elt = Vh_elt + Vv_elt + else: + cell = self.mesh.ufl_cell().cellname() + V_elt = FiniteElement(family, cell, degree+1) + return FunctionSpace(self.mesh, V_elt, name='HDiv') + + def build_dg_space(self, degree, variant=None): + if self.extruded_mesh: + if not self._initialised_base_spaces or self.T1.degree() != degree or self.T1.variant() != variant: + cell = self.mesh._base_mesh.ufl_cell().cellname() + S2 = FiniteElement("DG", cell, degree, variant=variant) + T1 = FiniteElement("DG", interval, degree, variant=variant) + else: + S2 = self.S2 + T1 = self.T1 + V_elt = TensorProductElement(S2, T1) + else: + cell = self.mesh.ufl_cell().cellname() + V_elt = FiniteElement("DG", cell, degree, variant=variant) + name = f'DG{degree}_equispaced' if variant == 'equispaced' else f'DG{degree}' + return FunctionSpace(self.mesh, V_elt, name=name) + + def build_theta_space(self, degree): + assert self.extruded_mesh + if not self._initialised_base_spaces: + cell = self.mesh._base_mesh.ufl_cell().cellname() + self.S2 = FiniteElement("DG", cell, degree) + self.T0 = FiniteElement("CG", interval, degree+1) + V_elt = TensorProductElement(self.S2, self.T0) + return FunctionSpace(self.mesh, V_elt, name='Vtheta') + + def build_cg_space(self, degree): + return FunctionSpace(self.mesh, "CG", degree, name=f'CG{degree}') + + +class FieldCreator(object): + + def __init__(self, equations): + self.fields = [] + for eqn in equations: + subfield_names = eqn.field_names if hasattr(eqn, "field_names") else None + self.add_field(eqn.field_name, eqn.function_space, subfield_names) + + def add_field(self, name, space, subfield_names=None): + value = Function(space, name=name) + setattr(self, name, value) + self.fields.append(value) + + if len(space) > 1: + assert len(space) == len(subfield_names) + for field_name, field in zip(subfield_names, value.split()): + setattr(self, field_name, field) + field.rename(field_name) + self.fields.append(field) + + def __call__(self, name): + return getattr(self, name) + + def __iter__(self): + return iter(self.fields) + + +class StateFields(FieldCreator): + + def __init__(self, *fields_to_dump): + self.fields = [] + self.output_specified = len(fields_to_dump) > 0 + self.to_dump = set((fields_to_dump)) + self.to_pickup = set(()) + + def __call__(self, name, space=None, subfield_names=None, dump=True, + pickup=False): + try: + return getattr(self, name) + except AttributeError: + self.add_field(name, space, subfield_names) + if dump: + if subfield_names is not None: + self.to_dump.update(subfield_names) + else: + self.to_dump.add(name) + if pickup: + if subfield_names is not None: + self.to_pickup.update(subfield_names) + else: + self.to_pickup.add(name) + return getattr(self, name) + + +class PointDataOutput(object): + def __init__(self, filename, ndt, field_points, description, + field_creator, comm, tolerance=None, create=True): + """Create a dump file that stores fields evaluated at points. + + :arg filename: The filename. + :arg field_points: Iterable of pairs (field_name, evaluation_points). + :arg description: Description of the simulation. + :arg field_creator: The field creator (only used to determine + datatype and shape of fields). + :kwarg create: If False, assume that filename already exists + """ + # Overwrite on creation. + self.dump_count = 0 + self.filename = filename + self.field_points = field_points + self.tolerance = tolerance + self.comm = comm + if not create: + return + if self.comm.rank == 0: + with Dataset(filename, "w") as dataset: + dataset.description = "Point data for simulation {desc}".format(desc=description) + dataset.history = "Created {t}".format(t=time.ctime()) + # FIXME add versioning information. + dataset.source = "Output from Gusto model" + # Appendable dimension, timesteps in the model + dataset.createDimension("time", None) + + var = dataset.createVariable("time", np.float64, ("time")) + var.units = "seconds" + # Now create the variable group for each field + for field_name, points in field_points: + group = dataset.createGroup(field_name) + npts, dim = points.shape + group.createDimension("points", npts) + group.createDimension("geometric_dimension", dim) + var = group.createVariable("points", points.dtype, + ("points", "geometric_dimension")) + var[:] = points + + # Get the UFL shape of the field + field_shape = field_creator(field_name).ufl_shape + # Number of geometric dimension occurences should be the same as the length of the UFL shape + field_len = len(field_shape) + field_count = field_shape.count(dim) + assert field_len == field_count, "Geometric dimension occurrences do not match UFL shape" + # Create the variable with the required shape + dimensions = ("time", "points") + field_count*("geometric_dimension",) + group.createVariable(field_name, field_creator(field_name).dat.dtype, dimensions) + + def dump(self, field_creator, t): + """Evaluate and dump field data at points. + + :arg field_creator: :class:`FieldCreator` for accessing + fields. + :arg t: Simulation time at which dump occurs. + """ + + val_list = [] + for field_name, points in self.field_points: + val_list.append((field_name, np.asarray(field_creator(field_name).at(points, tolerance=self.tolerance)))) + + if self.comm.rank == 0: + with Dataset(self.filename, "a") as dataset: + # Add new time index + dataset.variables["time"][self.dump_count] = t + for field_name, vals in val_list: + group = dataset.groups[field_name] + var = group.variables[field_name] + var[self.dump_count, :] = vals + + self.dump_count += 1 + + +class DiagnosticsOutput(object): + def __init__(self, filename, diagnostics, description, comm, create=True): + """Create a dump file that stores diagnostics. + + :arg filename: The filename. + :arg diagnostics: The :class:`Diagnostics` object. + :arg description: A description. + :kwarg create: If False, assume that filename already exists + """ + self.filename = filename + self.diagnostics = diagnostics + self.comm = comm + if not create: + return + if self.comm.rank == 0: + with Dataset(filename, "w") as dataset: + dataset.description = "Diagnostics data for simulation {desc}".format(desc=description) + dataset.history = "Created {t}".format(t=time.ctime()) + dataset.source = "Output from Gusto model" + dataset.createDimension("time", None) + var = dataset.createVariable("time", np.float64, ("time", )) + var.units = "seconds" + for name in diagnostics.fields: + group = dataset.createGroup(name) + for diagnostic in diagnostics.available_diagnostics: + group.createVariable(diagnostic, np.float64, ("time", )) + + def dump(self, state, t): + """Dump diagnostics. + + :arg state: The :class:`State` at which to compute the diagnostic. + :arg t: The current time. + """ + + diagnostics = [] + for fname in self.diagnostics.fields: + field = state.fields(fname) + for dname in self.diagnostics.available_diagnostics: + diagnostic = getattr(self.diagnostics, dname) + diagnostics.append((fname, dname, diagnostic(field))) + + if self.comm.rank == 0: + with Dataset(self.filename, "a") as dataset: + idx = dataset.dimensions["time"].size + dataset.variables["time"][idx:idx + 1] = t + for fname, dname, value in diagnostics: + group = dataset.groups[fname] + var = group.variables[dname] + var[idx:idx + 1] = value + + +class State(object): + """ + Build a model state to keep the variables in, and specify parameters. + + :arg mesh: The :class:`Mesh` to use. + :arg dt: The time step as a :class:`Constant`. If a float or int is passed, + it will be cast to a :class:`Constant`. + :arg output: class containing output parameters + :arg parameters: class containing physical parameters + :arg diagnostics: class containing diagnostic methods + :arg diagnostic_fields: list of diagnostic field classes + """ + + def __init__(self, mesh, dt, + output=None, + parameters=None, + diagnostics=None, + diagnostic_fields=None): + + if output is None: + raise RuntimeError("You must provide a directory name for dumping results") + else: + self.output = output + self.parameters = parameters + + if diagnostics is not None: + self.diagnostics = diagnostics + else: + self.diagnostics = Diagnostics() + if diagnostic_fields is not None: + self.diagnostic_fields = diagnostic_fields + else: + self.diagnostic_fields = [] + + # The mesh + self.mesh = mesh + + self.spaces = SpaceCreator(mesh) + + if self.output.dumplist is None: + + self.output.dumplist = [] + + self.fields = StateFields(*self.output.dumplist) + + self.dumpdir = None + self.dumpfile = None + self.to_pickup = None + + # figure out if we're on a sphere + try: + self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3 and mesh._base_mesh.topological_dimension() == 2) + except AttributeError: + self.on_sphere = (mesh.geometric_dimension() == 3 and mesh.topological_dimension() == 2) + + # build the vertical normal and define perp for 2d geometries + dim = mesh.topological_dimension() + if self.on_sphere: + x = SpatialCoordinate(mesh) + R = sqrt(inner(x, x)) + self.k = interpolate(x/R, mesh.coordinates.function_space()) + if dim == 2: + outward_normals = CellNormal(mesh) + self.perp = lambda u: cross(outward_normals, u) + else: + kvec = [0.0]*dim + kvec[dim-1] = 1.0 + self.k = Constant(kvec) + if dim == 2: + self.perp = lambda u: as_vector([-u[1], u[0]]) + + # setup logger + logger.setLevel(output.log_level) + set_log_handler(mesh.comm) + if parameters is not None: + logger.info("Physical parameters that take non-default values:") + logger.info(", ".join("%s: %s" % (k, float(v)) for (k, v) in vars(parameters).items())) + + # Constant to hold current time + self.t = Constant(0.0) + if type(dt) is Constant: + self.dt = dt + elif type(dt) in (float, int): + self.dt = Constant(dt) + else: + raise TypeError(f'dt must be a Constant, float or int, not {type(dt)}') + + def setup_diagnostics(self): + """ + Add special case diagnostic fields + """ + for name in self.output.perturbation_fields: + f = Perturbation(name) + self.diagnostic_fields.append(f) + + for name in self.output.steady_state_error_fields: + f = SteadyStateError(self, name) + self.diagnostic_fields.append(f) + + fields = set([f.name() for f in self.fields]) + field_deps = [(d, sorted(set(d.required_fields).difference(fields),)) for d in self.diagnostic_fields] + schedule = topo_sort(field_deps) + self.diagnostic_fields = schedule + for diagnostic in self.diagnostic_fields: + diagnostic.setup(self) + self.diagnostics.register(diagnostic.name) + + def setup_dump(self, t, tmax, pickup=False): + """ + Setup dump files + Check for existence of directory so as not to overwrite + output files + Setup checkpoint file + + :arg tmax: model stop time + :arg pickup: recover state from the checkpointing file if true, + otherwise dump and checkpoint to disk. (default is False). + """ + + if any([self.output.dump_vtus, self.output.dumplist_latlon, + self.output.dump_diagnostics, self.output.point_data, + self.output.checkpoint and not pickup]): + # setup output directory and check that it does not already exist + self.dumpdir = path.join("results", self.output.dirname) + running_tests = '--running-tests' in sys.argv or "pytest" in self.output.dirname + if self.mesh.comm.rank == 0: + if not running_tests and path.exists(self.dumpdir) and not pickup: + raise IOError("results directory '%s' already exists" + % self.dumpdir) + else: + if not running_tests: + makedirs(self.dumpdir) + + if self.output.dump_vtus: + + # setup pvd output file + outfile = path.join(self.dumpdir, "field_output.pvd") + self.dumpfile = File( + outfile, project_output=self.output.project_fields, + comm=self.mesh.comm) + + # make list of fields to dump + self.to_dump = [f for f in self.fields if f.name() in self.fields.to_dump] + + # make dump counter + self.dumpcount = itertools.count() + + # if there are fields to be dumped in latlon coordinates, + # setup the latlon coordinate mesh and make output file + if len(self.output.dumplist_latlon) > 0: + mesh_ll = get_latlon_mesh(self.mesh) + outfile_ll = path.join(self.dumpdir, "field_output_latlon.pvd") + self.dumpfile_ll = File(outfile_ll, + project_output=self.output.project_fields, + comm=self.mesh.comm) + + # make functions on latlon mesh, as specified by dumplist_latlon + self.to_dump_latlon = [] + for name in self.output.dumplist_latlon: + f = self.fields(name) + field = Function( + functionspaceimpl.WithGeometry.create( + f.function_space(), mesh_ll), + val=f.topological, name=name+'_ll') + self.to_dump_latlon.append(field) + + # we create new netcdf files to write to, unless pickup=True, in + # which case we just need the filenames + if self.output.dump_diagnostics: + diagnostics_filename = self.dumpdir+"/diagnostics.nc" + self.diagnostic_output = DiagnosticsOutput(diagnostics_filename, + self.diagnostics, + self.output.dirname, + self.mesh.comm, + create=not pickup) + + if len(self.output.point_data) > 0: + # set up point data output + pointdata_filename = self.dumpdir+"/point_data.nc" + ndt = int(tmax/float(self.dt)) + self.pointdata_output = PointDataOutput(pointdata_filename, ndt, + self.output.point_data, + self.output.dirname, + self.fields, + self.mesh.comm, + self.output.tolerance, + create=not pickup) + + # make point data dump counter + self.pddumpcount = itertools.count() + + # set frequency of point data output - defaults to + # dumpfreq if not set by user + if self.output.pddumpfreq is None: + self.output.pddumpfreq = self.output.dumpfreq + + # if we want to checkpoint and are not picking up from a previous + # checkpoint file, setup the checkpointing + if self.output.checkpoint: + if not pickup: + self.chkpt = DumbCheckpoint(path.join(self.dumpdir, "chkpt"), + mode=FILE_CREATE) + # make list of fields to pickup (this doesn't include + # diagnostic fields) + self.to_pickup = [f for f in self.fields if f.name() in self.fields.to_pickup] + + # if we want to checkpoint then make a checkpoint counter + if self.output.checkpoint: + self.chkptcount = itertools.count() + + # dump initial fields + self.dump(t) + + def pickup_from_checkpoint(self): + """ + :arg t: the current model time (default is zero). + """ + # TODO: this duplicates some code from setup_dump. Can this be avoided? + # It is because we don't know if we are picking up or setting dump first + if self.to_pickup is None: + self.to_pickup = [f for f in self.fields if f.name() in self.fields.to_pickup] + # Set dumpdir if has not been done already + if self.dumpdir is None: + self.dumpdir = path.join("results", self.output.dirname) + + if self.output.checkpoint: + # Open the checkpointing file for writing + if self.output.checkpoint_pickup_filename is not None: + chkfile = self.output.checkpoint_pickup_filename + else: + chkfile = path.join(self.dumpdir, "chkpt") + with DumbCheckpoint(chkfile, mode=FILE_READ) as chk: + # Recover all the fields from the checkpoint + for field in self.to_pickup: + chk.load(field) + t = chk.read_attribute("/", "time") + # Setup new checkpoint + self.chkpt = DumbCheckpoint(path.join(self.dumpdir, "chkpt"), mode=FILE_CREATE) + else: + raise ValueError("Must set checkpoint True if pickup") + + return t + + def dump(self, t): + """ + Dump output + """ + output = self.output + + # Diagnostics: + # Compute diagnostic fields + for field in self.diagnostic_fields: + field(self) + + if output.dump_diagnostics: + # Output diagnostic data + self.diagnostic_output.dump(self, t) + + if len(output.point_data) > 0 and (next(self.pddumpcount) % output.pddumpfreq) == 0: + # Output pointwise data + self.pointdata_output.dump(self.fields, t) + + # Dump all the fields to the checkpointing file (backup version) + if output.checkpoint and (next(self.chkptcount) % output.chkptfreq) == 0: + for field in self.to_pickup: + self.chkpt.store(field) + self.chkpt.write_attribute("/", "time", t) + + if output.dump_vtus and (next(self.dumpcount) % output.dumpfreq) == 0: + # dump fields + self.dumpfile.write(*self.to_dump) + + # dump fields on latlon mesh + if len(output.dumplist_latlon) > 0: + self.dumpfile_ll.write(*self.to_dump_latlon) + + def initialise(self, initial_conditions): + """ + Initialise state variables + + :arg initial_conditions: An iterable of pairs (field_name, pointwise_value) + """ + for name, ic in initial_conditions: + f_init = getattr(self.fields, name) + f_init.assign(ic) + f_init.rename(name) + + def set_reference_profiles(self, reference_profiles): + """ + Initialise reference profiles + + :arg reference_profiles: An iterable of pairs (field_name, interpolatory_value) + """ + for name, profile in reference_profiles: + if name+'bar' in self.fields: + # For reference profiles already added to state, allow + # interpolation from expressions + ref = self.fields(name+'bar') + elif isinstance(profile, Function): + # Need to add reference profile to state so profile must be + # a Function + ref = self.fields(name+'bar', space=profile.function_space(), dump=False) + else: + raise ValueError(f'When initialising reference profile {name}' + + ' the passed profile must be a Function') + ref.interpolate(profile) + + +def get_latlon_mesh(mesh): + coords_orig = mesh.coordinates + coords_fs = coords_orig.function_space() + + if coords_fs.extruded: + cell = mesh._base_mesh.ufl_cell().cellname() + DG1_hori_elt = FiniteElement("DG", cell, 1, variant="equispaced") + DG1_vert_elt = FiniteElement("DG", interval, 1, variant="equispaced") + DG1_elt = TensorProductElement(DG1_hori_elt, DG1_vert_elt) + else: + cell = mesh.ufl_cell().cellname() + DG1_elt = FiniteElement("DG", cell, 1, variant="equispaced") + vec_DG1 = VectorFunctionSpace(mesh, DG1_elt) + coords_dg = Function(vec_DG1).interpolate(coords_orig) + coords_latlon = Function(vec_DG1) + shapes = {"nDOFs": vec_DG1.finat_element.space_dimension(), 'dim': 3} + + radius = np.min(np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2)) + # lat-lon 'x' = atan2(y, x) + coords_latlon.dat.data[:, 0] = np.arctan2(coords_dg.dat.data[:, 1], coords_dg.dat.data[:, 0]) + # lat-lon 'y' = asin(z/sqrt(x^2 + y^2 + z^2)) + coords_latlon.dat.data[:, 1] = np.arcsin(coords_dg.dat.data[:, 2]/np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2)) + # our vertical coordinate is radius - the minimum radius + coords_latlon.dat.data[:, 2] = np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2) - radius + +# We need to ensure that all points in a cell are on the same side of the branch cut in longitude coords +# This kernel amends the longitude coords so that all longitudes in one cell are close together + kernel = op2.Kernel(""" +#define PI 3.141592653589793 +#define TWO_PI 6.283185307179586 +void splat_coords(double *coords) {{ + double max_diff = 0.0; + double diff = 0.0; + + for (int i=0; i<{nDOFs}; i++) {{ + for (int j=0; j<{nDOFs}; j++) {{ + diff = coords[i*{dim}] - coords[j*{dim}]; + if (fabs(diff) > max_diff) {{ + max_diff = diff; + }} + }} + }} + + if (max_diff > PI) {{ + for (int i=0; i<{nDOFs}; i++) {{ + if (coords[i*{dim}] < 0) {{ + coords[i*{dim}] += TWO_PI; + }} + }} + }} +}} +""".format(**shapes), "splat_coords") + + op2.par_loop(kernel, coords_latlon.cell_set, + coords_latlon.dat(op2.RW, coords_latlon.cell_node_map())) + return Mesh(coords_latlon) + + +def topo_sort(field_deps): + name2field = dict((f.name, f) for f, _ in field_deps) + # map node: (input_deps, output_deps) + graph = dict((f.name, (list(deps), [])) for f, deps in field_deps) + roots = [] + for f, input_deps in field_deps: + if len(input_deps) == 0: + # No dependencies, candidate for evaluation + roots.append(f.name) + for d in input_deps: + # add f as output dependency + graph[d][1].append(f.name) + + schedule = [] + while roots: + n = roots.pop() + schedule.append(n) + output_deps = list(graph[n][1]) + for m in output_deps: + # Remove edge + graph[m][0].remove(n) + graph[n][1].remove(m) + # If m now as no input deps, candidate for evaluation + if len(graph[m][0]) == 0: + roots.append(m) + if any(len(i) for i, _ in graph.values()): + cycle = "\n".join("%s -> %s" % (f, i) for f, (i, _) in graph.items() + if f not in schedule) + raise RuntimeError("Field dependencies have a cycle:\n\n%s" % cycle) + return list(map(name2field.__getitem__, schedule)) From 6bc50b37905876a8ce7e30ae41c6754a1b2ca890 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Tue, 6 Dec 2022 21:50:52 +0000 Subject: [PATCH 02/12] split state into domain and io objects, and update use of domain throughout most of the model --- gusto/__init__.py | 1 - gusto/diffusion.py | 9 +- gusto/domain.py | 50 +++ gusto/equations.py | 382 +++++++++++++---------- gusto/function_spaces.py | 201 ++++++++++++ gusto/initialisation_tools.py | 113 +++---- gusto/io.py | 556 ++++++++++++++++++++++++++++++++++ gusto/linear_solvers.py | 62 ++-- gusto/physics.py | 37 ++- gusto/transport_forms.py | 97 +++--- 10 files changed, 1184 insertions(+), 324 deletions(-) create mode 100644 gusto/domain.py create mode 100644 gusto/function_spaces.py create mode 100644 gusto/io.py diff --git a/gusto/__init__.py b/gusto/__init__.py index 2b1d6c40c..c85012ef0 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -3,7 +3,6 @@ from gusto.diagnostics import * # noqa from gusto.diffusion import * # noqa from gusto.equations import * # noqa -from gusto.fields import * # noqa from gusto.fml import * # noqa from gusto.forcing import * # noqa from gusto.initialisation_tools import * # noqa diff --git a/gusto/diffusion.py b/gusto/diffusion.py index 621ada6bd..175038e85 100644 --- a/gusto/diffusion.py +++ b/gusto/diffusion.py @@ -8,7 +8,7 @@ __all__ = ["interior_penalty_diffusion_form"] -def interior_penalty_diffusion_form(state, test, q, parameters): +def interior_penalty_diffusion_form(domain, test, q, parameters): u""" Form for the interior penalty discretisation of a diffusion term, ∇.(κ∇q) @@ -16,7 +16,8 @@ def interior_penalty_diffusion_form(state, test, q, parameters): weight function. Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the equation's test function. q (:class:`Function`): the variable being diffused. parameters (:class:`DiffusionParameters`): object containing metadata @@ -26,11 +27,11 @@ def interior_penalty_diffusion_form(state, test, q, parameters): :class:`ufl.Form`: the diffusion form. """ - dS_ = (dS_v + dS_h) if state.mesh.extruded else dS + dS_ = (dS_v + dS_h) if domain.mesh.extruded else dS kappa = parameters.kappa mu = parameters.mu - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) form = inner(grad(test), grad(q)*kappa)*dx diff --git a/gusto/domain.py b/gusto/domain.py new file mode 100644 index 000000000..c7e5667c0 --- /dev/null +++ b/gusto/domain.py @@ -0,0 +1,50 @@ +""" +The Domain object that is provided in this module contains the model's mesh and +the set of compatible function spaces defined upon it. +""" + +from gusto.spaces import Spaces +from firedrake import (Constant, SpatialCoordinate, sqrt, CellNormal, cross, + as_vector) + +class Domain(object): + """The Domain holds the model's mesh and its compatible function spaces.""" + def __init__(self, mesh, family, degree): + """ + Args: + mesh (:class:`Mesh`): the model's mesh. + family (str): the finite element space family used for the velocity + field. This determines the other finite element spaces used via + the de Rham complex. + degree (int): the element degree used for the velocity space. + """ + + self.mesh = mesh + self.spaces = [space for space in self._build_spaces(state, family, degree)] + + # figure out if we're on a sphere + try: + self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3 and mesh._base_mesh.topological_dimension() == 2) + except AttributeError: + self.on_sphere = (mesh.geometric_dimension() == 3 and mesh.topological_dimension() == 2) + + # build the vertical normal and define perp for 2d geometries + dim = mesh.topological_dimension() + if self.on_sphere: + x = SpatialCoordinate(mesh) + R = sqrt(inner(x, x)) + self.k = interpolate(x/R, mesh.coordinates.function_space()) + if dim == 2: + outward_normals = CellNormal(mesh) + self.perp = lambda u: cross(outward_normals, u) + else: + kvec = [0.0]*dim + kvec[dim-1] = 1.0 + self.k = Constant(kvec) + if dim == 2: + self.perp = lambda u: as_vector([-u[1], u[0]]) + + # TODO: why have this as a separate routine? + def _build_spaces(self, family, degree): + spaces = Spaces(self.mesh) + return spaces.build_compatible_spaces(family, degree) diff --git a/gusto/equations.py b/gusto/equations.py index 7d645312f..314516448 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -6,6 +6,7 @@ TrialFunction, FacetNormal, jump, avg, dS_v, DirichletBC, conditional, SpatialCoordinate, split, Constant, action) +from gusto.fields import StateFields from gusto.fml.form_manipulation_labelling import Term, all_terms, identity, drop from gusto.labels import (subject, time_derivative, transport, prognostic, transporting_velocity, replace_subject, linearisation, @@ -28,30 +29,30 @@ class PrognosticEquation(object, metaclass=ABCMeta): """Base class for prognostic equations.""" - def __init__(self, state, function_space, field_name): + def __init__(self, domain, function_space, field_name): """ Args: - state (:class:`State`): the model's state object. + 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.state = state + self.domain = domain self.function_space = function_space self.field_name = field_name + self.fields = StateFields() # TODO: should there be an argument here? self.bcs = {} if len(function_space) > 1: assert hasattr(self, "field_names") - state.fields(field_name, function_space, + self.fields(field_name, function_space, subfield_names=self.field_names, pickup=True) for fname in self.field_names: - state.diagnostics.register(fname) self.bcs[fname] = [] else: - state.fields(field_name, function_space) - state.diagnostics.register(field_name) + self.fields(field_name, function_space) self.bcs[field_name] = [] @@ -59,11 +60,12 @@ def __init__(self, state, function_space, field_name): class AdvectionEquation(PrognosticEquation): u"""Discretises the advection equation, ∂q/∂t + (u.∇)q = 0""" - def __init__(self, state, function_space, field_name, + def __init__(self, domain, function_space, field_name, ufamily=None, udegree=None, Vu=None, **kwargs): """ Args: - state (:class:`State`): the model's state object. + 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. @@ -77,33 +79,34 @@ def __init__(self, state, function_space, field_name, velocity field. If this is Defaults to None. **kwargs: any keyword arguments to be passed to the advection form. """ - super().__init__(state, function_space, field_name) + super().__init__(domain, function_space, field_name) - if not hasattr(state.fields, "u"): + if not hasattr(self.fields, "u"): if Vu is not None: - V = state.spaces("HDiv", V=Vu) + V = domain.spaces("HDiv", V=Vu) else: assert ufamily is not None, "Specify the family for u" assert udegree is not None, "Specify the degree of the u space" - V = state.spaces("HDiv", ufamily, udegree) - state.fields("u", V) + V = domain.spaces("HDiv", ufamily, udegree) + self.fields("u", V) test = TestFunction(function_space) q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) self.residual = subject( - mass_form + advection_form(state, test, q, **kwargs), q + mass_form + advection_form(domain, test, q, **kwargs), q ) class ContinuityEquation(PrognosticEquation): u"""Discretises the continuity equation, ∂q/∂t + ∇(u*q) = 0""" - def __init__(self, state, function_space, field_name, + def __init__(self, domain, function_space, field_name, ufamily=None, udegree=None, Vu=None, **kwargs): """ Args: - state (:class:`State`): the model's state object. + 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. @@ -117,40 +120,41 @@ def __init__(self, state, function_space, field_name, velocity field. If this is Defaults to None. **kwargs: any keyword arguments to be passed to the advection form. """ - super().__init__(state, function_space, field_name) + super().__init__(domain, function_space, field_name) - if not hasattr(state.fields, "u"): + if not hasattr(self.fields, "u"): if Vu is not None: - V = state.spaces("HDiv", V=Vu) + V = domain.spaces("HDiv", V=Vu) else: assert ufamily is not None, "Specify the family for u" assert udegree is not None, "Specify the degree of the u space" - V = state.spaces("HDiv", ufamily, udegree) - state.fields("u", V) + V = domain.spaces("HDiv", ufamily, udegree) + self.fields("u", V) test = TestFunction(function_space) q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) self.residual = subject( - mass_form + continuity_form(state, test, q, **kwargs), q + mass_form + continuity_form(domain, test, q, **kwargs), q ) class DiffusionEquation(PrognosticEquation): u"""Discretises the diffusion equation, ∂q/∂t = ∇.(κ∇q)""" - def __init__(self, state, function_space, field_name, + def __init__(self, domain, function_space, field_name, diffusion_parameters): """ Args: - state (:class:`State`): the model's state object. + 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__(state, function_space, field_name) + super().__init__(domain, function_space, field_name) test = TestFunction(function_space) q = Function(function_space) @@ -159,19 +163,20 @@ def __init__(self, state, function_space, field_name, self.residual = subject( mass_form + interior_penalty_diffusion_form( - state, test, q, diffusion_parameters), q + domain, test, q, diffusion_parameters), q ) class AdvectionDiffusionEquation(PrognosticEquation): u"""The advection-diffusion equation, ∂q/∂t + (u.∇)q = ∇.(κ∇q)""" - def __init__(self, state, function_space, field_name, + def __init__(self, domain, function_space, field_name, ufamily=None, udegree=None, Vu=None, diffusion_parameters=None, **kwargs): """ Args: - state (:class:`State`): the model's state object. + 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. @@ -188,25 +193,25 @@ def __init__(self, state, function_space, field_name, **kwargs: any keyword arguments to be passed to the advection form. """ - super().__init__(state, function_space, field_name) + super().__init__(domain, function_space, field_name) - if not hasattr(state.fields, "u"): + if not hasattr(self.fields, "u"): if Vu is not None: - V = state.spaces("HDiv", V=Vu) + V = domain.spaces("HDiv", V=Vu) else: assert ufamily is not None, "Specify the family for u" assert udegree is not None, "Specify the degree of the u space" - V = state.spaces("HDiv", ufamily, udegree) - state.fields("u", V) + V = domain.spaces("HDiv", ufamily, udegree) + self.fields("u", V) test = TestFunction(function_space) q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) self.residual = subject( mass_form - + advection_form(state, test, q, **kwargs) + + advection_form(domain, test, q, **kwargs) + interior_penalty_diffusion_form( - state, test, q, diffusion_parameters), q + domain, test, q, diffusion_parameters), q ) @@ -219,18 +224,14 @@ class PrognosticEquationSet(PrognosticEquation, metaclass=ABCMeta): contains common routines for these equation sets. """ - def __init__(self, field_names, state, family, degree, - linearisation_map=None, no_normal_flow_bc_ids=None, - active_tracers=None): + def __init__(self, field_names, domain, 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. - state (:class:`State`): the model's state object. - family (str): the finite element space family used for the velocity - field. This determines the other finite element spaces used via - the de Rham complex. - degree (int): the element degree used for the velocity space. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. 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 @@ -244,21 +245,19 @@ def __init__(self, field_names, state, family, degree, self.field_names = field_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 = [space for space in self._build_spaces(state, family, degree)] + self.reference_profiles_initialised = False # Add active tracers to the list of prognostics if active_tracers is None: active_tracers = [] - self.add_tracers_to_prognostics(state, 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__(state, W, full_field_name) + super().__init__(domain, W, full_field_name) # Set up test functions, trials and prognostics self.tests = TestFunctions(W) @@ -269,10 +268,7 @@ def __init__(self, field_names, state, family, degree, # 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(state, no_normal_flow_bc_ids) - - def _build_spaces(self, state, family, degree): - return state.spaces.build_compatible_spaces(family, degree) + self.set_no_normal_flow_bcs(domain, no_normal_flow_bc_ids) # ======================================================================== # # Set up time derivative / mass terms @@ -367,11 +363,41 @@ def linearise_equation_set(self): self.residual = self.residual.label_map( all_terms, replace_trial_function(self.X)) + # ======================================================================== # + # Reference Profile Routines + # ======================================================================== # + + def set_reference_profiles(self, reference_profiles): + """ + Initialise the equation'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. + """ + # TODO: come back and consider all aspects of this + for name, profile in reference_profiles: + if name+'bar' in self.fields: + # For reference profiles already added to state, allow + # interpolation from expressions + ref = self.fields(name+'bar') + elif isinstance(profile, Function): + # Need to add reference profile to state so profile must be + # a Function + ref = self.fields(name+'bar', space=profile.function_space(), dump=False) + else: + raise ValueError(f'When initialising reference profile {name}' + + ' the passed profile must be a Function') + ref.interpolate(profile) + + self.reference_profiles_initialised = True + # ======================================================================== # # Boundary Condition Routines # ======================================================================== # - def set_no_normal_flow_bcs(self, state, no_normal_flow_bc_ids): + 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. @@ -380,7 +406,8 @@ def set_no_normal_flow_bcs(self, state, no_normal_flow_bc_ids): a velocity variable named 'u' to apply the boundary conditions to. Args: - state (:class:`State`): the model's state. + 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. @@ -394,7 +421,7 @@ def set_no_normal_flow_bcs(self, state, no_normal_flow_bc_ids): 'No-normal-flow boundary conditions can only be applied ' + 'when there is a variable called "u" and none was found') - Vu = state.spaces("HDiv") + Vu = domain.spaces("HDiv") if Vu.extruded: self.bcs['u'].append(DirichletBC(Vu, 0.0, "bottom")) self.bcs['u'].append(DirichletBC(Vu, 0.0, "top")) @@ -412,12 +439,13 @@ def set_no_normal_flow_bcs(self, state, no_normal_flow_bc_ids): # Active Tracer Routines # ======================================================================== # - def add_tracers_to_prognostics(self, state, active_tracers): + def add_tracers_to_prognostics(self, domain, active_tracers): """ Augments the equation set with specified active tracer variables. Args: - state (:class:`State`): the model's state. + 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. @@ -433,16 +461,17 @@ def add_tracers_to_prognostics(self, state, active_tracers): self.field_names.append(tracer.name) else: raise ValueError(f'There is already a field named {tracer.name}') - self.spaces.append(state.spaces(tracer.space)) + self.spaces.append(domain.spaces(tracer.space)) else: raise TypeError(f'Tracers must be ActiveTracer objects, not {type(tracer)}') - def generate_tracer_transport_terms(self, state, active_tracers): + def generate_tracer_transport_terms(self, domain, active_tracers): """ Adds the transport forms for the active tracers to the equation set. Args: - state (:class:`State`): the model's state. + 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. @@ -465,9 +494,9 @@ def generate_tracer_transport_terms(self, state, active_tracers): tracer_prog = split(self.X)[idx] tracer_test = self.tests[idx] if tracer.transport_eqn == TransportEquationType.advective: - tracer_adv = prognostic(advection_form(state, tracer_test, tracer_prog), tracer.name) + tracer_adv = prognostic(advection_form(domain, tracer_test, tracer_prog), tracer.name) elif tracer.transport_eqn == TransportEquationType.conservative: - tracer_adv = prognostic(continuity_form(state, tracer_test, tracer_prog), tracer.name) + tracer_adv = prognostic(continuity_form(domain, tracer_test, tracer_prog), tracer.name) else: raise ValueError(f'Transport eqn {tracer.transport_eqn} not recognised') @@ -482,38 +511,62 @@ def generate_tracer_transport_terms(self, state, active_tracers): class ForcedAdvectionEquation(PrognosticEquationSet): - - def __init__(self, state, function_space, field_name, + u""" + Discretises the advection equation with a source/sink term, + ∂q/∂t + (u.∇)q = F, + which can also be augmented with active tracers. + """ + def __init__(self, domain, function_space, field_name, ufamily=None, udegree=None, Vu=None, active_tracers=None, **kwargs): + """ + 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. + ufamily (str, optional): the family of the function space to use + for the velocity field. Only used if `Vu` is not provided. + Defaults to None. + udegree (int, optional): the degree of the function space to use for + the velocity field. Only used if `Vu` is not provided. Defaults + to None. + Vu (:class:`FunctionSpace`, optional): the function space for the + velocity field. If this is 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. + **kwargs: any keyword arguments to be passed to the advection form. + """ self.field_names = [field_name] self.active_tracers = active_tracers self.terms_to_linearise = {} # Build finite element spaces - self.spaces = [state.spaces("tracer", V=function_space)] + self.spaces = [domain.spaces("tracer", V=function_space)] # Add active tracers to the list of prognostics if active_tracers is None: active_tracers = [] - self.add_tracers_to_prognostics(state, 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) - PrognosticEquation.__init__(self, state, W, full_field_name) + PrognosticEquation.__init__(self, domain, W, full_field_name) - if not hasattr(state.fields, "u"): + if not hasattr(self.fields, "u"): if Vu is not None: - V = state.spaces("HDiv", V=Vu) + V = domain.spaces("HDiv", V=Vu) else: assert ufamily is not None, "Specify the family for u" assert udegree is not None, "Specify the degree of the u space" - V = state.spaces("HDiv", ufamily, udegree) - state.fields("u", V) + V = domain.spaces("HDiv", ufamily, udegree) + self.fields("u", V) self.tests = TestFunctions(W) self.X = Function(W) @@ -521,7 +574,7 @@ def __init__(self, state, function_space, field_name, mass_form = self.generate_mass_terms() self.residual = subject( - mass_form + advection_form(state, self.tests[0], split(self.X)[0], **kwargs), self.X + mass_form + advection_form(domain, self.tests[0], split(self.X)[0], **kwargs), self.X ) # ============================================================================ # @@ -538,17 +591,16 @@ class ShallowWaterEquations(PrognosticEquationSet): for Coriolis parameter 'f' and bottom surface 'b'. """ - def __init__(self, state, family, degree, fexpr=None, bexpr=None, + def __init__(self, domain, parameters, fexpr=None, bexpr=None, linearisation_map='default', u_transport_option='vector_invariant_form', no_normal_flow_bc_ids=None, active_tracers=None): """ Args: - state (:class:`State`): the model's state object. - family (str): the finite element space family used for the velocity - field. This determines the other finite element spaces used via - the de Rham complex. - degree (int): the element degree used for the velocity space. + 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 @@ -586,13 +638,14 @@ def __init__(self, state, family, degree, fexpr=None, bexpr=None, (any(t.has_label(time_derivative, pressure_gradient)) or (t.get(prognostic) == "D" and t.has_label(transport))) - super().__init__(field_names, state, family, degree, + super().__init__(field_names, domain, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) - g = state.parameters.g - H = state.parameters.H + self.parameters = parameters + g = parameters.g + H = parameters.H w, phi = self.tests[0:2] u, D = split(self.X)[0:2] @@ -608,28 +661,28 @@ def __init__(self, state, family, degree, fexpr=None, bexpr=None, # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": - u_adv = prognostic(vector_invariant_form(state, w, u), "u") + u_adv = prognostic(vector_invariant_form(domain, w, u), "u") elif u_transport_option == "vector_advection_form": - u_adv = prognostic(advection_form(state, w, u), "u") + u_adv = prognostic(advection_form(domain, w, u), "u") elif u_transport_option == "vector_manifold_advection_form": - u_adv = prognostic(vector_manifold_advection_form(state, w, u), "u") + u_adv = prognostic(vector_manifold_advection_form(domain, w, u), "u") elif u_transport_option == "circulation_form": - ke_form = prognostic(kinetic_energy_form(state, w, u), "u") + ke_form = prognostic(kinetic_energy_form(domain, w, u), "u") ke_form = transport.remove(ke_form) ke_form = ke_form.label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) - u_adv = prognostic(advection_equation_circulation_form(state, w, u), "u") + ke_form + u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) # Depth transport term - D_adv = prognostic(continuity_form(state, phi, D), "D") + D_adv = prognostic(continuity_form(domain, phi, D), "D") # Transport term needs special linearisation if self.linearisation_map(D_adv.terms[0]): - linear_D_adv = linear_continuity_form(state, phi, H).label_map( + linear_D_adv = linear_continuity_form(domain, phi, H).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u_trial}), t.labels)) @@ -640,7 +693,7 @@ def __init__(self, state, family, degree, fexpr=None, bexpr=None, # Add transport of tracers if len(active_tracers) > 0: - adv_form += self.generate_tracer_transport_terms(state, active_tracers) + adv_form += self.generate_tracer_transport_terms(domain, active_tracers) # -------------------------------------------------------------------- # # Pressure Gradient Term @@ -654,19 +707,19 @@ def __init__(self, state, family, degree, fexpr=None, bexpr=None, # Extra Terms (Coriolis and Topography) # -------------------------------------------------------------------- # if fexpr is not None: - V = FunctionSpace(state.mesh, "CG", 1) - f = state.fields("coriolis", space=V) + V = FunctionSpace(domain.mesh, "CG", 1) + f = self.fields("coriolis", space=V) f.interpolate(fexpr) coriolis_form = coriolis( - subject(prognostic(f*inner(state.perp(u), w)*dx, "u"), self.X)) + subject(prognostic(f*inner(domain.perp(u), w)*dx, "u"), self.X)) # Add linearisation linear_coriolis = coriolis( - subject(prognostic(f*inner(state.perp(u_trial), w)*dx, "u"), self.X)) + 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: - b = state.fields("topography", state.spaces("DG")) + b = self.fields("topography", domain.spaces("DG")) b.interpolate(bexpr) topography_form = subject(prognostic(-g*div(w)*b*dx, "u"), self.X) residual += topography_form @@ -696,17 +749,16 @@ class LinearShallowWaterEquations(ShallowWaterEquations): which is then linearised. """ - def __init__(self, state, family, degree, fexpr=None, bexpr=None, + def __init__(self, domain, parameters, fexpr=None, bexpr=None, linearisation_map='default', u_transport_option="vector_invariant_form", no_normal_flow_bc_ids=None, active_tracers=None): """ Args: - state (:class:`State`): the model's state object. - family (str): the finite element space family used for the velocity - field. This determines the other finite element spaces used via - the de Rham complex. - degree (int): the element degree used for the velocity space. + 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 @@ -736,7 +788,7 @@ def __init__(self, state, family, degree, fexpr=None, bexpr=None, (any(t.has_label(time_derivative, pressure_gradient, coriolis)) or (t.get(prognostic) == "D" and t.has_label(transport))) - super().__init__(state, family, degree, fexpr=fexpr, bexpr=bexpr, + super().__init__(domain, parameters, fexpr=fexpr, bexpr=bexpr, linearisation_map=linearisation_map, u_transport_option=u_transport_option, no_normal_flow_bc_ids=no_normal_flow_bc_ids, @@ -748,7 +800,7 @@ def __init__(self, state, family, degree, fexpr=None, bexpr=None, # D transport term is a special case -- add facet term _, D = split(self.X) _, phi = self.tests - D_adv = prognostic(linear_continuity_form(state, phi, D, facet_term=True), "D") + D_adv = prognostic(linear_continuity_form(domain, phi, D, facet_term=True), "D") self.residual = self.residual.label_map( lambda t: t.has_label(transport) and t.get(prognostic) == "D", map_if_true=lambda t: Term(D_adv.form, t.labels) @@ -768,7 +820,7 @@ class CompressibleEulerEquations(PrognosticEquationSet): pressure. """ - def __init__(self, state, family, degree, Omega=None, sponge=None, + def __init__(self, domain, parameters, Omega=None, sponge=None, extra_terms=None, linearisation_map='default', u_transport_option="vector_invariant_form", diffusion_options=None, @@ -776,11 +828,10 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, active_tracers=None): """ Args: - state (:class:`State`): the model's state object. - family (str): the finite element space family used for the velocity - field. This determines the other finite element spaces used via - the de Rham complex. - degree (int): the element degree used for the velocity space. + 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 @@ -822,22 +873,23 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, (t.has_label(time_derivative) or (t.get(prognostic) != "u" and t.has_label(transport))) - super().__init__(field_names, state, family, degree, + super().__init__(field_names, domain, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) - g = state.parameters.g - cp = state.parameters.cp + 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] - rhobar = state.fields("rhobar", space=state.spaces("DG"), dump=False) - thetabar = state.fields("thetabar", space=state.spaces("theta"), dump=False) + rhobar = self.fields("rhobar", space=domain.spaces("DG"), dump=False) + thetabar = self.fields("thetabar", space=domain.spaces("theta"), dump=False) zero_expr = Constant(0.0)*theta - exner = exner_pressure(state.parameters, rho, theta) - n = FacetNormal(state.mesh) + exner = exner_pressure(parameters, rho, theta) + n = FacetNormal(domain.mesh) # -------------------------------------------------------------------- # # Time Derivative Terms @@ -849,38 +901,38 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": - u_adv = prognostic(vector_invariant_form(state, w, u), "u") + u_adv = prognostic(vector_invariant_form(domain, w, u), "u") elif u_transport_option == "vector_advection_form": - u_adv = prognostic(advection_form(state, w, u), "u") + u_adv = prognostic(advection_form(domain, w, u), "u") elif u_transport_option == "vector_manifold_advection_form": - u_adv = prognostic(vector_manifold_advection_form(state, w, u), "u") + u_adv = prognostic(vector_manifold_advection_form(domain, w, u), "u") elif u_transport_option == "circulation_form": - ke_form = prognostic(kinetic_energy_form(state, w, u), "u") + ke_form = prognostic(kinetic_energy_form(domain, w, u), "u") ke_form = transport.remove(ke_form) ke_form = ke_form.label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) - u_adv = prognostic(advection_equation_circulation_form(state, w, u), "u") + ke_form + u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) # Density transport (conservative form) - rho_adv = prognostic(continuity_form(state, phi, rho), "rho") + rho_adv = prognostic(continuity_form(domain, phi, rho), "rho") # Transport term needs special linearisation if self.linearisation_map(rho_adv.terms[0]): - linear_rho_adv = linear_continuity_form(state, phi, rhobar).label_map( + linear_rho_adv = linear_continuity_form(domain, phi, rhobar).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u_trial}), t.labels)) rho_adv = linearisation(rho_adv, linear_rho_adv) # Potential temperature transport (advective form) - theta_adv = prognostic(advection_form(state, gamma, theta), "theta") + theta_adv = prognostic(advection_form(domain, gamma, theta), "theta") # Transport term needs special linearisation if self.linearisation_map(theta_adv.terms[0]): - linear_theta_adv = linear_advection_form(state, gamma, thetabar).label_map( + linear_theta_adv = linear_advection_form(domain, gamma, thetabar).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u_trial}), t.labels)) @@ -890,7 +942,7 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, # Add transport of tracers if len(active_tracers) > 0: - adv_form += self.generate_tracer_transport_terms(state, active_tracers) + adv_form += self.generate_tracer_transport_terms(domain, active_tracers) # -------------------------------------------------------------------- # # Pressure Gradient Term @@ -912,7 +964,7 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, # -------------------------------------------------------------------- # # Gravitational Term # -------------------------------------------------------------------- # - gravity_form = subject(prognostic(Term(g*inner(state.k, w)*dx), "u"), self.X) + gravity_form = subject(prognostic(Term(g*inner(domain.k, w)*dx), "u"), self.X) residual = (mass_form + adv_form + pressure_gradient_form + gravity_form) @@ -920,12 +972,12 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, # Moist Thermodynamic Divergence Term # -------------------------------------------------------------------- # if len(active_tracers) > 0: - cv = state.parameters.cv - c_vv = state.parameters.c_vv - c_pv = state.parameters.c_pv - c_pl = state.parameters.c_pl - R_d = state.parameters.R_d - R_v = state.parameters.R_v + 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 @@ -959,8 +1011,8 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, inner(w, cross(2*Omega, u))*dx, "u"), self.X) if sponge is not None: - W_DG = FunctionSpace(state.mesh, "DG", 2) - x = SpatialCoordinate(state.mesh) + W_DG = FunctionSpace(domain.mesh, "DG", 2) + x = SpatialCoordinate(domain.mesh) z = x[len(x)-1] H = sponge.H zc = sponge.z_level @@ -972,7 +1024,7 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, self.mu = Function(W_DG).interpolate(muexpr) residual += name(subject(prognostic( - self.mu*inner(w, state.k)*inner(u, state.k)*dx, "u"), self.X), "sponge") + self.mu*inner(w, domain.k)*inner(u, domain.k)*dx, "u"), self.X), "sponge") if diffusion_options is not None: for field, diffusion in diffusion_options: @@ -981,7 +1033,7 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, fn = split(self.X)[idx] residual += subject( prognostic(interior_penalty_diffusion_form( - state, test, fn, diffusion), field), self.X) + domain, test, fn, diffusion), field), self.X) if extra_terms is not None: for field, term in extra_terms: @@ -1016,7 +1068,7 @@ class HydrostaticCompressibleEulerEquations(CompressibleEulerEquations): equations. """ - def __init__(self, state, family, degree, Omega=None, sponge=None, + def __init__(self, domain, parameters, Omega=None, sponge=None, extra_terms=None, linearisation_map='default', u_transport_option="vector_invariant_form", diffusion_options=None, @@ -1024,11 +1076,10 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, active_tracers=None): """ Args: - state (:class:`State`): the model's state object. - family (str): the finite element space family used for the velocity - field. This determines the other finite element spaces used via - the de Rham complex. - degree (int): the element degree used for the velocity space. + 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 @@ -1059,7 +1110,7 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, NotImplementedError: only mixing ratio tracers are implemented. """ - super().__init__(state, family, degree, Omega=Omega, sponge=sponge, + super().__init__(domain, parameters, Omega=Omega, sponge=sponge, extra_terms=extra_terms, linearisation_map=linearisation_map, u_transport_option=u_transport_option, @@ -1072,7 +1123,7 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, map_if_true=lambda t: hydrostatic(t, self.hydrostatic_projection(t)) ) - k = self.state.k + k = self.domain.k u = split(self.X)[0] self.residual += name( subject( @@ -1099,8 +1150,8 @@ def hydrostatic_projection(self, t): """ # TODO: make this more general, i.e. should work on the sphere - assert not self.state.on_sphere, "the hydrostatic projection is not yet implemented for spherical geometry" - k = Constant((*self.state.k, 0, 0)) + assert not self.domain.on_sphere, "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) @@ -1121,18 +1172,17 @@ class IncompressibleBoussinesqEquations(PrognosticEquationSet): where k is the vertical unit vector and, Ω is the planet's rotation vector. """ - def __init__(self, state, family, degree, Omega=None, + def __init__(self, domain, parameters, Omega=None, linearisation_map='default', u_transport_option="vector_invariant_form", no_normal_flow_bc_ids=None, active_tracers=None): """ Args: - state (:class:`State`): the model's state object. - family (str): the finite element space family used for the velocity - field. This determines the other finite element spaces used via - the de Rham complex. - degree (int): the element degree used for the velocity space. + 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. linearisation_map (func, optional): a function specifying which @@ -1170,16 +1220,18 @@ def __init__(self, state, family, degree, Omega=None, (t.has_label(time_derivative) or (t.get(prognostic) not in ["u", "p"] and t.has_label(transport))) - super().__init__(field_names, state, family, degree, + super().__init__(field_names, domain, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) + self.parameters = parameters + w, phi, gamma = self.tests[0:3] u, p, b = split(self.X) u_trial = split(self.trials)[0] - bbar = state.fields("bbar", space=state.spaces("theta"), dump=False) - bbar = state.fields("pbar", space=state.spaces("DG"), dump=False) + bbar = self.fields("bbar", space=domain.spaces("theta"), dump=False) + bbar = self.fields("pbar", space=domain.spaces("DG"), dump=False) # -------------------------------------------------------------------- # # Time Derivative Terms @@ -1191,27 +1243,27 @@ def __init__(self, state, family, degree, Omega=None, # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": - u_adv = prognostic(vector_invariant_form(state, w, u), "u") + u_adv = prognostic(vector_invariant_form(domain, w, u), "u") elif u_transport_option == "vector_advection_form": - u_adv = prognostic(advection_form(state, w, u), "u") + u_adv = prognostic(advection_form(domain, w, u), "u") elif u_transport_option == "vector_manifold_advection_form": - u_adv = prognostic(vector_manifold_advection_form(state, w, u), "u") + u_adv = prognostic(vector_manifold_advection_form(domain, w, u), "u") elif u_transport_option == "circulation_form": - ke_form = prognostic(kinetic_energy_form(state, w, u), "u") + ke_form = prognostic(kinetic_energy_form(domain, w, u), "u") ke_form = transport.remove(ke_form) ke_form = ke_form.label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) - u_adv = prognostic(advection_equation_circulation_form(state, w, u), "u") + ke_form + u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) # Buoyancy transport - b_adv = prognostic(advection_form(state, gamma, b), "b") + b_adv = prognostic(advection_form(domain, gamma, b), "b") if self.linearisation_map(b_adv.terms[0]): - linear_b_adv = linear_advection_form(state, gamma, bbar).label_map( + linear_b_adv = linear_advection_form(domain, gamma, bbar).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u_trial}), t.labels)) @@ -1221,7 +1273,7 @@ def __init__(self, state, family, degree, Omega=None, # Add transport of tracers if len(active_tracers) > 0: - adv_form += self.generate_tracer_transport_terms(state, active_tracers) + adv_form += self.generate_tracer_transport_terms(domain, active_tracers) # -------------------------------------------------------------------- # # Pressure Gradient Term @@ -1231,7 +1283,7 @@ def __init__(self, state, family, degree, Omega=None, # -------------------------------------------------------------------- # # Gravitational Term # -------------------------------------------------------------------- # - gravity_form = subject(prognostic(-b*inner(w, state.k)*dx, "u"), self.X) + gravity_form = subject(prognostic(-b*inner(w, domain.k)*dx, "u"), self.X) # -------------------------------------------------------------------- # # Divergence Term diff --git a/gusto/function_spaces.py b/gusto/function_spaces.py new file mode 100644 index 000000000..8d1f1f066 --- /dev/null +++ b/gusto/function_spaces.py @@ -0,0 +1,201 @@ +""" +This module contains routines to generate the compatible function spaces to be +used by the model. +""" + +from firedrake import (HDiv, FunctionSpace, FiniteElement, TensorProductElement, + interval) + +class Spaces(object): + """Object to create and hold the model's finite element spaces.""" + def __init__(self, mesh): + """ + Args: + mesh (:class:`Mesh`): the model's mesh. + """ + self.mesh = mesh + self.extruded_mesh = hasattr(mesh, "_base_mesh") + self._initialised_base_spaces = False + + def __call__(self, name, family=None, degree=None, V=None): + """ + Returns a space, and also creates it if it is not created yet. + + If a space needs creating, it may be that more arguments (such as the + family and degree) need to be provided. Alternatively a space can be + passed in to be stored in the creator. + + Args: + name (str): the name of the space. + family (str, optional): name of the finite element family to be + created. Defaults to None. + degree (int, optional): the degree of the finite element space to be + created. Defaults to None. + V (:class:`FunctionSpace`, optional): an existing space, to be + stored in the creator object. If this is provided, it will be + added to the creator and no other action will be taken. This + space will be returned. Defaults to None. + + Returns: + :class:`FunctionSpace`: the desired function space. + """ + + try: + return getattr(self, name) + except AttributeError: + if V is not None: + value = V + elif name == "HDiv" and family in ["BDM", "RT", "CG", "RTCF"]: + value = self.build_hdiv_space(family, degree) + elif name == "theta": + value = self.build_theta_space(degree) + elif name == "DG1_equispaced": + value = self.build_dg_space(1, variant='equispaced') + elif family == "DG": + value = self.build_dg_space(degree) + elif family == "CG": + value = self.build_cg_space(degree) + else: + raise ValueError(f'State has no space corresponding to {name}') + setattr(self, name, value) + return value + + def build_compatible_spaces(self, family, degree): + """ + Builds the sequence of compatible finite element spaces for the mesh. + + If the mesh is not extruded, this builds and returns the spaces: + (HDiv, DG). + If the mesh is extruded, this builds and returns the following spaces: + (HDiv, DG, theta). + The 'theta' space corresponds to the vertical component of the velocity. + + Args: + family (str): the family of the horizontal part of the HDiv space. + degree (int): the polynomial degree of the DG space. + + Returns: + tuple: the created compatible :class:`FunctionSpace` objects. + """ + if self.extruded_mesh and not self._initialised_base_spaces: + self.build_base_spaces(family, degree) + Vu = self.build_hdiv_space(family, degree) + setattr(self, "HDiv", Vu) + Vdg = self.build_dg_space(degree) + setattr(self, "DG", Vdg) + Vth = self.build_theta_space(degree) + setattr(self, "theta", Vth) + return Vu, Vdg, Vth + else: + Vu = self.build_hdiv_space(family, degree) + setattr(self, "HDiv", Vu) + Vdg = self.build_dg_space(degree) + setattr(self, "DG", Vdg) + return Vu, Vdg + + def build_base_spaces(self, family, degree): + """ + Builds the :class:`FiniteElement` objects for the base mesh. + + Args: + family (str): the family of the horizontal part of the HDiv space. + degree (int): the polynomial degree of the DG space. + """ + cell = self.mesh._base_mesh.ufl_cell().cellname() + + # horizontal base spaces + self.S1 = FiniteElement(family, cell, degree+1) + self.S2 = FiniteElement("DG", cell, degree) + + # vertical base spaces + self.T0 = FiniteElement("CG", interval, degree+1) + self.T1 = FiniteElement("DG", interval, degree) + + self._initialised_base_spaces = True + + def build_hdiv_space(self, family, degree): + """ + Builds and returns the HDiv :class:`FunctionSpace`. + + Args: + family (str): the family of the horizontal part of the HDiv space. + degree (int): the polynomial degree of the space. + + Returns: + :class:`FunctionSpace`: the HDiv space. + """ + if self.extruded_mesh: + if not self._initialised_base_spaces: + self.build_base_spaces(family, degree) + Vh_elt = HDiv(TensorProductElement(self.S1, self.T1)) + Vt_elt = TensorProductElement(self.S2, self.T0) + Vv_elt = HDiv(Vt_elt) + V_elt = Vh_elt + Vv_elt + else: + cell = self.mesh.ufl_cell().cellname() + V_elt = FiniteElement(family, cell, degree+1) + return FunctionSpace(self.mesh, V_elt, name='HDiv') + + def build_dg_space(self, degree, variant=None): + """ + Builds and returns the DG :class:`FunctionSpace`. + + Args: + degree (int): the polynomial degree of the space. + variant (str): the variant of the underlying :class:`FiniteElement` + to use. Defaults to None, which will call the default variant. + + Returns: + :class:`FunctionSpace`: the DG space. + """ + if self.extruded_mesh: + if not self._initialised_base_spaces or self.T1.degree() != degree or self.T1.variant() != variant: + cell = self.mesh._base_mesh.ufl_cell().cellname() + S2 = FiniteElement("DG", cell, degree, variant=variant) + T1 = FiniteElement("DG", interval, degree, variant=variant) + else: + S2 = self.S2 + T1 = self.T1 + V_elt = TensorProductElement(S2, T1) + else: + cell = self.mesh.ufl_cell().cellname() + V_elt = FiniteElement("DG", cell, degree, variant=variant) + name = f'DG{degree}_equispaced' if variant == 'equispaced' else f'DG{degree}' + return FunctionSpace(self.mesh, V_elt, name=name) + + def build_theta_space(self, degree): + """ + Builds and returns the 'theta' space. + + This corresponds to the non-Piola mapped space of the vertical component + of the velocity. The space will be discontinuous in the horizontal but + continuous in the vertical. + + Args: + degree (int): degree of the corresponding density space. + + Raises: + AssertionError: the mesh is not extruded. + + Returns: + :class:`FunctionSpace`: the 'theta' space. + """ + assert self.extruded_mesh + if not self._initialised_base_spaces: + cell = self.mesh._base_mesh.ufl_cell().cellname() + self.S2 = FiniteElement("DG", cell, degree) + self.T0 = FiniteElement("CG", interval, degree+1) + V_elt = TensorProductElement(self.S2, self.T0) + return FunctionSpace(self.mesh, V_elt, name='Vtheta') + + def build_cg_space(self, degree): + """ + Builds the continuous scalar space at the top of the de Rham complex. + + Args: + degree (int): degree of the continuous space. + + Returns: + :class:`FunctionSpace`: the continuous space. + """ + return FunctionSpace(self.mesh, "CG", degree, name=f'CG{degree}') diff --git a/gusto/initialisation_tools.py b/gusto/initialisation_tools.py index 1c8aa9b32..65adeb661 100644 --- a/gusto/initialisation_tools.py +++ b/gusto/initialisation_tools.py @@ -60,7 +60,7 @@ def sphere_to_cartesian(mesh, u_zonal, u_merid): return as_vector((cartesian_u_expr, cartesian_v_expr, cartesian_w_expr)) -def incompressible_hydrostatic_balance(state, b0, p0, top=False, params=None): +def incompressible_hydrostatic_balance(equation, b0, p0, top=False, params=None): """ Gives a pressure field in hydrostatic-balance for the Incompressible eqns. @@ -70,7 +70,7 @@ def incompressible_hydrostatic_balance(state, b0, p0, top=False, params=None): with zero flow enforced at one of the boundaries. Args: - state (:class:`State`): the model's state. + equation (:class:`PrognosticEquation`): the model's equation object. b0 (:class:`ufl.Expr`): the input buoyancy field. p0 (:class:`Function`): the pressure to be returned. top (bool, optional): whether the no-flow boundary condition is enforced @@ -81,8 +81,9 @@ def incompressible_hydrostatic_balance(state, b0, p0, top=False, params=None): """ # get F - Vu = state.spaces("HDiv") - Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1]) + domain = equation.domain + Vu = domain.spaces("HDiv") + Vv = FunctionSpace(equation.domain.mesh, Vu.ufl_element()._elements[-1]) v = TrialFunction(Vv) w = TestFunction(Vv) @@ -94,13 +95,13 @@ def incompressible_hydrostatic_balance(state, b0, p0, top=False, params=None): bcs = [DirichletBC(Vv, 0.0, bstring)] a = inner(w, v)*dx - L = inner(state.k, w)*b0*dx + L = inner(equation.domain.k, w)*b0*dx F = Function(Vv) solve(a == L, F, bcs=bcs) # define mixed function space - VDG = state.spaces("DG") + VDG = domain.spaces("DG") WV = (Vv)*(VDG) # get pprime @@ -136,7 +137,7 @@ def incompressible_hydrostatic_balance(state, b0, p0, top=False, params=None): p0.project(pprime) -def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, +def compressible_hydrostatic_balance(equation, theta0, rho0, exner0=None, top=False, exner_boundary=Constant(1.0), mr_t=None, solve_for_rho=False, @@ -151,7 +152,7 @@ def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, procedure for solving the resulting discrete systems. Args: - state (:class:`State`): the model's state. + equation (:class:`PrognosticEquation`): the model's equation object. theta0 (:class:`ufl.Expr`): the input (dry) potential temperature field. rho0 (:class:`Function`): the hydrostatically-balanced density to be found. @@ -175,16 +176,18 @@ def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, """ # Calculate hydrostatic Pi - VDG = state.spaces("DG") - Vu = state.spaces("HDiv") - Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1]) + domain = equation.domain + parameters = equation.parameters + VDG = domain.spaces("DG") + Vu = domain.spaces("HDiv") + Vv = FunctionSpace(equation.domain.mesh, Vu.ufl_element()._elements[-1]) W = MixedFunctionSpace((Vv, VDG)) v, exner = TrialFunctions(W) dv, dexner = TestFunctions(W) - n = FacetNormal(state.mesh) + n = FacetNormal(equation.domain.mesh) - cp = state.parameters.cp + cp = parameters.cp # add effect of density of water upon theta theta = theta0 @@ -207,9 +210,9 @@ def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, arhs = -cp*inner(dv, n)*theta*exner_boundary*bmeasure # Possibly make g vary with spatial coordinates? - g = state.parameters.g + g = parameters.g - arhs -= g*inner(dv, state.k)*dx + arhs -= g*inner(dv, equation.domain.k)*dx bcs = [DirichletBC(W.sub(0), zero(), bstring)] @@ -239,16 +242,16 @@ def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, if solve_for_rho: w1 = Function(W) v, rho = w1.split() - rho.interpolate(thermodynamics.rho(state.parameters, theta0, exner)) + rho.interpolate(thermodynamics.rho(parameters, theta0, exner)) v, rho = split(w1) dv, dexner = TestFunctions(W) - exner = thermodynamics.exner_pressure(state.parameters, rho, theta0) + exner = thermodynamics.exner_pressure(parameters, rho, theta0) F = ( (cp*inner(v, dv) - cp*div(dv*theta)*exner)*dx + dexner*div(theta0*v)*dx + cp*inner(dv, n)*theta*exner_boundary*bmeasure ) - F += g*inner(dv, state.k)*dx + F += g*inner(dv, equation.domain.k)*dx rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs) rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params, options_prefix="rhosolver") @@ -256,7 +259,7 @@ def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, v, rho_ = w1.split() rho0.assign(rho_) else: - rho0.interpolate(thermodynamics.rho(state.parameters, theta0, exner)) + rho0.interpolate(thermodynamics.rho(parameters, theta0, exner)) def remove_initial_w(u): @@ -276,7 +279,7 @@ def remove_initial_w(u): u.assign(uin) -def saturated_hydrostatic_balance(state, theta_e, mr_t, exner0=None, +def saturated_hydrostatic_balance(equation, theta_e, mr_t, exner0=None, top=False, exner_boundary=Constant(1.0), max_outer_solve_count=40, max_theta_solve_count=5, @@ -296,8 +299,7 @@ def saturated_hydrostatic_balance(state, theta_e, mr_t, exner0=None, converge to a solution. Args: - state (:class:`State`): the model's state object, through which the - prognostic variables are accessed. + equation (:class:`PrognosticEquation`): the model's equation object. theta_e (:class:`ufl.Expr`): expression for the desired wet equivalent potential temperature field. mr_t (:class:`ufl.Expr`): expression for the total moisture content. @@ -324,15 +326,17 @@ def saturated_hydrostatic_balance(state, theta_e, mr_t, exner0=None, number of iterations. """ - theta0 = state.fields('theta') - rho0 = state.fields('rho') - mr_v0 = state.fields('water_vapour') + theta0 = equation.fields('theta') + rho0 = equation.fields('rho') + mr_v0 = equation.fields('water_vapour') # Calculate hydrostatic exner pressure + domain = equation.domain + parameters = equation.parameters Vt = theta0.function_space() Vr = rho0.function_space() - VDG = state.spaces("DG") + VDG = domain.spaces("DG") if any(deg > 2 for deg in VDG.ufl_element().degree()): logger.warning("default quadrature degree most likely not sufficient for this degree element") @@ -353,15 +357,15 @@ def saturated_hydrostatic_balance(state, theta_e, mr_t, exner0=None, delta = 0.8 # expressions for finding theta0 and mr_v0 from theta_e and mr_t - exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) - p = thermodynamics.p(state.parameters, exner) - T = thermodynamics.T(state.parameters, theta0, exner, mr_v0) - r_v_expr = thermodynamics.r_sat(state.parameters, T, p) - theta_e_expr = thermodynamics.theta_e(state.parameters, T, p, mr_v0, mr_t) + exner = thermodynamics.exner_pressure(parameters, rho_averaged, theta0) + p = thermodynamics.p(parameters, exner) + T = thermodynamics.T(parameters, theta0, exner, mr_v0) + r_v_expr = thermodynamics.r_sat(parameters, T, p) + theta_e_expr = thermodynamics.theta_e(parameters, T, p, mr_v0, mr_t) for i in range(max_outer_solve_count): # solve for rho with theta_vd and w_v guesses - compressible_hydrostatic_balance(state, theta0, rho_h, top=top, + compressible_hydrostatic_balance(equation, theta0, rho_h, top=top, exner_boundary=exner_boundary, mr_t=mr_t, solve_for_rho=True) @@ -396,16 +400,16 @@ def saturated_hydrostatic_balance(state, theta_e, mr_t, exner0=None, raise RuntimeError('Hydrostatic balance solve has not converged within %i' % i, 'iterations') if exner0 is not None: - exner = thermodynamics.exner(state.parameters, rho0, theta0) + exner = thermodynamics.exner(parameters, rho0, theta0) exner0.interpolate(exner) # do one extra solve for rho - compressible_hydrostatic_balance(state, theta0, rho0, top=top, + compressible_hydrostatic_balance(equation, theta0, rho0, top=top, exner_boundary=exner_boundary, mr_t=mr_t, solve_for_rho=True) -def unsaturated_hydrostatic_balance(state, theta_d, H, exner0=None, +def unsaturated_hydrostatic_balance(equation, theta_d, H, exner0=None, top=False, exner_boundary=Constant(1.0), max_outer_solve_count=40, max_inner_solve_count=20): @@ -423,8 +427,7 @@ def unsaturated_hydrostatic_balance(state, theta_d, H, exner0=None, These steps are iterated until we (hopefully) converge to a solution. Args: - state (:class:`State`): the model's state object, through which the - prognostic variables are accessed. + equation (:class:`PrognosticEquation`): the model's equation object. theta_d (:class:`ufl.Expr`): the specified dry potential temperature field. H (:class:`ufl.Expr`): the specified relative humidity field. @@ -449,18 +452,20 @@ def unsaturated_hydrostatic_balance(state, theta_d, H, exner0=None, number of iterations. """ - theta0 = state.fields('theta') - rho0 = state.fields('rho') - mr_v0 = state.fields('water_vapour') + theta0 = equation.fields('theta') + rho0 = equation.fields('rho') + mr_v0 = equation.fields('water_vapour') # Calculate hydrostatic exner pressure + domain = equation.domain + parameters = equation.parameters Vt = theta0.function_space() Vr = rho0.function_space() - R_d = state.parameters.R_d - R_v = state.parameters.R_v + R_d = parameters.R_d + R_v = parameters.R_v epsilon = R_d / R_v - VDG = state.spaces("DG") + VDG = domain.spaces("DG") if any(deg > 2 for deg in VDG.ufl_element().degree()): logger.warning("default quadrature degree most likely not sufficient for this degree element") @@ -480,21 +485,21 @@ def unsaturated_hydrostatic_balance(state, theta_d, H, exner0=None, delta = 1.0 # make expressions for determining mr_v0 - exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) - p = thermodynamics.p(state.parameters, exner) - T = thermodynamics.T(state.parameters, theta0, exner, mr_v0) - r_v_expr = thermodynamics.r_v(state.parameters, H, T, p) + exner = thermodynamics.exner_pressure(parameters, rho_averaged, theta0) + p = thermodynamics.p(parameters, exner) + T = thermodynamics.T(parameters, theta0, exner, mr_v0) + r_v_expr = thermodynamics.r_v(parameters, H, T, p) # make expressions to evaluate residual - exner_ev = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) - p_ev = thermodynamics.p(state.parameters, exner_ev) - T_ev = thermodynamics.T(state.parameters, theta0, exner_ev, mr_v0) - RH_ev = thermodynamics.RH(state.parameters, mr_v0, T_ev, p_ev) + exner_ev = thermodynamics.exner_pressure(parameters, rho_averaged, theta0) + p_ev = thermodynamics.p(parameters, exner_ev) + T_ev = thermodynamics.T(parameters, theta0, exner_ev, mr_v0) + RH_ev = thermodynamics.RH(parameters, mr_v0, T_ev, p_ev) RH = Function(Vt) for i in range(max_outer_solve_count): # solve for rho with theta_vd and w_v guesses - compressible_hydrostatic_balance(state, theta0, rho_h, top=top, + compressible_hydrostatic_balance(equation, theta0, rho_h, top=top, exner_boundary=exner_boundary, mr_t=mr_v0, solve_for_rho=True) @@ -525,10 +530,10 @@ def unsaturated_hydrostatic_balance(state, theta_d, H, exner0=None, raise RuntimeError('Hydrostatic balance solve has not converged within %i' % i, 'iterations') if exner0 is not None: - exner = thermodynamics.exner_pressure(state.parameters, rho0, theta0) + exner = thermodynamics.exner_pressure(parameters, rho0, theta0) exner0.interpolate(exner) # do one extra solve for rho - compressible_hydrostatic_balance(state, theta0, rho0, top=top, + compressible_hydrostatic_balance(equation, theta0, rho0, top=top, exner_boundary=exner_boundary, mr_t=mr_v0, solve_for_rho=True) diff --git a/gusto/io.py b/gusto/io.py new file mode 100644 index 000000000..f8fb7682e --- /dev/null +++ b/gusto/io.py @@ -0,0 +1,556 @@ +"""Provides the model's IO, which controls input, output and diagnostics.""" + +from os import path, makedirs +import itertools +from netCDF4 import Dataset +import sys +import time +from gusto.diagnostics import Diagnostics, Perturbation, SteadyStateError +from firedrake import (FiniteElement, TensorProductElement, VectorFunctionSpace, + interval, Function, Mesh, functionspaceimpl, File, + Constant, op2, DumbCheckpoint, FILE_CREATE, FILE_READ) +import numpy as np +from gusto.configuration import logger, set_log_handler + +__all__ = ["IO"] + + +class PointDataOutput(object): + """Object for outputting field point data.""" + def __init__(self, filename, ndt, field_points, description, + field_creator, comm, tolerance=None, create=True): + """ + Args: + filename (str): name of file to output to. + ndt (int): number of time points to output at. TODO: remove as this + is unused. + field_points (list): some iterable of pairs, matching fields with + arrays of evaluation points: (field_name, evaluation_points). + description (str): a description of the simulation to be included in + the output. + field_creator (:class:`FieldCreator`): the field creator, used to + determine the datatype and shape of fields. + comm (:class:`MPI.Comm`): MPI communicator. + tolerance (float, optional): tolerance to use for the evaluation of + fields at points. Defaults to None. + create (bool, optional): whether the output file needs creating, or + if it already exists. Defaults to True. + """ + # Overwrite on creation. + self.dump_count = 0 + self.filename = filename + self.field_points = field_points + self.tolerance = tolerance + self.comm = comm + if not create: + return + if self.comm.rank == 0: + with Dataset(filename, "w") as dataset: + dataset.description = "Point data for simulation {desc}".format(desc=description) + dataset.history = "Created {t}".format(t=time.ctime()) + # FIXME add versioning information. + dataset.source = "Output from Gusto model" + # Appendable dimension, timesteps in the model + dataset.createDimension("time", None) + + var = dataset.createVariable("time", np.float64, ("time")) + var.units = "seconds" + # Now create the variable group for each field + for field_name, points in field_points: + group = dataset.createGroup(field_name) + npts, dim = points.shape + group.createDimension("points", npts) + group.createDimension("geometric_dimension", dim) + var = group.createVariable("points", points.dtype, + ("points", "geometric_dimension")) + var[:] = points + + # Get the UFL shape of the field + field_shape = field_creator(field_name).ufl_shape + # Number of geometric dimension occurences should be the same as the length of the UFL shape + field_len = len(field_shape) + field_count = field_shape.count(dim) + assert field_len == field_count, "Geometric dimension occurrences do not match UFL shape" + # Create the variable with the required shape + dimensions = ("time", "points") + field_count*("geometric_dimension",) + group.createVariable(field_name, field_creator(field_name).dat.dtype, dimensions) + + def dump(self, field_creator, t): + """ + Evaluate and output field data at points. + + Args: + field_creator (:class:`FieldCreator`): gives access to the fields. + t (float): simulation time at which the output occurs. + """ + + val_list = [] + for field_name, points in self.field_points: + val_list.append((field_name, np.asarray(field_creator(field_name).at(points, tolerance=self.tolerance)))) + + if self.comm.rank == 0: + with Dataset(self.filename, "a") as dataset: + # Add new time index + dataset.variables["time"][self.dump_count] = t + for field_name, vals in val_list: + group = dataset.groups[field_name] + var = group.variables[field_name] + var[self.dump_count, :] = vals + + self.dump_count += 1 + + +class DiagnosticsOutput(object): + """Object for outputting global diagnostic data.""" + def __init__(self, filename, diagnostics, description, comm, create=True): + """ + Args: + filename (str): name of file to output to. + diagnostics (:class:`Diagnostics`): the object holding and + controlling the diagnostic evaluation. + description (str): a description of the simulation to be included in + the output. + comm (:class:`MPI.Comm`): MPI communicator. + create (bool, optional): whether the output file needs creating, or + if it already exists. Defaults to True. + """ + self.filename = filename + self.diagnostics = diagnostics + self.comm = comm + if not create: + return + if self.comm.rank == 0: + with Dataset(filename, "w") as dataset: + dataset.description = "Diagnostics data for simulation {desc}".format(desc=description) + dataset.history = "Created {t}".format(t=time.ctime()) + dataset.source = "Output from Gusto model" + dataset.createDimension("time", None) + var = dataset.createVariable("time", np.float64, ("time", )) + var.units = "seconds" + for name in diagnostics.fields: + group = dataset.createGroup(name) + for diagnostic in diagnostics.available_diagnostics: + group.createVariable(diagnostic, np.float64, ("time", )) + + def dump(self, equation, t): + """ + Output the global diagnostics. + + equation (:class:`PrognosticEquation`): the model's equation object. + t (float): simulation time at which the output occurs. + """ + + diagnostics = [] + for fname in self.diagnostics.fields: + field = equation.fields(fname) + for dname in self.diagnostics.available_diagnostics: + diagnostic = getattr(self.diagnostics, dname) + diagnostics.append((fname, dname, diagnostic(field))) + + if self.comm.rank == 0: + with Dataset(self.filename, "a") as dataset: + idx = dataset.dimensions["time"].size + dataset.variables["time"][idx:idx + 1] = t + for fname, dname, value in diagnostics: + group = dataset.groups[fname] + var = group.variables[dname] + var[idx:idx + 1] = value + + +class IO(object): + """Controls the model's input, output and diagnostics.""" + + def __init__(self, mesh, dt, + output=None, + parameters=None, + diagnostics=None, + diagnostic_fields=None): + """ + Args: + dt (:class:`Constant`): the time taken to perform a single model + step. If a float or int is passed, it will be cast to a + :class:`Constant`. + output (:class:`OutputParameters`, optional): holds and describes + the options for outputting. Defaults to None. + diagnostics (:class:`Diagnostics`, optional): object holding and + controlling the model's diagnostics. Defaults to None. + diagnostic_fields (list, optional): an iterable of `DiagnosticField` + objects. Defaults to None. + + Raises: + RuntimeError: if no output is provided. + TypeError: if `dt` cannot be cast to a :class:`Constant`. + """ + + if output is None: + # TODO: maybe this shouldn't be an optional argument then? + raise RuntimeError("You must provide a directory name for dumping results") + else: + self.output = output + self.parameters = parameters + + if diagnostics is not None: + self.diagnostics = diagnostics + else: + self.diagnostics = Diagnostics() + if diagnostic_fields is not None: + self.diagnostic_fields = diagnostic_fields + else: + self.diagnostic_fields = [] + + # TODO: quick way of ensuring that diagnostics are registered + if hasattr(self, "field_names"): + for fname in self.field_names: + self.diagnostics.register(fname) + self.bcs[fname] = [] + else: + self.diagnostics.register(field_name) + + # The mesh + self.mesh = mesh + + if self.output.dumplist is None: + self.output.dumplist = [] + + self.dumpdir = None + self.dumpfile = None + self.to_pickup = None + + # setup logger + logger.setLevel(output.log_level) + set_log_handler(mesh.comm) + if parameters is not None: + logger.info("Physical parameters that take non-default values:") + logger.info(", ".join("%s: %s" % (k, float(v)) for (k, v) in vars(parameters).items())) + + # Constant to hold current time + self.t = Constant(0.0) + if type(dt) is Constant: + self.dt = dt + elif type(dt) in (float, int): + self.dt = Constant(dt) + else: + raise TypeError(f'dt must be a Constant, float or int, not {type(dt)}') + + def setup_diagnostics(self): + """Concatenates the various types of diagnostic field.""" + for name in self.output.perturbation_fields: + f = Perturbation(name) + self.diagnostic_fields.append(f) + + for name in self.output.steady_state_error_fields: + f = SteadyStateError(self, name) + self.diagnostic_fields.append(f) + + fields = set([f.name() for f in self.fields]) + field_deps = [(d, sorted(set(d.required_fields).difference(fields),)) for d in self.diagnostic_fields] + schedule = topo_sort(field_deps) + self.diagnostic_fields = schedule + for diagnostic in self.diagnostic_fields: + diagnostic.setup(self) + self.diagnostics.register(diagnostic.name) + + def setup_dump(self, t, tmax, pickup=False): + """ + Sets up a series of things used for outputting. + + This prepares the model for outputting. First it checks for the + existence the specified outputting directory, so prevent it being + overwritten unintentionally. It then sets up the output files and the + checkpointing file. + + Args: + t (float): the current model time. + tmax (float): the end time of the model's simulation. + pickup (bool, optional): whether to pick up the model's initial + state from a checkpointing file. Defaults to False. + + Raises: + IOError: if the results directory already exists, and the model is + not picking up or running in test mode. + """ + + if any([self.output.dump_vtus, self.output.dumplist_latlon, + self.output.dump_diagnostics, self.output.point_data, + self.output.checkpoint and not pickup]): + # setup output directory and check that it does not already exist + self.dumpdir = path.join("results", self.output.dirname) + running_tests = '--running-tests' in sys.argv or "pytest" in self.output.dirname + if self.mesh.comm.rank == 0: + if not running_tests and path.exists(self.dumpdir) and not pickup: + raise IOError("results directory '%s' already exists" + % self.dumpdir) + else: + if not running_tests: + makedirs(self.dumpdir) + + if self.output.dump_vtus: + + # setup pvd output file + outfile = path.join(self.dumpdir, "field_output.pvd") + self.dumpfile = File( + outfile, project_output=self.output.project_fields, + comm=self.mesh.comm) + + # make list of fields to dump + self.to_dump = [f for f in self.fields if f.name() in self.fields.to_dump] + + # make dump counter + self.dumpcount = itertools.count() + + # if there are fields to be dumped in latlon coordinates, + # setup the latlon coordinate mesh and make output file + if len(self.output.dumplist_latlon) > 0: + mesh_ll = get_latlon_mesh(self.mesh) + outfile_ll = path.join(self.dumpdir, "field_output_latlon.pvd") + self.dumpfile_ll = File(outfile_ll, + project_output=self.output.project_fields, + comm=self.mesh.comm) + + # make functions on latlon mesh, as specified by dumplist_latlon + self.to_dump_latlon = [] + for name in self.output.dumplist_latlon: + f = self.fields(name) + field = Function( + functionspaceimpl.WithGeometry.create( + f.function_space(), mesh_ll), + val=f.topological, name=name+'_ll') + self.to_dump_latlon.append(field) + + # we create new netcdf files to write to, unless pickup=True, in + # which case we just need the filenames + if self.output.dump_diagnostics: + diagnostics_filename = self.dumpdir+"/diagnostics.nc" + self.diagnostic_output = DiagnosticsOutput(diagnostics_filename, + self.diagnostics, + self.output.dirname, + self.mesh.comm, + create=not pickup) + + if len(self.output.point_data) > 0: + # set up point data output + pointdata_filename = self.dumpdir+"/point_data.nc" + ndt = int(tmax/float(self.dt)) + self.pointdata_output = PointDataOutput(pointdata_filename, ndt, + self.output.point_data, + self.output.dirname, + self.fields, + self.mesh.comm, + self.output.tolerance, + create=not pickup) + + # make point data dump counter + self.pddumpcount = itertools.count() + + # set frequency of point data output - defaults to + # dumpfreq if not set by user + if self.output.pddumpfreq is None: + self.output.pddumpfreq = self.output.dumpfreq + + # if we want to checkpoint and are not picking up from a previous + # checkpoint file, setup the checkpointing + if self.output.checkpoint: + if not pickup: + self.chkpt = DumbCheckpoint(path.join(self.dumpdir, "chkpt"), + mode=FILE_CREATE) + # make list of fields to pickup (this doesn't include + # diagnostic fields) + self.to_pickup = [f for f in self.fields if f.name() in self.fields.to_pickup] + + # if we want to checkpoint then make a checkpoint counter + if self.output.checkpoint: + self.chkptcount = itertools.count() + + # dump initial fields + self.dump(t) + + def pickup_from_checkpoint(self): + """Picks up the model's variables from a checkpoint file.""" + # TODO: this duplicates some code from setup_dump. Can this be avoided? + # It is because we don't know if we are picking up or setting dump first + if self.to_pickup is None: + self.to_pickup = [f for f in self.fields if f.name() in self.fields.to_pickup] + # Set dumpdir if has not been done already + if self.dumpdir is None: + self.dumpdir = path.join("results", self.output.dirname) + + if self.output.checkpoint: + # Open the checkpointing file for writing + if self.output.checkpoint_pickup_filename is not None: + chkfile = self.output.checkpoint_pickup_filename + else: + chkfile = path.join(self.dumpdir, "chkpt") + with DumbCheckpoint(chkfile, mode=FILE_READ) as chk: + # Recover all the fields from the checkpoint + for field in self.to_pickup: + chk.load(field) + t = chk.read_attribute("/", "time") + # Setup new checkpoint + self.chkpt = DumbCheckpoint(path.join(self.dumpdir, "chkpt"), mode=FILE_CREATE) + else: + raise ValueError("Must set checkpoint True if pickup") + + return t + + def dump(self, t): + """ + Dumps all of the required model output. + + This includes point data, global diagnostics and general field data to + paraview data files. Also writes the model's prognostic variables to + a checkpoint file if specified. + + Args: + t (float): the simulation's current time. + """ + output = self.output + + # Diagnostics: + # Compute diagnostic fields + for field in self.diagnostic_fields: + field(self) + + if output.dump_diagnostics: + # Output diagnostic data + self.diagnostic_output.dump(self, t) + + if len(output.point_data) > 0 and (next(self.pddumpcount) % output.pddumpfreq) == 0: + # Output pointwise data + self.pointdata_output.dump(self.fields, t) + + # Dump all the fields to the checkpointing file (backup version) + if output.checkpoint and (next(self.chkptcount) % output.chkptfreq) == 0: + for field in self.to_pickup: + self.chkpt.store(field) + self.chkpt.write_attribute("/", "time", t) + + if output.dump_vtus and (next(self.dumpcount) % output.dumpfreq) == 0: + # dump fields + self.dumpfile.write(*self.to_dump) + + # dump fields on latlon mesh + if len(output.dumplist_latlon) > 0: + self.dumpfile_ll.write(*self.to_dump_latlon) + + def initialise(self, initial_conditions): + """ + Initialise the state's prognostic variables. + + Args: + initial_conditions (list): an iterable of pairs: (field_name, expr), + where 'field_name' is the string giving the name of the + prognostic field and expr is the :class:`ufl.Expr` whose value + is used to set the initial field. + """ + for name, ic in initial_conditions: + f_init = getattr(self.fields, name) + f_init.assign(ic) + f_init.rename(name) + +def get_latlon_mesh(mesh): + """ + Construct a planar latitude-longitude mesh from a spherical mesh. + + Args: + mesh (:class:`Mesh`): the mesh on which the simulation is performed. + """ + coords_orig = mesh.coordinates + coords_fs = coords_orig.function_space() + + if coords_fs.extruded: + cell = mesh._base_mesh.ufl_cell().cellname() + DG1_hori_elt = FiniteElement("DG", cell, 1, variant="equispaced") + DG1_vert_elt = FiniteElement("DG", interval, 1, variant="equispaced") + DG1_elt = TensorProductElement(DG1_hori_elt, DG1_vert_elt) + else: + cell = mesh.ufl_cell().cellname() + DG1_elt = FiniteElement("DG", cell, 1, variant="equispaced") + vec_DG1 = VectorFunctionSpace(mesh, DG1_elt) + coords_dg = Function(vec_DG1).interpolate(coords_orig) + coords_latlon = Function(vec_DG1) + shapes = {"nDOFs": vec_DG1.finat_element.space_dimension(), 'dim': 3} + + radius = np.min(np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2)) + # lat-lon 'x' = atan2(y, x) + coords_latlon.dat.data[:, 0] = np.arctan2(coords_dg.dat.data[:, 1], coords_dg.dat.data[:, 0]) + # lat-lon 'y' = asin(z/sqrt(x^2 + y^2 + z^2)) + coords_latlon.dat.data[:, 1] = np.arcsin(coords_dg.dat.data[:, 2]/np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2)) + # our vertical coordinate is radius - the minimum radius + coords_latlon.dat.data[:, 2] = np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2) - radius + +# We need to ensure that all points in a cell are on the same side of the branch cut in longitude coords +# This kernel amends the longitude coords so that all longitudes in one cell are close together + kernel = op2.Kernel(""" +#define PI 3.141592653589793 +#define TWO_PI 6.283185307179586 +void splat_coords(double *coords) {{ + double max_diff = 0.0; + double diff = 0.0; + + for (int i=0; i<{nDOFs}; i++) {{ + for (int j=0; j<{nDOFs}; j++) {{ + diff = coords[i*{dim}] - coords[j*{dim}]; + if (fabs(diff) > max_diff) {{ + max_diff = diff; + }} + }} + }} + + if (max_diff > PI) {{ + for (int i=0; i<{nDOFs}; i++) {{ + if (coords[i*{dim}] < 0) {{ + coords[i*{dim}] += TWO_PI; + }} + }} + }} +}} +""".format(**shapes), "splat_coords") + + op2.par_loop(kernel, coords_latlon.cell_set, + coords_latlon.dat(op2.RW, coords_latlon.cell_node_map())) + return Mesh(coords_latlon) + + +def topo_sort(field_deps): + """ + Perform a topological sort to determine the order to evaluate diagnostics. + + Args: + field_deps (list): a list of tuples, pairing diagnostic fields with the + fields that they are to be evaluated from. + + Raises: + RuntimeError: if there is a cyclic dependency in the diagnostic fields. + + Returns: + list: a list specifying the order in which to evaluate the diagnostics. + """ + name2field = dict((f.name, f) for f, _ in field_deps) + # map node: (input_deps, output_deps) + graph = dict((f.name, (list(deps), [])) for f, deps in field_deps) + roots = [] + for f, input_deps in field_deps: + if len(input_deps) == 0: + # No dependencies, candidate for evaluation + roots.append(f.name) + for d in input_deps: + # add f as output dependency + graph[d][1].append(f.name) + + schedule = [] + while roots: + n = roots.pop() + schedule.append(n) + output_deps = list(graph[n][1]) + for m in output_deps: + # Remove edge + graph[m][0].remove(n) + graph[n][1].remove(m) + # If m now as no input deps, candidate for evaluation + if len(graph[m][0]) == 0: + roots.append(m) + if any(len(i) for i, _ in graph.values()): + cycle = "\n".join("%s -> %s" % (f, i) for f, (i, _) in graph.items() + if f not in schedule) + raise RuntimeError("Field dependencies have a cycle:\n\n%s" % cycle) + return list(map(name2field.__getitem__, schedule)) diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index 3ddd5b13d..29dd344f2 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -28,11 +28,10 @@ class TimesteppingSolver(object, metaclass=ABCMeta): """Base class for timestepping linear solvers for Gusto.""" - def __init__(self, state, equations, alpha=0.5, solver_parameters=None, + def __init__(self, equations, alpha=0.5, solver_parameters=None, overwrite_solver_parameters=False): """ Args: - state (:class:`State`): the model's state object. equations (:class:`PrognosticEquation`): the model's equation. alpha (float, optional): the semi-implicit off-centring factor. Defaults to 0.5. A value of 1 is fully-implicit. @@ -44,7 +43,6 @@ def __init__(self, state, equations, alpha=0.5, solver_parameters=None, update the default parameters with the `solver_parameters` passed in. Defaults to False. """ - self.state = state self.equations = equations self.alpha = alpha @@ -122,12 +120,11 @@ class CompressibleSolver(TimesteppingSolver): 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}}} - def __init__(self, state, equations, alpha=0.5, + def __init__(self, equations, alpha=0.5, quadrature_degree=None, solver_parameters=None, overwrite_solver_parameters=False, moisture=None): """ Args: - state (:class:`State`): the model's state object. equations (:class:`PrognosticEquation`): the model's equation. alpha (float, optional): the semi-implicit off-centring factor. Defaults to 0.5. A value of 1 is fully-implicit. @@ -149,7 +146,7 @@ def __init__(self, state, equations, alpha=0.5, if quadrature_degree is not None: self.quadrature_degree = quadrature_degree else: - dgspace = state.spaces("DG") + dgspace = equations.domain.spaces("DG") if any(deg > 2 for deg in dgspace.ufl_element().degree()): logger.warning("default quadrature degree most likely not sufficient for this degree element") self.quadrature_degree = (5, 5) @@ -164,20 +161,19 @@ def __init__(self, state, equations, alpha=0.5, # Turn monitor on for the trace system self.solver_parameters["condensed_field"]["ksp_monitor_true_residual"] = None - super().__init__(state, equations, alpha, solver_parameters, + super().__init__(equations, alpha, solver_parameters, overwrite_solver_parameters) @timed_function("Gusto:SolverSetup") def _setup_solver(self): - state = self.state dt = state.dt beta_ = dt*self.alpha - cp = state.parameters.cp - Vu = state.spaces("HDiv") - Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element())) - Vtheta = state.spaces("theta") - Vrho = state.spaces("DG") + cp = equations.parameters.cp + Vu = equations.domain.spaces("HDiv") + Vu_broken = FunctionSpace(equations.domain.mesh, BrokenElement(Vu.ufl_element())) + Vtheta = equations.domain.spaces("theta") + Vrho = equations.domain.spaces("DG") # Store time-stepping coefficients as UFL Constants beta = Constant(beta_) @@ -185,7 +181,7 @@ def _setup_solver(self): h_deg = Vrho.ufl_element().degree()[0] v_deg = Vrho.ufl_element().degree()[1] - Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg)) + Vtrace = FunctionSpace(equations.domain.mesh, "HDiv Trace", degree=(h_deg, v_deg)) # Split up the rhs vector (symbolically) self.xrhs = Function(self.equations.function_space) @@ -196,17 +192,17 @@ def _setup_solver(self): w, phi, dl = TestFunctions(M) u, rho, l0 = TrialFunctions(M) - n = FacetNormal(state.mesh) + n = FacetNormal(equations.domin.mesh) # Get background fields - thetabar = state.fields("thetabar") - rhobar = state.fields("rhobar") - exnerbar = thermodynamics.exner_pressure(state.parameters, rhobar, thetabar) - exnerbar_rho = thermodynamics.dexner_drho(state.parameters, rhobar, thetabar) - exnerbar_theta = thermodynamics.dexner_dtheta(state.parameters, rhobar, thetabar) + thetabar = equations.fields("thetabar") + rhobar = equations.fields("rhobar") + exnerbar = thermodynamics.exner_pressure(equations.parameters, rhobar, thetabar) + exnerbar_rho = thermodynamics.dexner_drho(equations.parameters, rhobar, thetabar) + exnerbar_theta = thermodynamics.dexner_dtheta(equations.parameters, rhobar, thetabar) # Analytical (approximate) elimination of theta - k = state.k # Upward pointing unit vector + k = equations.domain.k # Upward pointing unit vector theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in # Only include theta' (rather than exner') in the vertical @@ -232,10 +228,11 @@ def V(u): + ds_b(degree=(self.quadrature_degree))) # Add effect of density of water upon theta + # TODO: this has to be done using active tracers if self.moisture is not None: water_t = Function(Vtheta).assign(0.0) for water in self.moisture: - water_t += self.state.fields(water) + water_t += self.equations.fields(water) theta_w = theta / (1 + water_t) thetabar_w = thetabar / (1 + water_t) else: @@ -405,13 +402,6 @@ class IncompressibleSolver(TimesteppingSolver): (1) Analytically eliminate b (introduces error near topography) (2) Solve resulting system for (u,p) using a hybrid-mixed method (3) Reconstruct b - - :arg state: a :class:`.State` object containing everything else. - :arg solver_parameters: (optional) Solver parameters. - :arg overwrite_solver_parameters: boolean, if True use only the - solver_parameters that have been passed in, if False then - update the default solver parameters with the solver_parameters - passed in. """ solver_parameters = { @@ -430,12 +420,12 @@ class IncompressibleSolver(TimesteppingSolver): @timed_function("Gusto:SolverSetup") def _setup_solver(self): - state = self.state # just cutting down line length a bit + equation = self.equation # just cutting down line length a bit dt = state.dt beta_ = dt*self.alpha - Vu = state.spaces("HDiv") - Vb = state.spaces("theta") - Vp = state.spaces("DG") + Vu = equation.domain.spaces("HDiv") + Vb = equation.domain.spaces("theta") + Vp = equation.domain.spaces("DG") # Store time-stepping coefficients as UFL Constants beta = Constant(beta_) @@ -450,10 +440,10 @@ def _setup_solver(self): u, p = TrialFunctions(M) # Get background fields - bbar = state.fields("bbar") + bbar = equation.fields("bbar") # Analytical (approximate) elimination of theta - k = state.k # Upward pointing unit vector + k = equation.domain.k # Upward pointing unit vector b = -dot(k, u)*dot(k, grad(bbar))*beta + b_in # vertical projection @@ -570,7 +560,7 @@ def __init__(self, equation, alpha): lambda t: Term(t.get(linearisation).form, t.labels), drop) - dt = equation.state.dt + dt = state.dt W = equation.function_space beta = dt*alpha diff --git a/gusto/physics.py b/gusto/physics.py index 7685e4dc7..dac6e69de 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -56,13 +56,11 @@ class SaturationAdjustment(Physics): A filter is applied to prevent generation of negative mixing ratios. """ - def __init__(self, equation, parameters, vapour_name='water_vapour', + def __init__(self, equation, vapour_name='water_vapour', cloud_name='cloud_water', latent_heat=True): """ Args: equation (:class:`PrognosticEquationSet`): the model's equation. - parameters (:class:`Configuration`): an object containing the - model's physical parameters. vapour_name (str, optional): name of the water vapour variable. Defaults to 'water_vapour'. cloud_name (str, optional): name of the cloud water variable. @@ -87,6 +85,7 @@ def __init__(self, equation, parameters, vapour_name='water_vapour', # Make prognostic for physics scheme self.X = Function(equation.X.function_space()) self.equation = equation + parameters = equation.parameters self.latent_heat = latent_heat # Vapour and cloud variables are needed for every form of this scheme @@ -242,13 +241,14 @@ class Fallout(Physics): for Cartesian geometry. """ - def __init__(self, equation, rain_name, state, moments=AdvectedMoments.M3): + def __init__(self, equation, rain_name, domain, moments=AdvectedMoments.M3): """ Args: equation (:class:`PrognosticEquationSet`): the model's equation. rain_name (str, optional): name of the rain variable. Defaults to 'rain'. - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. 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`. @@ -265,14 +265,15 @@ def __init__(self, equation, rain_name, state, moments=AdvectedMoments.M3): rain = self.X.split()[rain_idx] test = equation.tests[rain_idx] - Vu = state.spaces("HDiv") - v = state.fields('rainfall_velocity', Vu) + Vu = domain.spaces("HDiv") + # TODO: how do we allow this to be output? + v = Function(Vu, name='rainfall_velocity') # -------------------------------------------------------------------- # # Create physics term -- which is actually a transport term # -------------------------------------------------------------------- # - adv_term = advection_form(state, test, rain, outflow=True) + adv_term = advection_form(domain, test, rain, outflow=True) # Add rainfall velocity by replacing transport_velocity in term adv_term = adv_term.label_map(identity, map_if_true=lambda t: Term( @@ -289,7 +290,7 @@ def __init__(self, equation, rain_name, state, moments=AdvectedMoments.M3): if moments == AdvectedMoments.M0: # all rain falls at terminal velocity terminal_velocity = Constant(5) # in m/s - v.project(-terminal_velocity*state.k) + v.project(-terminal_velocity*domain.k) elif moments == AdvectedMoments.M3: # this advects the third moment M3 of the raindrop # distribution, which corresponds to the mean mass @@ -323,7 +324,7 @@ def __init__(self, equation, rain_name, state, moments=AdvectedMoments.M3): raise NotImplementedError('Currently we only have implementations for zero and one moment schemes for rainfall. Valid options are AdvectedMoments.M0 and AdvectedMoments.M3') if moments != AdvectedMoments.M0: - self.determine_v = Projector(-v_expression*state.k, v) + self.determine_v = Projector(-v_expression*domain.k, v) def evaluate(self, x_in, dt): """ @@ -450,13 +451,11 @@ class EvaporationOfRain(Physics): is the virtual dry potential temperature. """ - def __init__(self, equation, parameters, rain_name='rain', - vapour_name='water_vapour', latent_heat=True): + def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', + latent_heat=True): """ Args: equation (:class:`PrognosticEquationSet`): the model's equation. - parameters (:class:`Configuration`): an object containing the - model's physical parameters. cloud_name (str, optional): name of the rain variable. Defaults to 'rain'. vapour_name (str, optional): name of the water vapour variable. @@ -480,6 +479,7 @@ def __init__(self, equation, parameters, rain_name='rain', # Make prognostic for physics scheme self.X = Function(equation.X.function_space()) self.equation = equation + parameters = equation.parameters self.latent_heat = latent_heat # Vapour and cloud variables are needed for every form of this scheme @@ -616,8 +616,7 @@ class InstantRain(object): """ def __init__(self, equation, saturation_curve, vapour_name="water_vapour", - rain_name=None, parameters=None, convective_feedback=False, - set_tau_to_dt=False): + rain_name=None, convective_feedback=False, set_tau_to_dt=False): """ Args: equation (:class: 'PrognosticEquationSet'): the model's equation. @@ -627,9 +626,6 @@ def __init__(self, equation, saturation_curve, vapour_name="water_vapour", Defaults to "water_vapour". rain_name (str, optional): name of the rain variable. Defaults to None. - parameters (:class: 'Configuration', optional): an object - containing the model's physical parameters. Defaults to None - but required if convective_feedback is True. convective_feedback (bool, optional): True if the conversion of vapour affects the height equation. Defaults to False. set_tau_to_dt (bool, optional): True if the timescale for the @@ -638,6 +634,7 @@ def __init__(self, equation, saturation_curve, vapour_name="water_vapour", the parameters list. """ + parameters = equation.parameters self.convective_feedback = convective_feedback self.set_tau_to_dt = set_tau_to_dt @@ -650,7 +647,7 @@ def __init__(self, equation, saturation_curve, vapour_name="water_vapour", if self.convective_feedback: assert "D" in equation.field_names, "Depth field must exist for convective feedback" - assert parameters is not None, "You must provide parameters for convective feedback" + assert parameters.gamma is not None, "If convective feedback is used, gamma parameter must be specified" # obtain function space and functions; vapour needed for all cases W = equation.function_space diff --git a/gusto/transport_forms.py b/gusto/transport_forms.py index 3ea6a468b..ae9d14d7b 100644 --- a/gusto/transport_forms.py +++ b/gusto/transport_forms.py @@ -14,12 +14,13 @@ "advection_equation_circulation_form", "linear_continuity_form"] -def linear_advection_form(state, test, qbar): +def linear_advection_form(domain, test, qbar): """ The form corresponding to the linearised advective transport operator. Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. qbar (:class:`ufl.Expr`): the variable to be transported. @@ -27,22 +28,23 @@ def linear_advection_form(state, test, qbar): :class:`LabelledForm`: a labelled transport form. """ - ubar = Function(state.spaces("HDiv")) + ubar = Function(domain.spaces("HDiv")) # TODO: why is there a k here? - L = test*dot(ubar, state.k)*dot(state.k, grad(qbar))*dx + L = test*dot(ubar, domain.k)*dot(domain.k, grad(qbar))*dx form = transporting_velocity(L, ubar) return transport(form, TransportEquationType.advective) -def linear_continuity_form(state, test, qbar, facet_term=False): +def linear_continuity_form(domain, test, qbar, facet_term=False): """ The form corresponding to the linearised continuity transport operator. Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. qbar (:class:`ufl.Expr`): the variable to be transported. facet_term (bool, optional): whether to include interior facet terms. @@ -52,14 +54,14 @@ def linear_continuity_form(state, test, qbar, facet_term=False): :class:`LabelledForm`: a labelled transport form. """ - Vu = state.spaces("HDiv") + Vu = domain.spaces("HDiv") ubar = Function(Vu) L = qbar*test*div(ubar)*dx if facet_term: - n = FacetNormal(state.mesh) - Vu = state.spaces("HDiv") + n = FacetNormal(domain.mesh) + Vu = domain.spaces("HDiv") dS_ = (dS_v + dS_h) if Vu.extruded else dS L += jump(ubar*test, n)*avg(qbar)*dS_ @@ -68,7 +70,7 @@ def linear_continuity_form(state, test, qbar, facet_term=False): return transport(form, TransportEquationType.conservative) -def advection_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): +def advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): u""" The form corresponding to the advective transport operator. @@ -77,7 +79,8 @@ def advection_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): form is integrated by parts. Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. ibp (:class:`IntegrateByParts`, optional): an enumerator representing @@ -96,7 +99,7 @@ def advection_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): if outflow and ibp == IntegrateByParts.NEVER: raise ValueError("outflow is True and ibp is None are incompatible options") - Vu = state.spaces("HDiv") + Vu = domain.spaces("HDiv") dS_ = (dS_v + dS_h) if Vu.extruded else dS ubar = Function(Vu) @@ -106,7 +109,7 @@ def advection_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): L = inner(outer(test, ubar), grad(q))*dx if ibp != IntegrateByParts.NEVER: - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) L += dot(jump(test), (un('+')*q('+') - un('-')*q('-')))*dS_ @@ -116,7 +119,7 @@ def advection_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): + inner(test('-'), dot(ubar('-'), n('-'))*q('-')))*dS_ if outflow: - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) L += test*un*q*(ds_v + ds_t + ds_b) @@ -125,7 +128,7 @@ def advection_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): return ibp_label(transport(form, TransportEquationType.advective), ibp) -def continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): +def continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): u""" The form corresponding to the continuity transport operator. @@ -134,7 +137,8 @@ def continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): form is integrated by parts. Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. ibp (:class:`IntegrateByParts`, optional): an enumerator representing @@ -153,7 +157,7 @@ def continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): if outflow and ibp == IntegrateByParts.NEVER: raise ValueError("outflow is True and ibp is None are incompatible options") - Vu = state.spaces("HDiv") + Vu = domain.spaces("HDiv") dS_ = (dS_v + dS_h) if Vu.extruded else dS ubar = Function(Vu) @@ -163,7 +167,7 @@ def continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): L = inner(test, div(outer(q, ubar)))*dx if ibp != IntegrateByParts.NEVER: - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) L += dot(jump(test), (un('+')*q('+') - un('-')*q('-')))*dS_ @@ -173,7 +177,7 @@ def continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): + inner(test('-'), dot(ubar('-'), n('-'))*q('-')))*dS_ if outflow: - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) L += test*un*q*(ds_v + ds_t + ds_b) @@ -182,7 +186,7 @@ def continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): return ibp_label(transport(form, TransportEquationType.conservative), ibp) -def vector_manifold_advection_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): +def vector_manifold_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): """ Form for advective transport operator including vector manifold correction. @@ -192,7 +196,8 @@ def vector_manifold_advection_form(state, test, q, ibp=IntegrateByParts.ONCE, ou is based on that of Bernard, Remacle et al (2009). Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. ibp (:class:`IntegrateByParts`, optional): an enumerator representing @@ -205,13 +210,13 @@ def vector_manifold_advection_form(state, test, q, ibp=IntegrateByParts.ONCE, ou class:`LabelledForm`: a labelled transport form. """ - L = advection_form(state, test, q, ibp, outflow) + L = advection_form(domain, test, q, ibp, outflow) # TODO: there should maybe be a restriction on IBP here - Vu = state.spaces("HDiv") + Vu = domain.spaces("HDiv") dS_ = (dS_v + dS_h) if Vu.extruded else dS ubar = Function(Vu) - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) L += un('+')*inner(test('-'), n('+')+n('-'))*inner(q('+'), n('+'))*dS_ L += un('-')*inner(test('+'), n('+')+n('-'))*inner(q('-'), n('-'))*dS_ @@ -219,7 +224,7 @@ def vector_manifold_advection_form(state, test, q, ibp=IntegrateByParts.ONCE, ou return L -def vector_manifold_continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): +def vector_manifold_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): """ Form for continuity transport operator including vector manifold correction. @@ -229,7 +234,8 @@ def vector_manifold_continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, o is based on that of Bernard, Remacle et al (2009). Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. ibp (:class:`IntegrateByParts`, optional): an enumerator representing @@ -242,12 +248,12 @@ def vector_manifold_continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, o class:`LabelledForm`: a labelled transport form. """ - L = continuity_form(state, test, q, ibp, outflow) + L = continuity_form(domain, test, q, ibp, outflow) - Vu = state.spaces("HDiv") + Vu = domain.spaces("HDiv") dS_ = (dS_v + dS_h) if Vu.extruded else dS ubar = Function(Vu) - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) L += un('+')*inner(test('-'), n('+')+n('-'))*inner(q('+'), n('+'))*dS_ L += un('-')*inner(test('+'), n('+')+n('-'))*inner(q('-'), n('-'))*dS_ @@ -257,7 +263,7 @@ def vector_manifold_continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, o return transport(form) -def vector_invariant_form(state, test, q, ibp=IntegrateByParts.ONCE): +def vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE): u""" The form corresponding to the vector invariant transport operator. @@ -273,7 +279,8 @@ def vector_invariant_form(state, test, q, ibp=IntegrateByParts.ONCE): when integrating by parts. Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. ibp (:class:`IntegrateByParts`, optional): an enumerator representing @@ -287,13 +294,13 @@ def vector_invariant_form(state, test, q, ibp=IntegrateByParts.ONCE): class:`LabelledForm`: a labelled transport form. """ - Vu = state.spaces("HDiv") + Vu = domain.spaces("HDiv") dS_ = (dS_v + dS_h) if Vu.extruded else dS ubar = Function(Vu) - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) Upwind = 0.5*(sign(dot(ubar, n))+1) - if state.mesh.topological_dimension() == 3: + if domain.mesh.topological_dimension() == 3: if ibp != IntegrateByParts.ONCE: raise NotImplementedError @@ -313,9 +320,9 @@ def vector_invariant_form(state, test, q, ibp=IntegrateByParts.ONCE): else: - perp = state.perp - if state.on_sphere: - outward_normals = CellNormal(state.mesh) + perp = domain.perp + if domain.on_sphere: + outward_normals = CellNormal(domain.mesh) perp_u_upwind = lambda q: Upwind('+')*cross(outward_normals('+'), q('+')) + Upwind('-')*cross(outward_normals('-'), q('-')) else: perp_u_upwind = lambda q: Upwind('+')*perp(q('+')) + Upwind('-')*perp(q('-')) @@ -342,7 +349,7 @@ def vector_invariant_form(state, test, q, ibp=IntegrateByParts.ONCE): return transport(form, TransportEquationType.vector_invariant) -def kinetic_energy_form(state, test, q): +def kinetic_energy_form(domain, test, q): u""" The form corresponding to the kinetic energy term. @@ -351,7 +358,8 @@ def kinetic_energy_form(state, test, q): (1/2)∇(u.q). Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. @@ -359,7 +367,7 @@ def kinetic_energy_form(state, test, q): class:`LabelledForm`: a labelled transport form. """ - ubar = Function(state.spaces("HDiv")) + ubar = Function(domain.spaces("HDiv")) L = 0.5*div(test)*inner(q, ubar)*dx form = transporting_velocity(L, ubar) @@ -367,7 +375,7 @@ def kinetic_energy_form(state, test, q): return transport(form, TransportEquationType.vector_invariant) -def advection_equation_circulation_form(state, test, q, +def advection_equation_circulation_form(domain, test, q, ibp=IntegrateByParts.ONCE): u""" The circulation term in the transport of a vector-valued field. @@ -384,7 +392,8 @@ def advection_equation_circulation_form(state, test, q, term. An an upwind discretisation is used when integrating by parts. Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. ibp (:class:`IntegrateByParts`, optional): an enumerator representing @@ -399,8 +408,8 @@ def advection_equation_circulation_form(state, test, q, """ form = ( - vector_invariant_form(state, test, q, ibp=ibp) - - kinetic_energy_form(state, test, q) + vector_invariant_form(domain, test, q, ibp=ibp) + - kinetic_energy_form(domain, test, q) ) return form From ea40ff93809e9a74b8787b6df9a47a20d4c4e533 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Thu, 8 Dec 2022 19:40:27 +0000 Subject: [PATCH 03/12] get first transport tests working with new domain/io setup --- gusto/domain.py | 63 +++++-- gusto/equations.py | 68 +++----- gusto/fields.py | 13 +- gusto/function_spaces.py | 158 ++++++++++++------ gusto/io.py | 35 ++-- gusto/time_discretisation.py | 68 +++++--- gusto/timeloop.py | 69 ++++---- integration-tests/conftest.py | 14 +- .../transport/test_dg_transport.py | 48 +++--- 9 files changed, 314 insertions(+), 222 deletions(-) diff --git a/gusto/domain.py b/gusto/domain.py index c7e5667c0..ff8d9d8e0 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -3,29 +3,69 @@ the set of compatible function spaces defined upon it. """ -from gusto.spaces import Spaces +from gusto.function_spaces import Spaces from firedrake import (Constant, SpatialCoordinate, sqrt, CellNormal, cross, - as_vector) + as_vector, inner, interpolate) class Domain(object): - """The Domain holds the model's mesh and its compatible function spaces.""" - def __init__(self, mesh, family, degree): + """ + The Domain holds the model's mesh and its compatible function spaces. + + The compatible function spaces are given by the de Rham complex, and are + specified here through the family of the HDiv velocity space and the degree + of the DG space. + + For extruded meshes, it is possible to seperately specify the horizontal and + vertical degrees of the elements. Alternatively, if these degrees should be + the same then this can be specified through the "degree" argument. + """ + def __init__(self, mesh, family, degree=None, + horizontal_degree=None, vertical_degree=None): """ Args: mesh (:class:`Mesh`): the model's mesh. family (str): the finite element space family used for the velocity field. This determines the other finite element spaces used via the de Rham complex. - degree (int): the element degree used for the velocity space. + degree (int, optional): the element degree used for the DG space + Defaults to None, in which case the horizontal degree must be provided. + horizontal_degree (int, optional): the element degree used for the + horizontal part of the DG space. Defaults to None. + vertical_degree (int, optional): the element degree used for the + vertical part of the DG space. Defaults to None. + + Raises: + ValueError: if incompatible degrees are specified (e.g. specifying + both "degree" and "horizontal_degree"). """ + # Checks on degree arguments + if degree is None and horizontal_degree is None: + raise ValueError('Either "degree" or "horizontal_degree" must be passed to Domain') + if mesh.extruded and degree is None and vertical_degree is None: + raise ValueError('For extruded meshes, either degree or "vertical_degree" must be passed to Domain') + if degree is not None and horizontal_degree is not None: + raise ValueError('Cannot pass both "degree" and "horizontal_degree" to Domain') + if mesh.extruded and degree is not None and vertical_degree is not None: + raise ValueError('Cannot pass both "degree" and "vertical_degree" to Domain') + if not mesh.extruded and vertical_degree is not None: + raise ValueError('Cannot pass "vertical_degree" to Domain if mesh is not extruded') + + # Get degrees + self.horizontal_degree = degree if horizontal_degree is None else horizontal_degree + self.vertical_degree = degree if vertical_degree is None else vertical_degree + self.mesh = mesh - self.spaces = [space for space in self._build_spaces(state, family, degree)] + self.family = family + self.spaces = Spaces(mesh) + # Build and store compatible spaces + self.compatible_spaces = [space for space in self.spaces.build_compatible_spaces(self.family, self.horizontal_degree, self.vertical_degree)] - # figure out if we're on a sphere - try: + # Figure out if we're on a sphere + # TODO: could we run on other domains that could confuse this? + if hasattr(mesh, "_base_mesh"): self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3 and mesh._base_mesh.topological_dimension() == 2) - except AttributeError: + else: self.on_sphere = (mesh.geometric_dimension() == 3 and mesh.topological_dimension() == 2) # build the vertical normal and define perp for 2d geometries @@ -43,8 +83,3 @@ def __init__(self, mesh, family, degree): self.k = Constant(kvec) if dim == 2: self.perp = lambda u: as_vector([-u[1], u[0]]) - - # TODO: why have this as a separate routine? - def _build_spaces(self, family, degree): - spaces = Spaces(self.mesh) - return spaces.build_compatible_spaces(family, degree) diff --git a/gusto/equations.py b/gusto/equations.py index 314516448..4318f800d 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -60,8 +60,7 @@ def __init__(self, domain, function_space, field_name): class AdvectionEquation(PrognosticEquation): u"""Discretises the advection equation, ∂q/∂t + (u.∇)q = 0""" - def __init__(self, domain, function_space, field_name, - ufamily=None, udegree=None, Vu=None, **kwargs): + def __init__(self, domain, function_space, field_name, Vu=None, **kwargs): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -69,14 +68,9 @@ def __init__(self, domain, function_space, field_name, function_space (:class:`FunctionSpace`): the function space that the equation's prognostic is defined on. field_name (str): name of the prognostic field. - ufamily (str, optional): the family of the function space to use - for the velocity field. Only used if `Vu` is not provided. - Defaults to None. - udegree (int, optional): the degree of the function space to use for - the velocity field. Only used if `Vu` is not provided. Defaults - to None. Vu (:class:`FunctionSpace`, optional): the function space for the - velocity field. If this is Defaults to None. + velocity field. If this is not specified, uses the HDiv spaces + set up by the domain. Defaults to None. **kwargs: any keyword arguments to be passed to the advection form. """ super().__init__(domain, function_space, field_name) @@ -85,9 +79,7 @@ def __init__(self, domain, function_space, field_name, if Vu is not None: V = domain.spaces("HDiv", V=Vu) else: - assert ufamily is not None, "Specify the family for u" - assert udegree is not None, "Specify the degree of the u space" - V = domain.spaces("HDiv", ufamily, udegree) + V = domain.spaces("HDiv") self.fields("u", V) test = TestFunction(function_space) q = Function(function_space) @@ -101,8 +93,7 @@ def __init__(self, domain, function_space, field_name, class ContinuityEquation(PrognosticEquation): u"""Discretises the continuity equation, ∂q/∂t + ∇(u*q) = 0""" - def __init__(self, domain, function_space, field_name, - ufamily=None, udegree=None, Vu=None, **kwargs): + def __init__(self, domain, function_space, field_name, Vu=None, **kwargs): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -110,14 +101,9 @@ def __init__(self, domain, function_space, field_name, function_space (:class:`FunctionSpace`): the function space that the equation's prognostic is defined on. field_name (str): name of the prognostic field. - ufamily (str, optional): the family of the function space to use - for the velocity field. Only used if `Vu` is not provided. - Defaults to None. - udegree (int, optional): the degree of the function space to use for - the velocity field. Only used if `Vu` is not provided. Defaults - to None. Vu (:class:`FunctionSpace`, optional): the function space for the - velocity field. If this is Defaults to None. + velocity field. If this is not specified, uses the HDiv spaces + set up by the domain. Defaults to None. **kwargs: any keyword arguments to be passed to the advection form. """ super().__init__(domain, function_space, field_name) @@ -126,9 +112,7 @@ def __init__(self, domain, function_space, field_name, if Vu is not None: V = domain.spaces("HDiv", V=Vu) else: - assert ufamily is not None, "Specify the family for u" - assert udegree is not None, "Specify the degree of the u space" - V = domain.spaces("HDiv", ufamily, udegree) + V = domain.spaces("HDiv") self.fields("u", V) test = TestFunction(function_space) q = Function(function_space) @@ -170,9 +154,8 @@ def __init__(self, domain, function_space, field_name, class AdvectionDiffusionEquation(PrognosticEquation): u"""The advection-diffusion equation, ∂q/∂t + (u.∇)q = ∇.(κ∇q)""" - def __init__(self, domain, function_space, field_name, - ufamily=None, udegree=None, Vu=None, diffusion_parameters=None, - **kwargs): + def __init__(self, domain, function_space, field_name, Vu=None, + diffusion_parameters=None, **kwargs): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -180,12 +163,6 @@ def __init__(self, domain, function_space, field_name, function_space (:class:`FunctionSpace`): the function space that the equation's prognostic is defined on. field_name (str): name of the prognostic field. - ufamily (str, optional): the family of the function space to use - for the velocity field. Only used if `Vu` is not provided. - Defaults to None. - udegree (int, optional): the degree of the function space to use for - the velocity field. Only used if `Vu` is not provided. Defaults - to None. Vu (:class:`FunctionSpace`, optional): the function space for the velocity field. If this is Defaults to None. diffusion_parameters (:class:`DiffusionParameters`, optional): @@ -199,9 +176,7 @@ def __init__(self, domain, function_space, field_name, if Vu is not None: V = domain.spaces("HDiv", V=Vu) else: - assert ufamily is not None, "Specify the family for u" - assert udegree is not None, "Specify the degree of the u space" - V = domain.spaces("HDiv", ufamily, udegree) + V = domain.spaces("HDiv") self.fields("u", V) test = TestFunction(function_space) q = Function(function_space) @@ -247,6 +222,9 @@ def __init__(self, field_names, domain, linearisation_map=None, self.linearisation_map = lambda t: False if linearisation_map is None else linearisation_map(t) self.reference_profiles_initialised = False + # Build finite element spaces + self.spaces = domain.compatible_spaces + # Add active tracers to the list of prognostics if active_tracers is None: active_tracers = [] @@ -516,9 +494,8 @@ class ForcedAdvectionEquation(PrognosticEquationSet): ∂q/∂t + (u.∇)q = F, which can also be augmented with active tracers. """ - def __init__(self, domain, function_space, field_name, - ufamily=None, udegree=None, Vu=None, active_tracers=None, - **kwargs): + def __init__(self, domain, function_space, field_name, Vu=None, + active_tracers=None, **kwargs): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -526,14 +503,9 @@ def __init__(self, domain, function_space, field_name, function_space (:class:`FunctionSpace`): the function space that the equation's prognostic is defined on. field_name (str): name of the prognostic field. - ufamily (str, optional): the family of the function space to use - for the velocity field. Only used if `Vu` is not provided. - Defaults to None. - udegree (int, optional): the degree of the function space to use for - the velocity field. Only used if `Vu` is not provided. Defaults - to None. Vu (:class:`FunctionSpace`, optional): the function space for the - velocity field. If this is Defaults to None. + velocity field. If this is not specified, uses the HDiv spaces + set up by the domain. 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. @@ -563,9 +535,7 @@ def __init__(self, domain, function_space, field_name, if Vu is not None: V = domain.spaces("HDiv", V=Vu) else: - assert ufamily is not None, "Specify the family for u" - assert udegree is not None, "Specify the degree of the u space" - V = domain.spaces("HDiv", ufamily, udegree) + V = domain.spaces("HDiv") self.fields("u", V) self.tests = TestFunctions(W) diff --git a/gusto/fields.py b/gusto/fields.py index ad231f39d..63e8b9bc6 100644 --- a/gusto/fields.py +++ b/gusto/fields.py @@ -58,7 +58,7 @@ def __iter__(self): class StateFields(Fields): - """Creates the prognostic fields for the :class:`State` object.""" + """Creates the prognostic fields for the model's equation.""" def __init__(self, *fields_to_dump): """ @@ -149,14 +149,17 @@ def add_fields(self, equation, levels=None): except AttributeError: setattr(self, level, Fields(equation)) - def initialise(self, state): + def initialise(self, equation): + # TODO: should this be IO? """ - Initialises the time fields from those currently in state + Initialises the time fields from those currently in the equation. - Args: state (:class:`State`): the model state object + Args: + equation (:class:`PrognosticEquation`): the model's prognostic + equation object. """ for field in self.n: - field.assign(state.fields(field.name())) + field.assign(equation.fields(field.name())) self.np1(field.name()).assign(field) def update(self): diff --git a/gusto/function_spaces.py b/gusto/function_spaces.py index 8d1f1f066..71f12ee4a 100644 --- a/gusto/function_spaces.py +++ b/gusto/function_spaces.py @@ -6,6 +6,13 @@ from firedrake import (HDiv, FunctionSpace, FiniteElement, TensorProductElement, interval) +# TODO: there is danger here for confusion about degree, particularly for the CG +# spaces -- does a "CG" space with degree = 1 mean the "CG" space in the de Rham +# complex of degree 1 ("CG3"), or "CG1"? +# TODO: would it be better to separate creation of specific named spaces from +# the creation of the de Rham complex spaces? +# TODO: how do we create HCurl spaces if we want them? + class Spaces(object): """Object to create and hold the model's finite element spaces.""" def __init__(self, mesh): @@ -17,7 +24,8 @@ def __init__(self, mesh): self.extruded_mesh = hasattr(mesh, "_base_mesh") self._initialised_base_spaces = False - def __call__(self, name, family=None, degree=None, V=None): + def __call__(self, name, family=None, horizontal_degree=None, + vertical_degree=None, V=None): """ Returns a space, and also creates it if it is not created yet. @@ -29,8 +37,10 @@ def __call__(self, name, family=None, degree=None, V=None): name (str): the name of the space. family (str, optional): name of the finite element family to be created. Defaults to None. - degree (int, optional): the degree of the finite element space to be - created. Defaults to None. + horizontal_degree (int, optional): the horizontal degree of the + finite element space to be created. Defaults to None. + vertical_degree (int, optional): the vertical degree of the + finite element space to be created. Defaults to None. V (:class:`FunctionSpace`, optional): an existing space, to be stored in the creator object. If this is provided, it will be added to the creator and no other action will be taken. This @@ -41,26 +51,38 @@ def __call__(self, name, family=None, degree=None, V=None): """ try: + # First attempt to return the space based on the name, if it exists return getattr(self, name) + except AttributeError: + + # Space does not exist in creator if V is not None: + # The space itself has been provided (to add it to the creator) value = V - elif name == "HDiv" and family in ["BDM", "RT", "CG", "RTCF"]: - value = self.build_hdiv_space(family, degree) - elif name == "theta": - value = self.build_theta_space(degree) - elif name == "DG1_equispaced": - value = self.build_dg_space(1, variant='equispaced') - elif family == "DG": - value = self.build_dg_space(degree) - elif family == "CG": - value = self.build_cg_space(degree) + else: - raise ValueError(f'State has no space corresponding to {name}') + # Need to create space, based on name/family/degree + assert horizontal_degree is not None + + # Loop through name and family combinations + if name == "HDiv" and family in ["BDM", "RT", "CG", "RTCF"]: + value = self.build_hdiv_space(family, horizontal_degree, vertical_degree) + elif name == "theta": + value = self.build_theta_space(horizontal_degree, vertical_degree) + elif name == "DG1_equispaced": + value = self.build_dg_space(1, variant='equispaced') + elif family == "DG": + value = self.build_dg_space(horizontal_degree, vertical_degree) + elif family == "CG": + value = self.build_cg_space(horizontal_degree, vertical_degree) + else: + raise ValueError(f'State has no space corresponding to {name}') setattr(self, name, value) return value - def build_compatible_spaces(self, family, degree): + def build_compatible_spaces(self, family, horizontal_degree, + vertical_degree=None): """ Builds the sequence of compatible finite element spaces for the mesh. @@ -72,76 +94,93 @@ def build_compatible_spaces(self, family, degree): Args: family (str): the family of the horizontal part of the HDiv space. - degree (int): the polynomial degree of the DG space. + horizontal_degree (int): the polynomial degree of the horizontal + part of the DG space. + vertical_degree (int, optional): the polynomial degree of the + vertical part of the DG space. Defaults to None. Must be + specified if the mesh is extruded. Returns: tuple: the created compatible :class:`FunctionSpace` objects. """ if self.extruded_mesh and not self._initialised_base_spaces: - self.build_base_spaces(family, degree) - Vu = self.build_hdiv_space(family, degree) + self.build_base_spaces(family, horizontal_degree, vertical_degree) + Vu = self.build_hdiv_space(family, horizontal_degree, vertical_degree) setattr(self, "HDiv", Vu) - Vdg = self.build_dg_space(degree) + Vdg = self.build_dg_space(horizontal_degree, vertical_degree) setattr(self, "DG", Vdg) - Vth = self.build_theta_space(degree) + Vth = self.build_theta_space(horizontal_degree, vertical_degree) setattr(self, "theta", Vth) return Vu, Vdg, Vth else: - Vu = self.build_hdiv_space(family, degree) + Vu = self.build_hdiv_space(family, horizontal_degree+1) setattr(self, "HDiv", Vu) - Vdg = self.build_dg_space(degree) + Vdg = self.build_dg_space(horizontal_degree, vertical_degree) setattr(self, "DG", Vdg) return Vu, Vdg - def build_base_spaces(self, family, degree): + def build_base_spaces(self, family, horizontal_degree, vertical_degree): """ Builds the :class:`FiniteElement` objects for the base mesh. Args: family (str): the family of the horizontal part of the HDiv space. - degree (int): the polynomial degree of the DG space. + horizontal_degree (int): the polynomial degree of the horizontal + part of the DG space. + vertical_degree (int): the polynomial degree of the vertical part of + the DG space. """ cell = self.mesh._base_mesh.ufl_cell().cellname() # horizontal base spaces - self.S1 = FiniteElement(family, cell, degree+1) - self.S2 = FiniteElement("DG", cell, degree) + self.S1 = FiniteElement(family, cell, horizontal_degree+1) + self.S2 = FiniteElement("DG", cell, horizontal_degree) # vertical base spaces - self.T0 = FiniteElement("CG", interval, degree+1) - self.T1 = FiniteElement("DG", interval, degree) + self.T0 = FiniteElement("CG", interval, vertical_degree+1) + self.T1 = FiniteElement("DG", interval, vertical_degree) self._initialised_base_spaces = True - def build_hdiv_space(self, family, degree): + def build_hdiv_space(self, family, horizontal_degree, vertical_degree=None): """ Builds and returns the HDiv :class:`FunctionSpace`. Args: family (str): the family of the horizontal part of the HDiv space. - degree (int): the polynomial degree of the space. + horizontal_degree (int): the polynomial degree of the horizontal + part of the DG space from the de Rham complex. + vertical_degree (int, optional): the polynomial degree of the + vertical part of the the DG space from the de Rham complex. + Defaults to None. Must be specified if the mesh is extruded. Returns: :class:`FunctionSpace`: the HDiv space. """ if self.extruded_mesh: if not self._initialised_base_spaces: - self.build_base_spaces(family, degree) + if vertical_degree is None: + raise ValueError('vertical_degree must be specified to create HDiv space on an extruded mesh') + self.build_base_spaces(family, horizontal_degree, vertical_degree) Vh_elt = HDiv(TensorProductElement(self.S1, self.T1)) Vt_elt = TensorProductElement(self.S2, self.T0) Vv_elt = HDiv(Vt_elt) V_elt = Vh_elt + Vv_elt else: cell = self.mesh.ufl_cell().cellname() - V_elt = FiniteElement(family, cell, degree+1) + V_elt = FiniteElement(family, cell, horizontal_degree) return FunctionSpace(self.mesh, V_elt, name='HDiv') - def build_dg_space(self, degree, variant=None): + def build_dg_space(self, horizontal_degree, vertical_degree=None, variant=None): """ Builds and returns the DG :class:`FunctionSpace`. Args: - degree (int): the polynomial degree of the space. + horizontal_degree (int): the polynomial degree of the horizontal + part of the DG space. + vertical_degree (int, optional): the polynomial degree of the + vertical part of the DG space. Defaults to None. Must be + specified if the mesh is extruded. variant (str): the variant of the underlying :class:`FiniteElement` to use. Defaults to None, which will call the default variant. @@ -149,21 +188,24 @@ def build_dg_space(self, degree, variant=None): :class:`FunctionSpace`: the DG space. """ if self.extruded_mesh: - if not self._initialised_base_spaces or self.T1.degree() != degree or self.T1.variant() != variant: + if vertical_degree is None: + raise ValueError('vertical_degree must be specified to create DG space on an extruded mesh') + if not self._initialised_base_spaces or self.T1.degree() != vertical_degree or self.T1.variant() != variant: cell = self.mesh._base_mesh.ufl_cell().cellname() - S2 = FiniteElement("DG", cell, degree, variant=variant) - T1 = FiniteElement("DG", interval, degree, variant=variant) + S2 = FiniteElement("DG", cell, horizontal_degree, variant=variant) + T1 = FiniteElement("DG", interval, vertical_degree, variant=variant) else: S2 = self.S2 T1 = self.T1 V_elt = TensorProductElement(S2, T1) else: cell = self.mesh.ufl_cell().cellname() - V_elt = FiniteElement("DG", cell, degree, variant=variant) - name = f'DG{degree}_equispaced' if variant == 'equispaced' else f'DG{degree}' + V_elt = FiniteElement("DG", cell, horizontal_degree, variant=variant) + # TODO: how should we name this if vertical degree is different? + name = f'DG{horizontal_degree}_equispaced' if variant == 'equispaced' else f'DG{horizontal_degree}' return FunctionSpace(self.mesh, V_elt, name=name) - def build_theta_space(self, degree): + def build_theta_space(self, horizontal_degree, vertical_degree): """ Builds and returns the 'theta' space. @@ -172,7 +214,10 @@ def build_theta_space(self, degree): continuous in the vertical. Args: - degree (int): degree of the corresponding density space. + horizontal_degree (int): the polynomial degree of the horizontal + part of the DG space from the de Rham complex. + vertical_degree (int): the polynomial degree of the vertical part of + the DG space from the de Rham complex. Raises: AssertionError: the mesh is not extruded. @@ -180,22 +225,41 @@ def build_theta_space(self, degree): Returns: :class:`FunctionSpace`: the 'theta' space. """ - assert self.extruded_mesh + assert self.extruded_mesh, 'Cannot create theta space if mesh is not extruded' if not self._initialised_base_spaces: cell = self.mesh._base_mesh.ufl_cell().cellname() - self.S2 = FiniteElement("DG", cell, degree) - self.T0 = FiniteElement("CG", interval, degree+1) + self.S2 = FiniteElement("DG", cell, horizontal_degree) + self.T0 = FiniteElement("CG", interval, vertical_degree+1) V_elt = TensorProductElement(self.S2, self.T0) return FunctionSpace(self.mesh, V_elt, name='Vtheta') - def build_cg_space(self, degree): + def build_cg_space(self, horizontal_degree, vertical_degree): """ Builds the continuous scalar space at the top of the de Rham complex. Args: - degree (int): degree of the continuous space. + horizontal_degree (int): the polynomial degree of the horizontal + part of the DG space from the de Rham complex. + vertical_degree (int, optional): the polynomial degree of the + vertical part of the the DG space from the de Rham complex. + Defaults to None. Must be specified if the mesh is extruded. Returns: :class:`FunctionSpace`: the continuous space. """ - return FunctionSpace(self.mesh, "CG", degree, name=f'CG{degree}') + + if self.extruded_mesh: + if vertical_degree is None: + raise ValueError('vertical_degree must be specified to create CG space on an extruded mesh') + cell = self.mesh._base_mesh.ufl_cell().cellname() + CG_hori = FiniteElement("CG", cell, horizontal_degree+2) + CG_vert = FiniteElement("CG", interval, vertical_degree+2) + V_elt = TensorProductElement(CG_hori, CG_vert) + else: + cell = self.mesh.ufl_cell().cellname() + V_elt = FiniteElement("DG", cell, horizontal_degree+2, variant=variant) + + # How should we name this if the horizontal and vertical degrees are different? + name = f'CG{horizontal_degree+2}' + + return FunctionSpace(self.mesh, V_elt, name=name) diff --git a/gusto/io.py b/gusto/io.py index f8fb7682e..c77cdb576 100644 --- a/gusto/io.py +++ b/gusto/io.py @@ -160,13 +160,16 @@ def dump(self, equation, t): class IO(object): """Controls the model's input, output and diagnostics.""" - def __init__(self, mesh, dt, + def __init__(self, domain, equation, dt, output=None, parameters=None, diagnostics=None, diagnostic_fields=None): """ Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + equation (:class:`PrognosticEquation`): the prognostic equation. dt (:class:`Constant`): the time taken to perform a single model step. If a float or int is passed, it will be cast to a :class:`Constant`. @@ -182,6 +185,9 @@ def __init__(self, mesh, dt, TypeError: if `dt` cannot be cast to a :class:`Constant`. """ + self.equation = equation + self.mesh = domain.mesh + if output is None: # TODO: maybe this shouldn't be an optional argument then? raise RuntimeError("You must provide a directory name for dumping results") @@ -200,14 +206,11 @@ def __init__(self, mesh, dt, # TODO: quick way of ensuring that diagnostics are registered if hasattr(self, "field_names"): - for fname in self.field_names: + for fname in equation.field_names: self.diagnostics.register(fname) self.bcs[fname] = [] else: - self.diagnostics.register(field_name) - - # The mesh - self.mesh = mesh + self.diagnostics.register(equation.field_name) if self.output.dumplist is None: self.output.dumplist = [] @@ -218,7 +221,7 @@ def __init__(self, mesh, dt, # setup logger logger.setLevel(output.log_level) - set_log_handler(mesh.comm) + set_log_handler(self.mesh.comm) if parameters is not None: logger.info("Physical parameters that take non-default values:") logger.info(", ".join("%s: %s" % (k, float(v)) for (k, v) in vars(parameters).items())) @@ -242,7 +245,7 @@ def setup_diagnostics(self): f = SteadyStateError(self, name) self.diagnostic_fields.append(f) - fields = set([f.name() for f in self.fields]) + fields = set([f.name() for f in self.equation.fields]) field_deps = [(d, sorted(set(d.required_fields).difference(fields),)) for d in self.diagnostic_fields] schedule = topo_sort(field_deps) self.diagnostic_fields = schedule @@ -293,7 +296,7 @@ def setup_dump(self, t, tmax, pickup=False): comm=self.mesh.comm) # make list of fields to dump - self.to_dump = [f for f in self.fields if f.name() in self.fields.to_dump] + self.to_dump = [f for f in self.equation.fields if f.name() in self.equation.fields.to_dump] # make dump counter self.dumpcount = itertools.count() @@ -310,7 +313,7 @@ def setup_dump(self, t, tmax, pickup=False): # make functions on latlon mesh, as specified by dumplist_latlon self.to_dump_latlon = [] for name in self.output.dumplist_latlon: - f = self.fields(name) + f = self.equation.fields(name) field = Function( functionspaceimpl.WithGeometry.create( f.function_space(), mesh_ll), @@ -334,7 +337,7 @@ def setup_dump(self, t, tmax, pickup=False): self.pointdata_output = PointDataOutput(pointdata_filename, ndt, self.output.point_data, self.output.dirname, - self.fields, + self.equation.fields, self.mesh.comm, self.output.tolerance, create=not pickup) @@ -355,7 +358,7 @@ def setup_dump(self, t, tmax, pickup=False): mode=FILE_CREATE) # make list of fields to pickup (this doesn't include # diagnostic fields) - self.to_pickup = [f for f in self.fields if f.name() in self.fields.to_pickup] + self.to_pickup = [f for f in self.equation.fields if f.name() in self.equation.fields.to_pickup] # if we want to checkpoint then make a checkpoint counter if self.output.checkpoint: @@ -369,7 +372,7 @@ def pickup_from_checkpoint(self): # TODO: this duplicates some code from setup_dump. Can this be avoided? # It is because we don't know if we are picking up or setting dump first if self.to_pickup is None: - self.to_pickup = [f for f in self.fields if f.name() in self.fields.to_pickup] + self.to_pickup = [f for f in self.equation.fields if f.name() in self.equation.fields.to_pickup] # Set dumpdir if has not been done already if self.dumpdir is None: self.dumpdir = path.join("results", self.output.dirname) @@ -412,11 +415,11 @@ def dump(self, t): if output.dump_diagnostics: # Output diagnostic data - self.diagnostic_output.dump(self, t) + self.diagnostic_output.dump(self.equation, t) if len(output.point_data) > 0 and (next(self.pddumpcount) % output.pddumpfreq) == 0: # Output pointwise data - self.pointdata_output.dump(self.fields, t) + self.pointdata_output.dump(self.equation.fields, t) # Dump all the fields to the checkpointing file (backup version) if output.checkpoint and (next(self.chkptcount) % output.chkptfreq) == 0: @@ -443,7 +446,7 @@ def initialise(self, initial_conditions): is used to set the initial field. """ for name, ic in initial_conditions: - f_init = getattr(self.fields, name) + f_init = getattr(self.equation.fields, name) f_init.assign(ic) f_init.rename(name) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 54a85742c..1e23e914e 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -73,11 +73,13 @@ def new_apply(self, x_out, x_in): class TimeDiscretisation(object, metaclass=ABCMeta): """Base class for time discretisation schemes.""" - def __init__(self, state, field_name=None, solver_parameters=None, + def __init__(self, domain, io, field_name=None, solver_parameters=None, limiter=None, options=None): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + io (:class:`IO`): the model's object for controlling input/output. field_name (str, optional): name of the field to be evolved. Defaults to None. solver_parameters (dict, optional): dictionary of parameters to @@ -89,10 +91,11 @@ def __init__(self, state, field_name=None, solver_parameters=None, to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. """ - self.state = state + self.domain = domain self.field_name = field_name + self.equation = None - self.dt = self.state.dt + self.dt = io.dt self.limiter = limiter @@ -125,11 +128,12 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): *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: self.idx = equation.field_names.index(self.field_name) - self.fs = self.state.fields(self.field_name).function_space() + self.fs = equation.fields(self.field_name).function_space() self.residual = self.residual.label_map( lambda t: t.get(prognostic) == self.field_name, lambda t: Term( @@ -172,7 +176,7 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): # construct the embedding space if not specified if options.embedding_space is None: V_elt = BrokenElement(self.fs.ufl_element()) - self.fs = FunctionSpace(self.state.mesh, V_elt) + self.fs = FunctionSpace(self.domain.mesh, V_elt) else: self.fs = options.embedding_space self.xdg_in = Function(self.fs) @@ -180,7 +184,7 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): if self.idx is None: self.x_projected = Function(equation.function_space) else: - self.x_projected = Function(self.state.fields(self.field_name).function_space()) + self.x_projected = Function(equation.fields(self.field_name).function_space()) new_test = TestFunction(self.fs) parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', @@ -204,7 +208,7 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): if self.discretisation_option == "supg": # construct tau, if it is not specified - dim = self.state.mesh.topological_dimension() + dim = self.domain.mesh.topological_dimension() if options.tau is not None: # if tau is provided, check that is has the right size tau = options.tau @@ -253,7 +257,7 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): if self.discretisation_option == "recovered": # set up the necessary functions - self.x_in = Function(self.state.fields(self.field_name).function_space()) + self.x_in = Function(equation.fields(self.field_name).function_space()) # Operator to recover to higher discontinuous space self.x_recoverer = ReversibleRecoverer(self.x_in, self.xdg_in, options) @@ -371,14 +375,14 @@ def replace_transport_term(self): # Do the options specify a different ibp to the old transport term? if old_transport_term.labels['ibp'] != self.options.ibp: # Set up a new transport term - field = self.state.fields(self.field_name) + field = self.equation.fields(self.field_name) test = TestFunction(self.fs) # Set up new transport term (depending on the type of transport equation) if old_transport_term.labels['transport'] == TransportEquationType.advective: - new_transport_term = advection_form(self.state, test, field, ibp=self.options.ibp) + new_transport_term = advection_form(self.domain, test, field, ibp=self.options.ibp) elif old_transport_term.labels['transport'] == TransportEquationType.conservative: - new_transport_term = continuity_form(self.state, test, field, ibp=self.options.ibp) + new_transport_term = continuity_form(self.domain, test, field, ibp=self.options.ibp) else: raise NotImplementedError(f'Replacement of transport term not implemented yet for {old_transport_term.labels["transport"]}') @@ -434,11 +438,13 @@ def apply(self, x_out, x_in): class ExplicitTimeDiscretisation(TimeDiscretisation): """Base class for explicit time discretisations.""" - def __init__(self, state, field_name=None, subcycles=None, + def __init__(self, domain, io, field_name=None, subcycles=None, solver_parameters=None, limiter=None, options=None): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + io (:class:`IO`): the model's object for controlling input/output. field_name (str, optional): name of the field to be evolved. Defaults to None. subcycles (int, optional): the number of sub-steps to perform. @@ -452,7 +458,7 @@ def __init__(self, state, field_name=None, subcycles=None, to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. """ - super().__init__(state, field_name, + super().__init__(domain, io, field_name, solver_parameters=solver_parameters, limiter=limiter, options=options) @@ -774,11 +780,13 @@ class BackwardEuler(TimeDiscretisation): The backward Euler method for operator F is the most simple implicit scheme: y^(n+1) = y^n + dt*F[y^(n+1)]. """ - def __init__(self, state, field_name=None, solver_parameters=None, + def __init__(self, domain, io, field_name=None, solver_parameters=None, limiter=None, options=None): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + io (:class:`IO`): the model's object for controlling input/output. field_name (str, optional): name of the field to be evolved. Defaults to None. subcycles (int, optional): the number of sub-steps to perform. @@ -797,7 +805,7 @@ def __init__(self, state, field_name=None, solver_parameters=None, """ if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") - super().__init__(state=state, field_name=field_name, + super().__init__(domain=domain, io=io, field_name=field_name, solver_parameters=solver_parameters, limiter=limiter, options=options) @@ -844,11 +852,13 @@ class ThetaMethod(TimeDiscretisation): for off-centring parameter theta. """ - def __init__(self, state, field_name=None, theta=None, + def __init__(self, domain, io, field_name=None, theta=None, solver_parameters=None, options=None): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + io (:class:`IO`): the model's object for controlling input/output. field_name (str, optional): name of the field to be evolved. Defaults to None. theta (float, optional): the off-centring parameter. theta = 1 @@ -876,7 +886,7 @@ def __init__(self, state, field_name=None, theta=None, 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} - super().__init__(state, field_name, + super().__init__(domain, io, field_name, solver_parameters=solver_parameters, options=options) @@ -926,11 +936,13 @@ class ImplicitMidpoint(ThetaMethod): It is equivalent to the "theta" method with theta = 1/2. """ - def __init__(self, state, field_name=None, solver_parameters=None, + def __init__(self, domain, io, field_name=None, solver_parameters=None, options=None): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + io (:class:`IO`): the model's object for controlling input/output. field_name (str, optional): name of the field to be evolved. Defaults to None. solver_parameters (dict, optional): dictionary of parameters to @@ -940,7 +952,7 @@ def __init__(self, state, field_name=None, solver_parameters=None, to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. """ - super().__init__(state, field_name, theta=0.5, + super().__init__(domain, io, field_name, theta=0.5, solver_parameters=solver_parameters, options=options) @@ -948,11 +960,13 @@ def __init__(self, state, field_name=None, solver_parameters=None, class MultilevelTimeDiscretisation(TimeDiscretisation): """Base class for multi-level timesteppers""" - def __init__(self, state, field_name=None, solver_parameters=None, + def __init__(self, domain, io, field_name=None, solver_parameters=None, limiter=None, options=None): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + io (:class:`IO`): the model's object for controlling input/output. field_name (str, optional): name of the field to be evolved. Defaults to None. solver_parameters (dict, optional): dictionary of parameters to @@ -966,7 +980,7 @@ def __init__(self, state, field_name=None, solver_parameters=None, """ if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") - super().__init__(state=state, field_name=field_name, + super().__init__(domain=domain, io=io, field_name=field_name, solver_parameters=solver_parameters, limiter=limiter, options=options) self.initial_timesteps = 0 diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 988677f70..3fce49b57 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -19,15 +19,15 @@ class BaseTimestepper(object, metaclass=ABCMeta): """Base class for timesteppers.""" - def __init__(self, equation, state): + def __init__(self, equation, io): """ Args: equation (:class:`PrognosticEquation`): the prognostic equation. - state (:class:`State`): the model's state object + io (:class:`IO`): the model's object for controlling input/output. """ self.equation = equation - self.state = state + self.io = io self.setup_fields() self.setup_scheme() @@ -61,39 +61,39 @@ def run(self, t, tmax, pickup=False): pickup: (bool): specify whether to pickup from a previous run """ - state = self.state + io = self.io if pickup: - t = state.pickup_from_checkpoint() + t = io.pickup_from_checkpoint() - state.setup_diagnostics() + io.setup_diagnostics() with timed_stage("Dump output"): - state.setup_dump(t, tmax, pickup) + io.setup_dump(t, tmax, pickup) - state.t.assign(t) + io.t.assign(t) - self.x.initialise(state) + self.x.initialise(self.equation) - while float(state.t) < tmax - 0.5*float(state.dt): - logger.info(f'at start of timestep, t={float(state.t)}, dt={float(state.dt)}') + while float(io.t) < tmax - 0.5*float(io.dt): + logger.info(f'at start of timestep, t={float(io.t)}, dt={float(io.dt)}') self.x.update() self.timestep() for field in self.x.np1: - state.fields(field.name()).assign(field) + self.equation.fields(field.name()).assign(field) - state.t.assign(state.t + state.dt) + io.t.assign(io.t + io.dt) with timed_stage("Dump output"): - state.dump(float(state.t)) + io.dump(float(io.t)) - if state.output.checkpoint: - state.chkpt.close() + if io.output.checkpoint: + io.chkpt.close() - logger.info(f'TIMELOOP complete. t={float(state.t)}, tmax={tmax}') + logger.info(f'TIMELOOP complete. t={float(io.t)}, tmax={tmax}') class Timestepper(BaseTimestepper): @@ -101,16 +101,16 @@ class Timestepper(BaseTimestepper): Implements a timeloop by applying a scheme to a prognostic equation. """ - def __init__(self, equation, scheme, state): + def __init__(self, equation, scheme, io): """ Args: equation (:class:`PrognosticEquation`): the prognostic equation scheme (:class:`TimeDiscretisation`): the scheme to use to timestep the prognostic equation - state (:class:`State`): the model's state object + io (:class:`IO`): the model's object for controlling input/output. """ self.scheme = scheme - super().__init__(equation=equation, state=state) + super().__init__(equation=equation, io=io) @property def transporting_velocity(self): @@ -140,13 +140,13 @@ class SplitPhysicsTimestepper(Timestepper): scheme to be applied to the physics terms than the prognostic equation. """ - def __init__(self, equation, scheme, state, physics_schemes=None): + def __init__(self, equation, scheme, io, physics_schemes=None): """ Args: equation (:class:`PrognosticEquation`): the prognostic equation scheme (:class:`TimeDiscretisation`): the scheme to use to timestep the prognostic equation - state (:class:`State`): the model's state object + 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. @@ -156,7 +156,7 @@ def __init__(self, equation, scheme, state, physics_schemes=None): self.equation = equation self.scheme = scheme - self.state = state + self.io = io self.setup_fields() self.setup_scheme() @@ -206,7 +206,7 @@ class SemiImplicitQuasiNewton(BaseTimestepper): terms. """ - def __init__(self, equation_set, state, transport_schemes, + def __init__(self, equation_set, io, transport_schemes, auxiliary_equations_and_schemes=None, linear_solver=None, diffusion_schemes=None, @@ -216,7 +216,7 @@ def __init__(self, equation_set, state, transport_schemes, Args: equation_set (:class:`PrognosticEquationSet`): the prognostic equation set to be solved - state (:class:`State`) the model's state object + 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 @@ -265,7 +265,10 @@ def __init__(self, equation_set, state, transport_schemes, assert scheme.field_name in equation_set.field_names self.diffusion_schemes.append((scheme.field_name, scheme)) - super().__init__(equation_set, state) + if not equation_set.reference_profile_initialised: + raise RuntimeError('Reference profiles for equation set must be initialised to use Semi-Implicit Timestepper') + + super().__init__(equation_set, io) if auxiliary_equations_and_schemes is not None: for eqn, scheme in auxiliary_equations_and_schemes: @@ -379,7 +382,7 @@ def timestep(self): xrhs -= xnp1(self.field_name) with timed_stage("Implicit solve"): - self.linear_solver.solve(xrhs, dy) # solves linear system and places result in state.dy + self.linear_solver.solve(xrhs, dy) # solves linear system and places result in dy xnp1X = xnp1(self.field_name) xnp1X += dy @@ -406,14 +409,14 @@ class PrescribedTransport(Timestepper): """ Implements a timeloop with a prescibed transporting velocity """ - def __init__(self, equation, scheme, state, physics_schemes=None, + def __init__(self, equation, scheme, io, 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 - state (:class:`State`): the model's state object + 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. @@ -427,7 +430,7 @@ def __init__(self, equation, scheme, state, physics_schemes=None, updated. Defaults to None. """ - super().__init__(equation, scheme, state) + super().__init__(equation, scheme, io) if physics_schemes is not None: self.physics_schemes = physics_schemes @@ -442,14 +445,14 @@ def __init__(self, equation, scheme, state, physics_schemes=None, if prescribed_transporting_velocity is not None: self.velocity_projection = Projector( - prescribed_transporting_velocity(self.state.t), - self.state.fields('u')) + prescribed_transporting_velocity(self.io.t), + self.equation.fields('u')) else: self.velocity_projection = None @property def transporting_velocity(self): - return self.state.fields('u') + return self.equation.fields('u') def setup_fields(self): self.x = TimeLevelFields(self.equation, self.scheme.nlevels) diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py index 883c10909..64ebc13d2 100644 --- a/integration-tests/conftest.py +++ b/integration-tests/conftest.py @@ -9,7 +9,7 @@ from collections import namedtuple import pytest -opts = ('state', 'tmax', 'f_init', 'f_end', 'family', 'degree', +opts = ('domain', 'dt', 'tmax', 'output', 'f_init', 'f_end', 'degree', 'uexpr', 'umax', 'radius', 'tol') TracerSetup = namedtuple('TracerSetup', opts) TracerSetup.__new__.__defaults__ = (None,)*len(opts) @@ -29,7 +29,7 @@ def tracer_sphere(tmpdir, degree): dt = pi/3. * 0.02 output = OutputParameters(dirname=str(tmpdir), dumpfreq=15) - state = State(mesh, dt=dt, output=output) + domain = Domain(mesh, family="BDM", degree=degree) umax = 1.0 uexpr = as_vector([- umax * x[1] / radius, umax * x[0] / radius, 0.0]) @@ -40,7 +40,7 @@ def tracer_sphere(tmpdir, degree): tol = 0.05 - return TracerSetup(state, tmax, f_init, f_end, "BDM", degree, + return TracerSetup(domain, dt, tmax, output, f_init, f_end, degree, uexpr, umax, radius, tol) @@ -56,7 +56,7 @@ def tracer_slice(tmpdir, degree): dt = 0.01 tmax = 0.75 output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) - state = State(mesh, dt=dt, output=output) + domain = Domain(mesh, family="CG", degree=degree) uexpr = as_vector([2.0, 0.0]) @@ -73,7 +73,7 @@ def tracer_slice(tmpdir, degree): tol = 0.12 - return TracerSetup(state, tmax, f_init, f_end, "CG", degree, uexpr, tol=tol) + return TracerSetup(domain, dt, tmax, output, f_init, f_end, degree, uexpr, tol=tol) def tracer_blob_slice(tmpdir, degree): @@ -83,13 +83,13 @@ def tracer_blob_slice(tmpdir, degree): mesh = ExtrudedMesh(m, layers=10, layer_height=1.) output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) - state = State(mesh, dt=dt, output=output) + domain = Domain(mesh, family="CG", degree=degree) tmax = 1. x = SpatialCoordinate(mesh) f_init = exp(-((x[0]-0.5*L)**2 + (x[1]-0.5*L)**2)) - return TracerSetup(state, tmax, f_init, family="CG", degree=degree) + return TracerSetup(domain, dt, tmax, output, f_init, degree=degree) @pytest.fixture() diff --git a/integration-tests/transport/test_dg_transport.py b/integration-tests/transport/test_dg_transport.py index 46e48372e..981dce817 100644 --- a/integration-tests/transport/test_dg_transport.py +++ b/integration-tests/transport/test_dg_transport.py @@ -8,29 +8,29 @@ import pytest -def run(eqn, transport_scheme, state, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, state) +def run(eqn, transport_scheme, io, tmax, f_end): + timestepper = PrescribedTransport(eqn, transport_scheme, io) timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(eqn.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("geometry", ["slice", "sphere"]) @pytest.mark.parametrize("equation_form", ["advective", "continuity"]) def test_dg_transport_scalar(tmpdir, geometry, equation_form, tracer_setup): setup = tracer_setup(tmpdir, geometry) - state = setup.state - V = state.spaces("DG", "DG", 1) + domain = setup.domain + V = domain.spaces("DG") if equation_form == "advective": - eqn = AdvectionEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) + eqn = AdvectionEquation(domain, V, "f") else: - eqn = ContinuityEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) - state.fields("f").interpolate(setup.f_init) - state.fields("u").project(setup.uexpr) + eqn = ContinuityEquation(domain, V, "f") - transport_scheme = SSPRK3(state) - error = run(eqn, transport_scheme, state, setup.tmax, setup.f_end) + io = IO(domain, eqn, dt=setup.dt, output=setup.output) + eqn.fields("f").interpolate(setup.f_init) + eqn.fields("u").project(setup.uexpr) + + transport_scheme = SSPRK3(domain, io) + error = run(eqn, transport_scheme, io, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' @@ -39,20 +39,20 @@ def test_dg_transport_scalar(tmpdir, geometry, equation_form, tracer_setup): @pytest.mark.parametrize("equation_form", ["advective", "continuity"]) def test_dg_transport_vector(tmpdir, geometry, equation_form, tracer_setup): setup = tracer_setup(tmpdir, geometry) - state = setup.state - gdim = state.mesh.geometric_dimension() + domain = setup.domain + gdim = domain.mesh.geometric_dimension() f_init = as_vector([setup.f_init]*gdim) - V = VectorFunctionSpace(state.mesh, "DG", 1) + V = VectorFunctionSpace(domain.mesh, "DG", 1) if equation_form == "advective": - eqn = AdvectionEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) + eqn = AdvectionEquation(domain, V, "f") else: - eqn = ContinuityEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) - state.fields("f").interpolate(f_init) - state.fields("u").project(setup.uexpr) - transport_schemes = SSPRK3(state) + eqn = ContinuityEquation(domain, V, "f") + + io = IO(domain, eqn, dt=setup.dt, output=setup.output) + eqn.fields("f").interpolate(f_init) + eqn.fields("u").project(setup.uexpr) + transport_schemes = SSPRK3(domain, io) f_end = as_vector([setup.f_end]*gdim) - error = run(eqn, transport_schemes, state, setup.tmax, f_end) + error = run(eqn, transport_schemes, io, setup.tmax, f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' From cff88671263fac00004b05a1e0c6eb3fc4037afe Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Thu, 8 Dec 2022 21:29:41 +0000 Subject: [PATCH 04/12] update all transport integration tests --- gusto/function_spaces.py | 17 ++++-- .../transport/test_dg_transport.py | 1 + .../transport/test_embedded_dg_advection.py | 28 ++++----- integration-tests/transport/test_limiters.py | 55 +++++++++-------- .../transport/test_recovered_transport.py | 27 ++++---- .../transport/test_subcycling.py | 27 ++++---- .../transport/test_supg_transport.py | 61 ++++++++++--------- .../transport/test_vector_recovered_space.py | 30 ++++----- 8 files changed, 131 insertions(+), 115 deletions(-) diff --git a/gusto/function_spaces.py b/gusto/function_spaces.py index 71f12ee4a..3b68e485d 100644 --- a/gusto/function_spaces.py +++ b/gusto/function_spaces.py @@ -61,6 +61,13 @@ def __call__(self, name, family=None, horizontal_degree=None, # The space itself has been provided (to add it to the creator) value = V + elif name == "DG1_equispaced": + # Special case based on name + if self.extruded_mesh: + value = self.build_dg_space(1, 1, variant='equispaced') + else: + value = self.build_dg_space(1, variant='equispaced') + else: # Need to create space, based on name/family/degree assert horizontal_degree is not None @@ -70,8 +77,6 @@ def __call__(self, name, family=None, horizontal_degree=None, value = self.build_hdiv_space(family, horizontal_degree, vertical_degree) elif name == "theta": value = self.build_theta_space(horizontal_degree, vertical_degree) - elif name == "DG1_equispaced": - value = self.build_dg_space(1, variant='equispaced') elif family == "DG": value = self.build_dg_space(horizontal_degree, vertical_degree) elif family == "CG": @@ -252,14 +257,14 @@ def build_cg_space(self, horizontal_degree, vertical_degree): if vertical_degree is None: raise ValueError('vertical_degree must be specified to create CG space on an extruded mesh') cell = self.mesh._base_mesh.ufl_cell().cellname() - CG_hori = FiniteElement("CG", cell, horizontal_degree+2) - CG_vert = FiniteElement("CG", interval, vertical_degree+2) + CG_hori = FiniteElement("CG", cell, horizontal_degree+1) + CG_vert = FiniteElement("CG", interval, vertical_degree+1) V_elt = TensorProductElement(CG_hori, CG_vert) else: cell = self.mesh.ufl_cell().cellname() - V_elt = FiniteElement("DG", cell, horizontal_degree+2, variant=variant) + V_elt = FiniteElement("DG", cell, horizontal_degree+1, variant=variant) # How should we name this if the horizontal and vertical degrees are different? - name = f'CG{horizontal_degree+2}' + name = f'CG{horizontal_degree+1}' return FunctionSpace(self.mesh, V_elt, name=name) diff --git a/integration-tests/transport/test_dg_transport.py b/integration-tests/transport/test_dg_transport.py index 981dce817..457296896 100644 --- a/integration-tests/transport/test_dg_transport.py +++ b/integration-tests/transport/test_dg_transport.py @@ -20,6 +20,7 @@ def test_dg_transport_scalar(tmpdir, geometry, equation_form, tracer_setup): setup = tracer_setup(tmpdir, geometry) domain = setup.domain V = domain.spaces("DG") + if equation_form == "advective": eqn = AdvectionEquation(domain, V, "f") else: diff --git a/integration-tests/transport/test_embedded_dg_advection.py b/integration-tests/transport/test_embedded_dg_advection.py index 1df300601..376da9ba4 100644 --- a/integration-tests/transport/test_embedded_dg_advection.py +++ b/integration-tests/transport/test_embedded_dg_advection.py @@ -8,10 +8,10 @@ import pytest -def run(eqn, transport_schemes, state, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_schemes, state) +def run(eqn, transport_schemes, io, tmax, f_end): + timestepper = PrescribedTransport(eqn, transport_schemes, io) timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(eqn.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("ibp", [IntegrateByParts.ONCE, IntegrateByParts.TWICE]) @@ -20,25 +20,25 @@ def run(eqn, transport_schemes, state, tmax, f_end): def test_embedded_dg_advection_scalar(tmpdir, ibp, equation_form, space, tracer_setup): setup = tracer_setup(tmpdir, geometry="slice") - state = setup.state - V = state.spaces("theta", degree=1) + domain = setup.domain + V = domain.spaces("theta") if space == "broken": opts = EmbeddedDGOptions() elif space == "dg": - opts = EmbeddedDGOptions(embedding_space=state.spaces("DG1", "DG", 1)) + opts = EmbeddedDGOptions(embedding_space=domain.spaces("DG")) if equation_form == "advective": - eqn = AdvectionEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree, ibp=ibp) + eqn = AdvectionEquation(domain, V, "f", ibp=ibp) else: - eqn = ContinuityEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree, ibp=ibp) - state.fields("f").interpolate(setup.f_init) - state.fields("u").project(setup.uexpr) + eqn = ContinuityEquation(domain, V, "f", ibp=ibp) - transport_schemes = SSPRK3(state, options=opts) + io = IO(domain, eqn, dt=setup.dt, output=setup.output) + eqn.fields("f").interpolate(setup.f_init) + eqn.fields("u").project(setup.uexpr) - error = run(eqn, transport_schemes, state, setup.tmax, setup.f_end) + transport_schemes = SSPRK3(domain, io, options=opts) + + error = run(eqn, transport_schemes, io, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_limiters.py b/integration-tests/transport/test_limiters.py index 5672c8f1c..1ca6dd1cb 100644 --- a/integration-tests/transport/test_limiters.py +++ b/integration-tests/transport/test_limiters.py @@ -1,5 +1,5 @@ """ -This tests three limiter options for different transport schemes. +This tests limiter options for different transport schemes. A sharp bubble of warm air is generated in a vertical slice and then transported by a prescribed transport scheme. If the limiter is working, the transport should have produced no new maxima or minima. @@ -33,39 +33,41 @@ def setup_limiters(dirname, space): mesh = ExtrudedMesh(m, layers=20, layer_height=(Ld/20)) output = OutputParameters(dirname=dirname+'/limiters', dumpfreq=1, dumplist=['u', 'tracer', 'true_tracer']) - parameters = CompressibleParameters() - state = State(mesh, - dt=dt, - output=output, - parameters=parameters) + degree = 0 if space in ['DG0', 'Vtheta_degree_0'] else 1 + + domain = Domain(mesh, family="CG", degree=degree) if space == 'DG0': - V = state.spaces('DG', 'DG', 0) + V = domain.spaces('DG') VCG1 = FunctionSpace(mesh, 'CG', 1) - VDG1 = state.spaces('DG1_equispaced') + VDG1 = domain.spaces('DG1_equispaced') + elif space == 'DG1': + V = domain.spaces('DG') elif space == 'DG1_equispaced': - V = state.spaces('DG1_equispaced') + V = domain.spaces('DG1_equispaced') elif space == 'Vtheta_degree_0': - V = state.spaces('theta', degree=0) + V = domain.spaces('theta') VCG1 = FunctionSpace(mesh, 'CG', 1) - VDG1 = state.spaces('DG1_equispaced') + VDG1 = domain.spaces('DG1_equispaced') elif space == 'Vtheta_degree_1': - V = state.spaces('theta', degree=1) + V = domain.spaces('theta') else: raise NotImplementedError - Vpsi = FunctionSpace(mesh, 'CG', 2) + Vpsi = domain.spaces('CG', 'CG', degree, degree) # set up the equation - eqn = AdvectionEquation(state, V, 'tracer', ufamily='CG', udegree=1) + eqn = AdvectionEquation(domain, V, 'tracer') + + io = IO(domain, eqn, dt=dt, output=output) # ------------------------------------------------------------------------ # # Initial condition # ------------------------------------------------------------------------ # - tracer0 = state.fields('tracer', V) - true_field = state.fields('true_tracer', V) + tracer0 = eqn.fields('tracer', V) + true_field = eqn.fields('true_tracer', V) x, z = SpatialCoordinate(mesh) @@ -125,7 +127,7 @@ def setup_limiters(dirname, space): # ------------------------------------------------------------------------ # psi = Function(Vpsi) - u = state.fields('u') + u = eqn.fields('u') # set up solid body rotation for transport # we do this slightly complicated stream function to make the velocity 0 at edges @@ -158,33 +160,36 @@ def setup_limiters(dirname, space): recovered_space=VCG1, project_low_method='recover', boundary_method=BoundaryMethod.taylor) - transport_schemes = SSPRK3(state, options=opts, + transport_schemes = SSPRK3(domain, io, options=opts, limiter=VertexBasedLimiter(VDG1)) + elif space == 'DG1': + transport_schemes = SSPRK3(domain, io, limiter=DG1Limiter(V)) + elif space == 'DG1_equispaced': - transport_schemes = SSPRK3(state, limiter=VertexBasedLimiter(V)) + transport_schemes = SSPRK3(domain, io, limiter=VertexBasedLimiter(V)) elif space == 'Vtheta_degree_1': opts = EmbeddedDGOptions() - transport_schemes = SSPRK3(state, options=opts, limiter=ThetaLimiter(V)) + transport_schemes = SSPRK3(domain, io, options=opts, limiter=ThetaLimiter(V)) else: raise NotImplementedError # build time stepper - stepper = PrescribedTransport(eqn, transport_schemes, state) + stepper = PrescribedTransport(eqn, transport_schemes, io) - return stepper, tmax, state, true_field + return stepper, tmax, eqn, true_field @pytest.mark.parametrize('space', ['Vtheta_degree_0', 'Vtheta_degree_1', - 'DG0', 'DG1_equispaced']) + 'DG0', 'DG1', 'DG1_equispaced']) def test_limiters(tmpdir, space): # Setup and run dirname = str(tmpdir) - stepper, tmax, state, true_field = setup_limiters(dirname, space) + stepper, tmax, eqn, true_field = setup_limiters(dirname, space) stepper.run(t=0, tmax=tmax) - final_field = state.fields('tracer') + final_field = eqn.fields('tracer') # Check tracer is roughly in the correct place assert norm(true_field - final_field) / norm(true_field) < 0.05, \ diff --git a/integration-tests/transport/test_recovered_transport.py b/integration-tests/transport/test_recovered_transport.py index 35b3881b2..057e93a06 100644 --- a/integration-tests/transport/test_recovered_transport.py +++ b/integration-tests/transport/test_recovered_transport.py @@ -9,41 +9,42 @@ import pytest -def run(eqn, transport_scheme, state, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, state) +def run(eqn, transport_scheme, io, tmax, f_end): + timestepper = PrescribedTransport(eqn, transport_scheme, io) timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(eqn.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("geometry", ["slice", "sphere"]) def test_recovered_space_setup(tmpdir, geometry, tracer_setup): - # Make mesh and state using routine from conftest + # Make domain using routine from conftest setup = tracer_setup(tmpdir, geometry, degree=0) - state = setup.state - mesh = state.mesh + domain = setup.domain + mesh = domain.mesh # Spaces for recovery VDG0 = FunctionSpace(mesh, "DG", 0) - VDG1 = state.spaces("DG1_equispaced") + VDG1 = domain.spaces("DG1_equispaced") VCG1 = FunctionSpace(mesh, "CG", 1) # Make equation - eqn = ContinuityEquation(state, VDG0, "f", - ufamily=setup.family, udegree=1) + eqn = ContinuityEquation(domain, VDG0, "f") + + io = IO(domain, eqn, dt=setup.dt, output=setup.output) # Initialise fields - state.fields("f").interpolate(setup.f_init) - state.fields("u").project(setup.uexpr) + eqn.fields("f").interpolate(setup.f_init) + eqn.fields("u").project(setup.uexpr) # Declare transport scheme recovery_opts = RecoveryOptions(embedding_space=VDG1, recovered_space=VCG1, boundary_method=BoundaryMethod.taylor) - transport_scheme = SSPRK3(state, options=recovery_opts) + transport_scheme = SSPRK3(domain, io, options=recovery_opts) # Run and check error - error = run(eqn, transport_scheme, state, setup.tmax, setup.f_end) + error = run(eqn, transport_scheme, io, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_subcycling.py b/integration-tests/transport/test_subcycling.py index 7a019c586..963d3ac9a 100644 --- a/integration-tests/transport/test_subcycling.py +++ b/integration-tests/transport/test_subcycling.py @@ -8,28 +8,29 @@ import pytest -def run(eqn, transport_scheme, state, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, state) +def run(eqn, transport_scheme, io, tmax, f_end): + timestepper = PrescribedTransport(eqn, transport_scheme, io) timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(eqn.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("equation_form", ["advective", "continuity"]) def test_subcyling(tmpdir, equation_form, tracer_setup): geometry = "slice" setup = tracer_setup(tmpdir, geometry) - state = setup.state - V = state.spaces("DG", "DG", 1) + domain = setup.domain + V = domain.spaces("DG") if equation_form == "advective": - eqn = AdvectionEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) + eqn = AdvectionEquation(domain, V, "f") else: - eqn = ContinuityEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) - state.fields("f").interpolate(setup.f_init) - state.fields("u").project(setup.uexpr) + eqn = ContinuityEquation(domain, V, "f") - transport_scheme = SSPRK3(state, subcycles=2) - error = run(eqn, transport_scheme, state, setup.tmax, setup.f_end) + io = IO(domain, eqn, dt=setup.dt, output=setup.output) + + eqn.fields("f").interpolate(setup.f_init) + eqn.fields("u").project(setup.uexpr) + + transport_scheme = SSPRK3(domain, io, subcycles=2) + error = run(eqn, transport_scheme, io, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_supg_transport.py b/integration-tests/transport/test_supg_transport.py index dab44a551..8608e4c8f 100644 --- a/integration-tests/transport/test_supg_transport.py +++ b/integration-tests/transport/test_supg_transport.py @@ -3,15 +3,15 @@ to the correct position. """ -from firedrake import norm, VectorFunctionSpace, as_vector +from firedrake import norm, FunctionSpace, VectorFunctionSpace, as_vector from gusto import * import pytest -def run(eqn, transport_scheme, state, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, state) +def run(eqn, transport_scheme, io, tmax, f_end): + timestepper = PrescribedTransport(eqn, transport_scheme, io) timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(eqn.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("equation_form", ["advective", "continuity"]) @@ -21,31 +21,33 @@ def test_supg_transport_scalar(tmpdir, equation_form, scheme, space, tracer_setup): setup = tracer_setup(tmpdir, geometry="slice") - state = setup.state + domain = setup.domain if space == "CG": - V = state.spaces("CG1", "CG", 1) + V = FunctionSpace(domain.mesh, "CG", 1) ibp = IntegrateByParts.NEVER else: - V = state.spaces("theta", degree=1) + V = domain.spaces("theta") ibp = IntegrateByParts.TWICE opts = SUPGOptions(ibp=ibp) if equation_form == "advective": - eqn = AdvectionEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) + eqn = AdvectionEquation(domain, V, "f") else: - eqn = ContinuityEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) - state.fields("f").interpolate(setup.f_init) - state.fields("u").project(setup.uexpr) + eqn = ContinuityEquation(domain, V, "f") + + io = IO(domain, eqn, dt=setup.dt, output=setup.output) + + eqn.fields("f").interpolate(setup.f_init) + eqn.fields("u").project(setup.uexpr) + if scheme == "ssprk": - transport_scheme = SSPRK3(state, options=opts) + transport_scheme = SSPRK3(domain, io, options=opts) elif scheme == "implicit_midpoint": - transport_scheme = ImplicitMidpoint(state, options=opts) + transport_scheme = ImplicitMidpoint(domain, io, options=opts) - error = run(eqn, transport_scheme, state, setup.tmax, setup.f_end) + error = run(eqn, transport_scheme, io, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' @@ -57,37 +59,38 @@ def test_supg_transport_vector(tmpdir, equation_form, scheme, space, tracer_setup): setup = tracer_setup(tmpdir, geometry="slice") - state = setup.state + domain = setup.domain - gdim = state.mesh.geometric_dimension() + gdim = domain.mesh.geometric_dimension() f_init = as_vector([setup.f_init]*gdim) if space == "CG": - V = VectorFunctionSpace(state.mesh, "CG", 1) + V = VectorFunctionSpace(domain.mesh, "CG", 1) ibp = IntegrateByParts.NEVER else: - V = state.spaces("HDiv", setup.family, setup.degree) + V = domain.spaces("HDiv") ibp = IntegrateByParts.TWICE opts = SUPGOptions(ibp=ibp) if equation_form == "advective": - eqn = AdvectionEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) + eqn = AdvectionEquation(domain, V, "f") else: - eqn = ContinuityEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) - f = state.fields("f") + eqn = ContinuityEquation(domain, V, "f") + + io = IO(domain, eqn, dt=setup.dt, output=setup.output) + + f = eqn.fields("f") if space == "CG": f.interpolate(f_init) else: f.project(f_init) - state.fields("u").project(setup.uexpr) + eqn.fields("u").project(setup.uexpr) if scheme == "ssprk": - transport_scheme = SSPRK3(state, options=opts) + transport_scheme = SSPRK3(domain, io, options=opts) elif scheme == "implicit_midpoint": - transport_scheme = ImplicitMidpoint(state, options=opts) + transport_scheme = ImplicitMidpoint(domain, io, options=opts) f_end = as_vector([setup.f_end]*gdim) - error = run(eqn, transport_scheme, state, setup.tmax, f_end) + error = run(eqn, transport_scheme, io, setup.tmax, f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_vector_recovered_space.py b/integration-tests/transport/test_vector_recovered_space.py index 3fd91e5dd..ee0f19023 100644 --- a/integration-tests/transport/test_vector_recovered_space.py +++ b/integration-tests/transport/test_vector_recovered_space.py @@ -9,26 +9,26 @@ import pytest -def run(eqn, transport_scheme, state, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, state) +def run(eqn, transport_scheme, io, tmax, f_end): + timestepper = PrescribedTransport(eqn, transport_scheme, io) timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(eqn.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("geometry", ["slice"]) def test_vector_recovered_space_setup(tmpdir, geometry, tracer_setup): - # Make mesh and state using routine from conftest + # Make domain using routine from conftest setup = tracer_setup(tmpdir, geometry, degree=0) - state = setup.state - mesh = state.mesh - gdim = state.mesh.geometric_dimension() + domain = setup.domain + mesh = domain.mesh + gdim = mesh.geometric_dimension() # Spaces for recovery - Vu = state.spaces("HDiv", family=setup.family, degree=setup.degree) + Vu = domain.spaces("HDiv") if geometry == "slice": - VDG1 = state.spaces("DG1_equispaced") + VDG1 = domain.spaces("DG1_equispaced") Vec_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element(), name='Vec_DG1') Vec_CG1 = VectorFunctionSpace(mesh, "CG", 1, name='Vec_CG1') @@ -40,19 +40,19 @@ def test_vector_recovered_space_setup(tmpdir, geometry, tracer_setup): f'Recovered spaces for geometry {geometry} have not been implemented') # Make equation - eqn = AdvectionEquation(state, Vu, "f", - ufamily=setup.family, udegree=1) + eqn = AdvectionEquation(domain, Vu, "f") + io = IO(domain, eqn, dt=setup.dt, output=setup.output) # Initialise fields f_init = as_vector([setup.f_init]*gdim) - state.fields("f").project(f_init) - state.fields("u").project(setup.uexpr) + eqn.fields("f").project(f_init) + eqn.fields("u").project(setup.uexpr) - transport_scheme = SSPRK3(state, options=rec_opts) + transport_scheme = SSPRK3(domain, io, options=rec_opts) f_end = as_vector([setup.f_end]*gdim) # Run and check error - error = run(eqn, transport_scheme, state, setup.tmax, f_end) + error = run(eqn, transport_scheme, io, setup.tmax, f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' From edf741cf19bcaabae348aa595765026bed4d994b Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 11 Dec 2022 21:43:57 +0000 Subject: [PATCH 05/12] roll out through diagnostics and most of remaining integration tests --- gusto/diagnostics.py | 479 +++++++++--------- gusto/domain.py | 14 +- gusto/forcing.py | 4 +- gusto/function_spaces.py | 69 ++- gusto/io.py | 8 +- gusto/linear_solvers.py | 18 +- gusto/physics.py | 9 +- gusto/timeloop.py | 6 +- .../balance/test_compressible_balance.py | 30 +- .../balance/test_saturated_balance.py | 57 +-- .../balance/test_unsaturated_balance.py | 51 +- integration-tests/diffusion/test_diffusion.py | 40 +- integration-tests/model/test_checkpointing.py | 59 +-- .../model/test_passive_tracer.py | 34 +- .../model/test_prescribed_transport.py | 26 +- .../model/test_time_discretisation.py | 30 +- .../physics/test_condensation.py | 53 +- .../physics/test_instant_rain.py | 41 +- .../physics/test_precipitation.py | 30 +- integration-tests/transport/test_limiters.py | 2 +- 20 files changed, 535 insertions(+), 525 deletions(-) diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index e6a8168b3..fc0d62e5b 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -146,102 +146,112 @@ def name(self): """The name of this diagnostic field""" pass - def setup(self, state, space=None): + def setup(self, eqn, space=None): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. 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 space is None: - space = state.spaces("DG0", "DG", 0) - self.field = state.fields(self.name, space, pickup=False) + space = eqn.domain.spaces("DG0", "DG", 0) + + # TODO: would it be better for these diagnostics to be stored in the IO? + self.field = eqn.fields(self.name, space, pickup=False) self._initialised = True @abstractmethod - def compute(self, state): + def compute(self, eqn): """ Compute the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ pass - def __call__(self, state): + def __call__(self, eqn): """ Compute the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ - return self.compute(state) + return self.compute(eqn) class CourantNumber(DiagnosticField): """Dimensionless Courant number diagnostic field.""" name = "CourantNumber" - def setup(self, state): + def __init__(self, dt): + """ + Args: + dt (float): time increment over a model time step. + """ + super().__init__() + self.dt = dt + + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ if not self._initialised: - super(CourantNumber, self).setup(state) + super(CourantNumber, self).setup(eqn) # set up area computation - V = state.spaces("DG0") + V = eqn.domain.spaces("DG0") test = TestFunction(V) self.area = Function(V) assemble(test*dx, tensor=self.area) - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - u = state.fields("u") - dt = Constant(state.dt) - return self.field.project(sqrt(dot(u, u))/sqrt(self.area)*dt) + # TODO: update dt here? + u = eqn.fields("u") + return self.field.project(sqrt(dot(u, u))/sqrt(self.area)*self.dt) class VelocityX(DiagnosticField): """The geocentric Cartesian X component of the velocity field.""" name = "VelocityX" - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ if not self._initialised: - space = state.spaces("CG1", "CG", 1) - super(VelocityX, self).setup(state, space=space) + space = eqn.domain.spaces("CG1", "CG", 1) + super(VelocityX, self).setup(eqn, space=space) - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - u = state.fields("u") + u = eqn.fields("u") uh = u[0] return self.field.interpolate(uh) @@ -250,28 +260,28 @@ class VelocityZ(DiagnosticField): """The geocentric Cartesian Z component of the velocity field.""" name = "VelocityZ" - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ if not self._initialised: - space = state.spaces("CG1", "CG", 1) - super(VelocityZ, self).setup(state, space=space) + space = eqn.domain.spaces("CG1", "CG", 1) + super(VelocityZ, self).setup(eqn, space=space) - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - u = state.fields("u") + u = eqn.fields("u") w = u[u.geometric_dimension() - 1] return self.field.interpolate(w) @@ -280,28 +290,28 @@ class VelocityY(DiagnosticField): """The geocentric Cartesian Y component of the velocity field.""" name = "VelocityY" - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ if not self._initialised: - space = state.spaces("CG1", "CG", 1) - super(VelocityY, self).setup(state, space=space) + space = eqn.domain.spaces("CG1", "CG", 1) + super(VelocityY, self).setup(eqn, space=space) - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - u = state.fields("u") + u = eqn.fields("u") v = u[1] return self.field.interpolate(v) @@ -321,27 +331,27 @@ def name(self): """Gives the name of this diagnostic field.""" return self.fname+"_gradient" - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ if not self._initialised: - mesh_dim = state.mesh.geometric_dimension() + mesh_dim = eqn.domain.mesh.geometric_dimension() try: - field_dim = state.fields(self.fname).ufl_shape[0] + field_dim = eqn.fields(self.fname).ufl_shape[0] except IndexError: field_dim = 1 shape = (mesh_dim, ) * field_dim - space = TensorFunctionSpace(state.mesh, "CG", 1, shape=shape) - super().setup(state, space=space) + space = TensorFunctionSpace(eqn.domain.mesh, "CG", 1, shape=shape) + super().setup(eqn, space=space) - f = state.fields(self.fname) + f = eqn.fields(self.fname) test = TestFunction(space) trial = TrialFunction(space) - n = FacetNormal(state.mesh) + n = FacetNormal(eqn.domain.mesh) a = inner(test, trial)*dx L = -inner(div(test), f)*dx if space.extruded: @@ -349,12 +359,12 @@ def setup(self, state): prob = LinearVariationalProblem(a, L, self.field) self.solver = LinearVariationalSolver(prob) - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. @@ -367,28 +377,28 @@ class Divergence(DiagnosticField): """Diagnostic for computing the divergence of vector-valued fields.""" name = "Divergence" - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ if not self._initialised: - space = state.spaces("DG1", "DG", 1) - super(Divergence, self).setup(state, space=space) + space = eqn.domain.spaces("DG1", "DG", 1) + super(Divergence, self).setup(eqn, space=space) - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - u = state.fields("u") + u = eqn.fields("u") return self.field.interpolate(div(u)) @@ -402,28 +412,28 @@ def __init__(self, name): super().__init__() self.fname = name - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ if not self._initialised: # check geometric dimension is 3D - if state.mesh.geometric_dimension() != 3: + if eqn.domain.mesh.geometric_dimension() != 3: raise ValueError('Spherical components only work when the geometric dimension is 3!') - space = FunctionSpace(state.mesh, "CG", 1) - super().setup(state, space=space) + space = FunctionSpace(eqn.domain.mesh, "CG", 1) + super().setup(eqn, space=space) - V = VectorFunctionSpace(state.mesh, "CG", 1) - self.x, self.y, self.z = SpatialCoordinate(state.mesh) + V = VectorFunctionSpace(eqn.domain.mesh, "CG", 1) + self.x, self.y, self.z = SpatialCoordinate(eqn.domain.mesh) self.x_hat = Function(V).interpolate(Constant(as_vector([1.0, 0.0, 0.0]))) self.y_hat = Function(V).interpolate(Constant(as_vector([0.0, 1.0, 0.0]))) self.z_hat = Function(V).interpolate(Constant(as_vector([0.0, 0.0, 1.0]))) self.R = sqrt(self.x**2 + self.y**2) # distance from z axis self.r = sqrt(self.x**2 + self.y**2 + self.z**2) # distance from origin - self.f = state.fields(self.fname) + self.f = eqn.fields(self.fname) if np.prod(self.f.ufl_shape) != 3: raise ValueError('Components can only be found of a vector function space in 3D.') @@ -435,12 +445,12 @@ def name(self): """Gives the name of this diagnostic field.""" return self.fname+"_meridional" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. @@ -458,12 +468,12 @@ def name(self): """Gives the name of this diagnostic field.""" return self.fname+"_zonal" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. @@ -479,12 +489,12 @@ def name(self): """Gives the name of this diagnostic field.""" return self.fname+"_radial" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. @@ -508,31 +518,31 @@ def __init__(self, density_field, factor=1.): self.density_field = density_field self.factor = Constant(factor) - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ rho_grad = self.density_field+"_gradient" - super().setup(state) - self.grad_density = state.fields(rho_grad) - self.gradu = state.fields("u_gradient") + super().setup(eqn) + self.grad_density = eqn.fields(rho_grad) + self.gradu = eqn.fields("u_gradient") - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ denom = 0. - z_dim = state.mesh.geometric_dimension() - 1 - u_dim = state.fields("u").ufl_shape[0] + z_dim = eqn.domain.mesh.geometric_dimension() - 1 + u_dim = eqn.fields("u").ufl_shape[0] for i in range(u_dim-1): denom += self.gradu[i, z_dim]**2 Nsq = self.factor*self.grad_density[z_dim] @@ -565,17 +575,17 @@ class KineticEnergy(Energy): """Diagnostic kinetic energy density.""" name = "KineticEnergy" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - u = state.fields("u") + u = eqn.fields("u") energy = self.kinetic(u) return self.field.interpolate(energy) @@ -584,18 +594,18 @@ class ShallowWaterKineticEnergy(Energy): """Diagnostic shallow-water kinetic energy density.""" name = "ShallowWaterKineticEnergy" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - u = state.fields("u") - D = state.fields("D") + u = eqn.fields("u") + D = eqn.fields("D") energy = self.kinetic(u, D) return self.field.interpolate(energy) @@ -604,18 +614,18 @@ class ShallowWaterPotentialEnergy(Energy): """Diagnostic shallow-water potential energy density.""" name = "ShallowWaterPotentialEnergy" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - g = state.parameters.g - D = state.fields("D") + g = eqn.parameters.g + D = eqn.fields("D") energy = 0.5*g*D**2 return self.field.interpolate(energy) @@ -637,28 +647,28 @@ def name(self): base_name = "SWPotentialEnstrophy" return "_from_".join((base_name, self.base_field_name)) - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ if self.base_field_name == "PotentialVorticity": - pv = state.fields("PotentialVorticity") - D = state.fields("D") + pv = eqn.fields("PotentialVorticity") + D = eqn.fields("D") enstrophy = 0.5*pv**2*D elif self.base_field_name == "RelativeVorticity": - zeta = state.fields("RelativeVorticity") - D = state.fields("D") - f = state.fields("coriolis") + zeta = eqn.fields("RelativeVorticity") + D = eqn.fields("D") + f = eqn.fields("coriolis") enstrophy = 0.5*(zeta + f)**2/D elif self.base_field_name == "AbsoluteVorticity": - zeta_abs = state.fields("AbsoluteVorticity") - D = state.fields("D") + zeta_abs = eqn.fields("AbsoluteVorticity") + D = eqn.fields("D") enstrophy = 0.5*(zeta_abs)**2/D else: raise ValueError("Don't know how to compute enstrophy with base_field_name=%s; base_field_name should be %s %s or %s." % (self.base_field_name, "RelativeVorticity", "AbsoluteVorticity", "PotentialVorticity")) @@ -669,18 +679,18 @@ class CompressibleKineticEnergy(Energy): """Diagnostic (dry) compressible kinetic energy density.""" name = "CompressibleKineticEnergy" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - u = state.fields("u") - rho = state.fields("rho") + u = eqn.fields("u") + rho = eqn.fields("rho") energy = self.kinetic(u, rho) return self.field.interpolate(energy) @@ -710,30 +720,30 @@ def name(self): else: return "Exner" - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ if not self._initialised: - space = state.spaces("CG1", "CG", 1) - super(Exner, self).setup(state, space=space) + space = eqn.domain.spaces("CG1", "CG", 1) + super(Exner, self).setup(eqn, space=space) - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - rho = state.fields(self.rho_name) - theta = state.fields(self.theta_name) - exner = thermodynamics.exner_pressure(state.parameters, rho, theta) + rho = eqn.fields(self.rho_name) + theta = eqn.fields(self.theta_name) + exner = thermodynamics.exner_pressure(eqn.parameters, rho, theta) return self.field.interpolate(exner) @@ -754,29 +764,29 @@ def name(self): """Gives the name of this diagnostic field.""" return self.field1+"_plus_"+self.field2 - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ if not self._initialised: - space = state.fields(self.field1).function_space() - super(Sum, self).setup(state, space=space) + space = eqn.fields(self.field1).function_space() + super(Sum, self).setup(eqn, space=space) - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - field1 = state.fields(self.field1) - field2 = state.fields(self.field2) + field1 = eqn.fields(self.field1) + field2 = eqn.fields(self.field2) return self.field.assign(field1 + field2) @@ -797,45 +807,45 @@ def name(self): """Gives the name of this diagnostic field.""" return self.field1+"_minus_"+self.field2 - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ if not self._initialised: - space = state.fields(self.field1).function_space() - super(Difference, self).setup(state, space=space) + space = eqn.fields(self.field1).function_space() + super(Difference, self).setup(eqn, space=space) - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - field1 = state.fields(self.field1) - field2 = state.fields(self.field2) + field1 = eqn.fields(self.field1) + field2 = eqn.fields(self.field2) return self.field.assign(field1 - field2) class SteadyStateError(Difference): """Base diagnostic for computing the steady-state error in a field.""" - def __init__(self, state, name): + def __init__(self, eqn, name): """ Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. name (str): name of the field to take the perturbation of. """ DiagnosticField.__init__(self) self.field1 = name self.field2 = name+'_init' - field1 = state.fields(name) - field2 = state.fields(self.field2, field1.function_space()) + field1 = eqn.fields(name) + field2 = eqn.fields(self.field2, field1.function_space()) field2.assign(field1) @property @@ -864,52 +874,53 @@ def name(self): class ThermodynamicDiagnostic(DiagnosticField): """Base thermodynamic diagnostic field, computing many common fields.""" - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquationSet`): the model's equation. """ if not self._initialised: - space = state.fields("theta").function_space() + domain = eqn.domain + space = domain.spaces('theta') h_deg = space.ufl_element().degree()[0] v_deg = space.ufl_element().degree()[1]-1 boundary_method = BoundaryMethod.extruded if (v_deg == 0 and h_deg == 0) else None - super().setup(state, space=space) + super().setup(eqn, space=space) # now let's attach all of our fields - self.u = state.fields("u") - self.rho = state.fields("rho") - self.theta = state.fields("theta") + self.u = eqn.fields("u") + self.rho = eqn.fields("rho") + self.theta = eqn.fields("theta") self.rho_averaged = Function(space) self.recoverer = Recoverer(self.rho, self.rho_averaged, boundary_method=boundary_method) try: - self.r_v = state.fields("water_vapour") + self.r_v = eqn.fields("water_vapour") except NotImplementedError: self.r_v = Constant(0.0) try: - self.r_c = state.fields("cloud_water") + self.r_c = eqn.fields("cloud_water") except NotImplementedError: self.r_c = Constant(0.0) try: - self.rain = state.fields("rain") + self.rain = eqn.fields("rain") except NotImplementedError: self.rain = Constant(0.0) # now let's store the most common expressions - self.exner = thermodynamics.exner_pressure(state.parameters, self.rho_averaged, self.theta) - self.T = thermodynamics.T(state.parameters, self.theta, self.exner, r_v=self.r_v) - self.p = thermodynamics.p(state.parameters, self.exner) + self.exner = thermodynamics.exner_pressure(eqn.parameters, self.rho_averaged, self.theta) + self.T = thermodynamics.T(eqn.parameters, self.theta, self.exner, r_v=self.r_v) + self.p = thermodynamics.p(eqn.parameters, self.exner) self.r_l = self.r_c + self.rain self.r_t = self.r_v + self.r_c + self.rain - def compute(self, state): + def compute(self, eqn): """ Compute thermodynamic auxiliary fields commonly used by diagnostics. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ self.recoverer.project() @@ -918,84 +929,84 @@ class Theta_e(ThermodynamicDiagnostic): """The moist equivalent potential temperature diagnostic field.""" name = "Theta_e" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - super().compute(state) + super().compute(eqn) - return self.field.interpolate(thermodynamics.theta_e(state.parameters, self.T, self.p, self.r_v, self.r_t)) + return self.field.interpolate(thermodynamics.theta_e(eqn.parameters, self.T, self.p, self.r_v, self.r_t)) class InternalEnergy(ThermodynamicDiagnostic): """The moist compressible internal energy density.""" name = "InternalEnergy" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - super().compute(state) + super().compute(eqn) - return self.field.interpolate(thermodynamics.internal_energy(state.parameters, self.rho_averaged, self.T, r_v=self.r_v, r_l=self.r_l)) + return self.field.interpolate(thermodynamics.internal_energy(eqn.parameters, self.rho_averaged, self.T, r_v=self.r_v, r_l=self.r_l)) class PotentialEnergy(ThermodynamicDiagnostic): """The moist compressible potential energy density.""" name = "PotentialEnergy" - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ - super().setup(state) - self.x = SpatialCoordinate(state.mesh) + super().setup(eqn) + self.x = SpatialCoordinate(eqn.domain.mesh) - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - super().compute(state) + super().compute(eqn) - return self.field.interpolate(self.rho_averaged * (1 + self.r_t) * state.parameters.g * dot(self.x, state.k)) + return self.field.interpolate(self.rho_averaged * (1 + self.r_t) * eqn.parameters.g * dot(self.x, eqn.domain.k)) class ThermodynamicKineticEnergy(ThermodynamicDiagnostic): """The moist compressible kinetic energy density.""" name = "ThermodynamicKineticEnergy" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - super().compute(state) + super().compute(eqn) return self.field.project(0.5 * self.rho_averaged * (1 + self.r_t) * dot(self.u, self.u)) @@ -1004,36 +1015,36 @@ class Dewpoint(ThermodynamicDiagnostic): """The dewpoint temperature diagnostic field.""" name = "Dewpoint" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - super().compute(state) + super().compute(eqn) - return self.field.interpolate(thermodynamics.T_dew(state.parameters, self.p, self.r_v)) + return self.field.interpolate(thermodynamics.T_dew(eqn.parameters, self.p, self.r_v)) class Temperature(ThermodynamicDiagnostic): """The absolute temperature diagnostic field.""" name = "Temperature" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - super().compute(state) + super().compute(eqn) return self.field.assign(self.T) @@ -1042,55 +1053,55 @@ class Theta_d(ThermodynamicDiagnostic): """The dry potential temperature diagnostic field.""" name = "Theta_d" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - super().compute(state) + super().compute(eqn) - return self.field.interpolate(self.theta / (1 + self.r_v * state.parameters.R_v / state.parameters.R_d)) + return self.field.interpolate(self.theta / (1 + self.r_v * eqn.parameters.R_v / eqn.parameters.R_d)) class RelativeHumidity(ThermodynamicDiagnostic): """The relative humidity diagnostic field.""" name = "RelativeHumidity" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - super().compute(state) + super().compute(eqn) - return self.field.interpolate(thermodynamics.RH(state.parameters, self.r_v, self.T, self.p)) + return self.field.interpolate(thermodynamics.RH(eqn.parameters, self.r_v, self.T, self.p)) class Pressure(ThermodynamicDiagnostic): """The pressure field computed in the 'theta' space.""" name = "Pressure_Vt" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - super().compute(state) + super().compute(eqn) return self.field.assign(self.p) @@ -1099,17 +1110,17 @@ class Exner_Vt(ThermodynamicDiagnostic): """The Exner pressure field computed in the 'theta' space.""" name = "Exner_Vt" - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - super().compute(state) + super().compute(eqn) return self.field.assign(self.exner) @@ -1118,26 +1129,26 @@ class HydrostaticImbalance(DiagnosticField): """Hydrostatic imbalance diagnostic field.""" name = "HydrostaticImbalance" - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ if not self._initialised: - Vu = state.spaces("HDiv") - space = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1]) - super().setup(state, space=space) - rho = state.fields("rho") - rhobar = state.fields("rhobar") - theta = state.fields("theta") - thetabar = state.fields("thetabar") - exner = thermodynamics.exner_pressure(state.parameters, rho, theta) - exnerbar = thermodynamics.exner_pressure(state.parameters, rhobar, thetabar) - - cp = Constant(state.parameters.cp) - n = FacetNormal(state.mesh) + Vu = eqn.domain.spaces("HDiv") + space = FunctionSpace(eqn.domain.mesh, Vu.ufl_element()._elements[-1]) + super().setup(eqn, space=space) + rho = eqn.fields("rho") + rhobar = eqn.fields("rhobar") + theta = eqn.fields("theta") + thetabar = eqn.fields("thetabar") + exner = thermodynamics.exner_pressure(eqn.parameters, rho, theta) + exnerbar = thermodynamics.exner_pressure(eqn.parameters, rhobar, thetabar) + + cp = Constant(eqn.parameters.cp) + n = FacetNormal(eqn.domain.mesh) F = TrialFunction(space) w = TestFunction(space) @@ -1153,12 +1164,12 @@ def setup(self, state): imbalanceproblem = LinearVariationalProblem(a, L, self.field, bcs=bcs) self.imbalance_solver = LinearVariationalSolver(imbalanceproblem) - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. @@ -1171,23 +1182,23 @@ class Precipitation(DiagnosticField): """The total precipitation falling through the domain's bottom surface.""" name = "Precipitation" - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ if not self._initialised: - space = state.spaces("DG0", "DG", 0) - super().setup(state, space=space) + space = eqn.domain.spaces("DG0", "DG", 0) + super().setup(eqn, space=space) - rain = state.fields('rain') - rho = state.fields('rho') - v = state.fields('rainfall_velocity') + rain = eqn.fields('rain') + rho = eqn.fields('rho') + v = eqn.fields('rainfall_velocity') self.phi = TestFunction(space) flux = TrialFunction(space) - n = FacetNormal(state.mesh) + n = FacetNormal(eqn.domain.mesh) un = 0.5 * (dot(v, n) + abs(dot(v, n))) self.flux = Function(space) @@ -1202,12 +1213,12 @@ def setup(self, state): problem = LinearVariationalProblem(a, L, self.flux) self.solver = LinearVariationalSolver(problem) - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. @@ -1220,12 +1231,12 @@ def compute(self, state): class Vorticity(DiagnosticField): """Base diagnostic field class for shallow-water vorticity variables.""" - def setup(self, state, vorticity_type=None): + def setup(self, eqn, vorticity_type=None): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. vorticity_type (str, optional): denotes which type of vorticity to be computed ('relative', 'absolute' or 'potential'). Defaults to None. @@ -1235,36 +1246,36 @@ def setup(self, state, vorticity_type=None): if vorticity_type not in vorticity_types: raise ValueError("vorticity type must be one of %s, not %s" % (vorticity_types, vorticity_type)) try: - space = state.spaces("CG") + space = eqn.domain.spaces("CG") except ValueError: - dgspace = state.spaces("DG") + dgspace = eqn.domain.spaces("DG") cg_degree = dgspace.ufl_element().degree() + 2 - space = FunctionSpace(state.mesh, "CG", cg_degree) - super().setup(state, space=space) - u = state.fields("u") + space = FunctionSpace(eqn.domain.mesh, "CG", cg_degree) + super().setup(eqn, space=space) + u = eqn.fields("u") gamma = TestFunction(space) q = TrialFunction(space) if vorticity_type == "potential": - D = state.fields("D") + D = eqn.fields("D") a = q*gamma*D*dx else: a = q*gamma*dx - L = (- inner(state.perp(grad(gamma)), u))*dx + L = (- inner(eqn.domain.perp(grad(gamma)), u))*dx if vorticity_type != "relative": - f = state.fields("coriolis") + f = eqn.fields("coriolis") L += gamma*f*dx problem = LinearVariationalProblem(a, L, self.field) self.solver = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"}) - def compute(self, state): + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. @@ -1277,39 +1288,39 @@ class PotentialVorticity(Vorticity): u"""Diagnostic field for shallow-water potential vorticity, q=(∇×(u+f))/D""" name = "PotentialVorticity" - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ - super().setup(state, vorticity_type="potential") + super().setup(eqn, vorticity_type="potential") class AbsoluteVorticity(Vorticity): u"""Diagnostic field for absolute vorticity, ζ=∇×(u+f)""" name = "AbsoluteVorticity" - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ - super().setup(state, vorticity_type="absolute") + super().setup(eqn, vorticity_type="absolute") class RelativeVorticity(Vorticity): u"""Diagnostic field for relative vorticity, ζ=∇×u""" name = "RelativeVorticity" - def setup(self, state): + def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. """ - super().setup(state, vorticity_type="relative") + super().setup(eqn, vorticity_type="relative") diff --git a/gusto/domain.py b/gusto/domain.py index ff8d9d8e0..346ef0d35 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -3,7 +3,7 @@ the set of compatible function spaces defined upon it. """ -from gusto.function_spaces import Spaces +from gusto.function_spaces import Spaces, check_degree_args from firedrake import (Constant, SpatialCoordinate, sqrt, CellNormal, cross, as_vector, inner, interpolate) @@ -39,17 +39,7 @@ def __init__(self, mesh, family, degree=None, both "degree" and "horizontal_degree"). """ - # Checks on degree arguments - if degree is None and horizontal_degree is None: - raise ValueError('Either "degree" or "horizontal_degree" must be passed to Domain') - if mesh.extruded and degree is None and vertical_degree is None: - raise ValueError('For extruded meshes, either degree or "vertical_degree" must be passed to Domain') - if degree is not None and horizontal_degree is not None: - raise ValueError('Cannot pass both "degree" and "horizontal_degree" to Domain') - if mesh.extruded and degree is not None and vertical_degree is not None: - raise ValueError('Cannot pass both "degree" and "vertical_degree" to Domain') - if not mesh.extruded and vertical_degree is not None: - raise ValueError('Cannot pass "vertical_degree" to Domain if mesh is not extruded') + check_degree_args('Domain', mesh, degree, horizontal_degree, vertical_degree) # Get degrees self.horizontal_degree = degree if horizontal_degree is None else horizontal_degree diff --git a/gusto/forcing.py b/gusto/forcing.py index ee9cd8378..ef97d2142 100644 --- a/gusto/forcing.py +++ b/gusto/forcing.py @@ -21,7 +21,7 @@ class Forcing(object): discretisation. """ - def __init__(self, equation, alpha): + def __init__(self, equation, io, alpha): """ Args: equation (:class:`PrognosticEquationSet`): the prognostic equations @@ -33,7 +33,7 @@ def __init__(self, equation, alpha): self.field_name = equation.field_name implicit_terms = ["incompressibility", "sponge"] - dt = equation.state.dt + dt = io.dt W = equation.function_space self.x0 = Function(W) diff --git a/gusto/function_spaces.py b/gusto/function_spaces.py index 3b68e485d..956c71cc5 100644 --- a/gusto/function_spaces.py +++ b/gusto/function_spaces.py @@ -24,8 +24,8 @@ def __init__(self, mesh): self.extruded_mesh = hasattr(mesh, "_base_mesh") self._initialised_base_spaces = False - def __call__(self, name, family=None, horizontal_degree=None, - vertical_degree=None, V=None): + def __call__(self, name, family=None, degree=None, + horizontal_degree=None, vertical_degree=None, V=None): """ Returns a space, and also creates it if it is not created yet. @@ -33,10 +33,18 @@ def __call__(self, name, family=None, horizontal_degree=None, family and degree) need to be provided. Alternatively a space can be passed in to be stored in the creator. + For extruded meshes, it is possible to seperately specify the horizontal + and vertical degrees of the elements. Alternatively, if these degrees + should be the same then this can be specified through the "degree" + argument. + Args: name (str): the name of the space. family (str, optional): name of the finite element family to be created. Defaults to None. + degree (int, optional): the element degree used for the space. + Defaults to None, in which case the horizontal degree must be + provided. horizontal_degree (int, optional): the horizontal degree of the finite element space to be created. Defaults to None. vertical_degree (int, optional): the vertical degree of the @@ -62,15 +70,15 @@ def __call__(self, name, family=None, horizontal_degree=None, value = V elif name == "DG1_equispaced": - # Special case based on name - if self.extruded_mesh: - value = self.build_dg_space(1, 1, variant='equispaced') - else: - value = self.build_dg_space(1, variant='equispaced') + # Special case as no degree arguments need providing + value = self.build_dg_space(1, 1, variant='equispaced') else: - # Need to create space, based on name/family/degree - assert horizontal_degree is not None + check_degree_args('Spaces', self.mesh, degree, horizontal_degree, vertical_degree) + + # Convert to horizontal and vertical degrees + horizontal_degree = degree if horizontal_degree is None else horizontal_degree + vertical_degree = degree if vertical_degree is None else vertical_degree # Loop through name and family combinations if name == "HDiv" and family in ["BDM", "RT", "CG", "RTCF"]: @@ -244,10 +252,10 @@ def build_cg_space(self, horizontal_degree, vertical_degree): Args: horizontal_degree (int): the polynomial degree of the horizontal - part of the DG space from the de Rham complex. + part of the CG space. vertical_degree (int, optional): the polynomial degree of the - vertical part of the the DG space from the de Rham complex. - Defaults to None. Must be specified if the mesh is extruded. + vertical part of the the CG space. Defaults to None. Must be + specified if the mesh is extruded. Returns: :class:`FunctionSpace`: the continuous space. @@ -257,14 +265,43 @@ def build_cg_space(self, horizontal_degree, vertical_degree): if vertical_degree is None: raise ValueError('vertical_degree must be specified to create CG space on an extruded mesh') cell = self.mesh._base_mesh.ufl_cell().cellname() - CG_hori = FiniteElement("CG", cell, horizontal_degree+1) - CG_vert = FiniteElement("CG", interval, vertical_degree+1) + CG_hori = FiniteElement("CG", cell, horizontal_degree) + CG_vert = FiniteElement("CG", interval, vertical_degree) V_elt = TensorProductElement(CG_hori, CG_vert) else: cell = self.mesh.ufl_cell().cellname() - V_elt = FiniteElement("DG", cell, horizontal_degree+1, variant=variant) + V_elt = FiniteElement("CG", cell, horizontal_degree) # How should we name this if the horizontal and vertical degrees are different? - name = f'CG{horizontal_degree+1}' + name = f'CG{horizontal_degree}' return FunctionSpace(self.mesh, V_elt, name=name) + + +def check_degree_args(name, mesh, degree, horizontal_degree, vertical_degree): + """ + Check the degree arguments passed to either the :class:`Domain` or the + :class:`Spaces` object. This will raise errors if the arguments are not + appropriate. + + Args: + name (str): name of object to print out. + mesh (:class:`Mesh`): the model's mesh. + degree (int): the element degree. + horizontal_degree (int): the element degree used for the horizontal part + of a space. + vertical_degree (int): the element degree used for the vertical part + of a space. + """ + + # Checks on degree arguments + if degree is None and horizontal_degree is None: + raise ValueError(f'Either "degree" or "horizontal_degree" must be passed to {name}') + if mesh.extruded and degree is None and vertical_degree is None: + raise ValueError(f'For extruded meshes, either degree or "vertical_degree" must be passed to {name}') + if degree is not None and horizontal_degree is not None: + raise ValueError(f'Cannot pass both "degree" and "horizontal_degree" to {name}') + if mesh.extruded and degree is not None and vertical_degree is not None: + raise ValueError(f'Cannot pass both "degree" and "vertical_degree" to {name}') + if not mesh.extruded and vertical_degree is not None: + raise ValueError(f'Cannot pass "vertical_degree" to {name} if mesh is not extruded') diff --git a/gusto/io.py b/gusto/io.py index c77cdb576..7cb91d095 100644 --- a/gusto/io.py +++ b/gusto/io.py @@ -205,10 +205,9 @@ def __init__(self, domain, equation, dt, self.diagnostic_fields = [] # TODO: quick way of ensuring that diagnostics are registered - if hasattr(self, "field_names"): + if hasattr(equation, "field_names"): for fname in equation.field_names: self.diagnostics.register(fname) - self.bcs[fname] = [] else: self.diagnostics.register(equation.field_name) @@ -250,7 +249,8 @@ def setup_diagnostics(self): schedule = topo_sort(field_deps) self.diagnostic_fields = schedule for diagnostic in self.diagnostic_fields: - diagnostic.setup(self) + # TODO: for diagnostics to see equation and IO, change the setup here + diagnostic.setup(self.equation) self.diagnostics.register(diagnostic.name) def setup_dump(self, t, tmax, pickup=False): @@ -411,7 +411,7 @@ def dump(self, t): # Diagnostics: # Compute diagnostic fields for field in self.diagnostic_fields: - field(self) + field(self.equation) if output.dump_diagnostics: # Output diagnostic data diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index 29dd344f2..c19d1ea5c 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -28,7 +28,7 @@ class TimesteppingSolver(object, metaclass=ABCMeta): """Base class for timestepping linear solvers for Gusto.""" - def __init__(self, equations, alpha=0.5, solver_parameters=None, + def __init__(self, equations, io, alpha=0.5, solver_parameters=None, overwrite_solver_parameters=False): """ Args: @@ -44,6 +44,7 @@ def __init__(self, equations, alpha=0.5, solver_parameters=None, passed in. Defaults to False. """ self.equations = equations + self.dt = io.dt self.alpha = alpha if solver_parameters is not None: @@ -120,7 +121,7 @@ class CompressibleSolver(TimesteppingSolver): 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}}} - def __init__(self, equations, alpha=0.5, + def __init__(self, equations, io, alpha=0.5, quadrature_degree=None, solver_parameters=None, overwrite_solver_parameters=False, moisture=None): """ @@ -161,13 +162,14 @@ def __init__(self, equations, alpha=0.5, # Turn monitor on for the trace system self.solver_parameters["condensed_field"]["ksp_monitor_true_residual"] = None - super().__init__(equations, alpha, solver_parameters, + super().__init__(equations, io, alpha, solver_parameters, overwrite_solver_parameters) @timed_function("Gusto:SolverSetup") def _setup_solver(self): - dt = state.dt + equations = self.equations + dt = self.dt beta_ = dt*self.alpha cp = equations.parameters.cp Vu = equations.domain.spaces("HDiv") @@ -192,7 +194,7 @@ def _setup_solver(self): w, phi, dl = TestFunctions(M) u, rho, l0 = TrialFunctions(M) - n = FacetNormal(equations.domin.mesh) + n = FacetNormal(equations.domain.mesh) # Get background fields thetabar = equations.fields("thetabar") @@ -421,7 +423,7 @@ class IncompressibleSolver(TimesteppingSolver): @timed_function("Gusto:SolverSetup") def _setup_solver(self): equation = self.equation # just cutting down line length a bit - dt = state.dt + dt = self.dt beta_ = dt*self.alpha Vu = equation.domain.spaces("HDiv") Vb = equation.domain.spaces("theta") @@ -548,7 +550,7 @@ class LinearTimesteppingSolver(object): 'sub_pc_type': 'ilu'}} } - def __init__(self, equation, alpha): + def __init__(self, equation, io, alpha): """ Args: equation (:class:`PrognosticEquation`): the model's equation object. @@ -560,7 +562,7 @@ def __init__(self, equation, alpha): lambda t: Term(t.get(linearisation).form, t.labels), drop) - dt = state.dt + dt = io.dt W = equation.function_space beta = dt*alpha diff --git a/gusto/physics.py b/gusto/physics.py index dac6e69de..f45c63cba 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -57,7 +57,7 @@ class SaturationAdjustment(Physics): """ def __init__(self, equation, vapour_name='water_vapour', - cloud_name='cloud_water', latent_heat=True): + cloud_name='cloud_water', latent_heat=True, parameters=None): """ Args: equation (:class:`PrognosticEquationSet`): the model's equation. @@ -67,6 +67,9 @@ def __init__(self, equation, vapour_name='water_vapour', 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 @@ -85,7 +88,7 @@ def __init__(self, equation, vapour_name='water_vapour', # Make prognostic for physics scheme self.X = Function(equation.X.function_space()) self.equation = equation - parameters = equation.parameters + parameters = equation.parameters if parameters is None else parameters self.latent_heat = latent_heat # Vapour and cloud variables are needed for every form of this scheme @@ -267,7 +270,7 @@ def __init__(self, equation, rain_name, domain, moments=AdvectedMoments.M3): Vu = domain.spaces("HDiv") # TODO: how do we allow this to be output? - v = Function(Vu, name='rainfall_velocity') + v = equation.fields(name='rainfall_velocity', space=Vu) # -------------------------------------------------------------------- # # Create physics term -- which is actually a transport term diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 3fce49b57..e72924061 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -265,7 +265,7 @@ def __init__(self, equation_set, io, transport_schemes, assert scheme.field_name in equation_set.field_names self.diffusion_schemes.append((scheme.field_name, scheme)) - if not equation_set.reference_profile_initialised: + if not equation_set.reference_profiles_initialised: raise RuntimeError('Reference profiles for equation set must be initialised to use Semi-Implicit Timestepper') super().__init__(equation_set, io) @@ -295,10 +295,10 @@ def __init__(self, equation_set, io, transport_schemes, self.xrhs = Function(W) self.dy = Function(W) if linear_solver is None: - self.linear_solver = LinearTimesteppingSolver(equation_set, self.alpha) + self.linear_solver = LinearTimesteppingSolver(equation_set, io, self.alpha) else: self.linear_solver = linear_solver - self.forcing = Forcing(equation_set, self.alpha) + self.forcing = Forcing(equation_set, io, self.alpha) self.bcs = equation_set.bcs def _apply_bcs(self): diff --git a/integration-tests/balance/test_compressible_balance.py b/integration-tests/balance/test_compressible_balance.py index bfd7829e4..3e62638a5 100644 --- a/integration-tests/balance/test_compressible_balance.py +++ b/integration-tests/balance/test_compressible_balance.py @@ -24,41 +24,37 @@ def setup_balance(dirname): m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) + domain = Domain(mesh, "CG", 1) output = OutputParameters(dirname=dirname+'/dry_balance', dumpfreq=10, dumplist=['u']) parameters = CompressibleParameters() - - state = State(mesh, - dt=dt, - output=output, - parameters=parameters) - - eqns = CompressibleEulerEquations(state, "CG", 1) + eqns = CompressibleEulerEquations(domain, parameters) + io = IO(domain, eqns, dt=dt, output=output) # Initial conditions - rho0 = state.fields("rho") - theta0 = state.fields("theta") + rho0 = eqns.fields("rho") + theta0 = eqns.fields("theta") # Isentropic background state Tsurf = Constant(300.) theta0.interpolate(Tsurf) # Calculate hydrostatic exner - compressible_hydrostatic_balance(state, theta0, rho0, solve_for_rho=True) + compressible_hydrostatic_balance(eqns, theta0, rho0, solve_for_rho=True) - state.set_reference_profiles([('rho', rho0), - ('theta', theta0)]) + eqns.set_reference_profiles([('rho', rho0), + ('theta', theta0)]) # Set up transport schemes - transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "rho"), - SSPRK3(state, "theta", options=EmbeddedDGOptions())] + transported_fields = [ImplicitMidpoint(domain, io, "u"), + SSPRK3(domain, io, "rho"), + SSPRK3(domain, io, "theta", options=EmbeddedDGOptions())] # Set up linear solver - linear_solver = CompressibleSolver(state, eqns) + linear_solver = CompressibleSolver(eqns, io) # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, linear_solver=linear_solver) return stepper, tmax diff --git a/integration-tests/balance/test_saturated_balance.py b/integration-tests/balance/test_saturated_balance.py index e5c044c35..e40e02ef7 100644 --- a/integration-tests/balance/test_saturated_balance.py +++ b/integration-tests/balance/test_saturated_balance.py @@ -24,22 +24,11 @@ def setup_saturated(dirname, recovered): nlayers = int(H/deltax) ncolumns = int(L/deltax) - m = PeriodicIntervalMesh(ncolumns, L) - mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) - - # option to easily change between recovered and not if necessary - # default should be to use lowest order set of spaces degree = 0 if recovered else 1 - output = OutputParameters(dirname=dirname+'/saturated_balance', dumpfreq=1, dumplist=['u']) - parameters = CompressibleParameters() - diagnostic_fields = [Theta_e()] - - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) + m = PeriodicIntervalMesh(ncolumns, L) + mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) + domain = Domain(mesh, "CG", degree) tracers = [WaterVapour(), CloudWater()] @@ -47,14 +36,19 @@ def setup_saturated(dirname, recovered): u_transport_option = "vector_advection_form" else: u_transport_option = "vector_invariant_form" + parameters = CompressibleParameters() eqns = CompressibleEulerEquations( - state, "CG", degree, u_transport_option=u_transport_option, active_tracers=tracers) + domain, parameters, u_transport_option=u_transport_option, active_tracers=tracers) + + output = OutputParameters(dirname=dirname+'/saturated_balance', dumpfreq=1, dumplist=['u']) + diagnostic_fields = [Theta_e()] + io = IO(domain, eqns, dt=dt, output=output, diagnostic_fields=diagnostic_fields) # Initial conditions - rho0 = state.fields("rho") - theta0 = state.fields("theta") - water_v0 = state.fields("water_vapour") - water_c0 = state.fields("cloud_water") + rho0 = eqns.fields("rho") + theta0 = eqns.fields("theta") + water_v0 = eqns.fields("water_vapour") + water_c0 = eqns.fields("cloud_water") moisture = ['water_vapour', 'cloud_water'] # spaces @@ -67,15 +61,14 @@ def setup_saturated(dirname, recovered): water_t = Function(Vt).interpolate(total_water) # Calculate hydrostatic exner - saturated_hydrostatic_balance(state, theta_e, water_t) + saturated_hydrostatic_balance(eqns, theta_e, water_t) water_c0.assign(water_t - water_v0) - state.set_reference_profiles([('rho', rho0), - ('theta', theta0)]) + eqns.set_reference_profiles([('rho', rho0), ('theta', theta0)]) # Set up transport schemes if recovered: - VDG1 = state.spaces("DG1_equispaced") + VDG1 = domain.spaces("DG1_equispaced") VCG1 = FunctionSpace(mesh, "CG", 1) Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) @@ -99,23 +92,23 @@ def setup_saturated(dirname, recovered): wv_opts = EmbeddedDGOptions() wc_opts = EmbeddedDGOptions() - transported_fields = [SSPRK3(state, 'rho', options=rho_opts), - SSPRK3(state, 'theta', options=theta_opts), - SSPRK3(state, 'water_vapour', options=wv_opts), - SSPRK3(state, 'cloud_water', options=wc_opts)] + transported_fields = [SSPRK3(domain, io, 'rho', options=rho_opts), + SSPRK3(domain, io, 'theta', options=theta_opts), + SSPRK3(domain, io, 'water_vapour', options=wv_opts), + SSPRK3(domain, io, 'cloud_water', options=wc_opts)] if recovered: - transported_fields.append(SSPRK3(state, 'u', options=u_opts)) + transported_fields.append(SSPRK3(domain, io, 'u', options=u_opts)) else: - transported_fields.append(ImplicitMidpoint(state, 'u')) + transported_fields.append(ImplicitMidpoint(domain, io, 'u')) - linear_solver = CompressibleSolver(state, eqns, moisture=moisture) + linear_solver = CompressibleSolver(eqns, io, moisture=moisture) # add physics - physics_schemes = [(SaturationAdjustment(eqns, parameters), ForwardEuler(state))] + physics_schemes = [(SaturationAdjustment(eqns), ForwardEuler(domain, io))] # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, linear_solver=linear_solver, physics_schemes=physics_schemes) diff --git a/integration-tests/balance/test_unsaturated_balance.py b/integration-tests/balance/test_unsaturated_balance.py index d6d2dd5c5..fa3715067 100644 --- a/integration-tests/balance/test_unsaturated_balance.py +++ b/integration-tests/balance/test_unsaturated_balance.py @@ -24,20 +24,11 @@ def setup_unsaturated(dirname, recovered): nlayers = int(H/deltax) ncolumns = int(L/deltax) - m = PeriodicIntervalMesh(ncolumns, L) - mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) - degree = 0 if recovered else 1 - output = OutputParameters(dirname=dirname+'/unsaturated_balance', dumpfreq=1) - parameters = CompressibleParameters() - diagnostic_fields = [Theta_d(), RelativeHumidity()] - - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) + m = PeriodicIntervalMesh(ncolumns, L) + mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) + domain = Domain(mesh, "CG", degree) tracers = [WaterVapour(), CloudWater()] @@ -45,12 +36,17 @@ def setup_unsaturated(dirname, recovered): u_transport_option = "vector_advection_form" else: u_transport_option = "vector_invariant_form" + parameters = CompressibleParameters() eqns = CompressibleEulerEquations( - state, "CG", degree, u_transport_option=u_transport_option, active_tracers=tracers) + domain, parameters, u_transport_option=u_transport_option, active_tracers=tracers) + + output = OutputParameters(dirname=dirname+'/unsaturated_balance', dumpfreq=1) + diagnostic_fields = [Theta_d(), RelativeHumidity()] + io = IO(domain, eqns, dt=dt, output=output, diagnostic_fields=diagnostic_fields) # Initial conditions - rho0 = state.fields("rho") - theta0 = state.fields("theta") + rho0 = eqns.fields("rho") + theta0 = eqns.fields("theta") moisture = ['water_vapour', 'cloud_water'] # spaces @@ -63,14 +59,13 @@ def setup_unsaturated(dirname, recovered): RH = Function(Vt).interpolate(humidity) # Calculate hydrostatic exner - unsaturated_hydrostatic_balance(state, theta_d, RH) + unsaturated_hydrostatic_balance(eqns, theta_d, RH) - state.set_reference_profiles([('rho', rho0), - ('theta', theta0)]) + eqns.set_reference_profiles([('rho', rho0), ('theta', theta0)]) # Set up transport schemes if recovered: - VDG1 = state.spaces("DG1_equispaced") + VDG1 = domain.spaces("DG1_equispaced") VCG1 = FunctionSpace(mesh, "CG", 1) Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) @@ -87,22 +82,22 @@ def setup_unsaturated(dirname, recovered): rho_opts = None theta_opts = EmbeddedDGOptions() - transported_fields = [SSPRK3(state, "rho", options=rho_opts), - SSPRK3(state, "theta", options=theta_opts), - SSPRK3(state, "water_vapour", options=theta_opts), - SSPRK3(state, "cloud_water", options=theta_opts)] + transported_fields = [SSPRK3(domain, io, "rho", options=rho_opts), + SSPRK3(domain, io, "theta", options=theta_opts), + SSPRK3(domain, io, "water_vapour", options=theta_opts), + SSPRK3(domain, io, "cloud_water", options=theta_opts)] if recovered: - transported_fields.append(SSPRK3(state, "u", options=u_opts)) + transported_fields.append(SSPRK3(domain, io, "u", options=u_opts)) else: - transported_fields.append(ImplicitMidpoint(state, "u")) + transported_fields.append(ImplicitMidpoint(domain, io, "u")) - linear_solver = CompressibleSolver(state, eqns, moisture=moisture) + linear_solver = CompressibleSolver(eqns, io, moisture=moisture) # Set up physics - physics_schemes = [(SaturationAdjustment(eqns, parameters), ForwardEuler(state))] + physics_schemes = [(SaturationAdjustment(eqns), ForwardEuler(domain, io))] # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, linear_solver=linear_solver, physics_schemes=physics_schemes) diff --git a/integration-tests/diffusion/test_diffusion.py b/integration-tests/diffusion/test_diffusion.py index 5be082e25..9560ac5dc 100644 --- a/integration-tests/diffusion/test_diffusion.py +++ b/integration-tests/diffusion/test_diffusion.py @@ -8,18 +8,18 @@ import pytest -def run(equation, diffusion_scheme, state, tmax): +def run(equation, diffusion_scheme, io, tmax): - timestepper = Timestepper(equation, diffusion_scheme, state) + timestepper = Timestepper(equation, diffusion_scheme, io) timestepper.run(0., tmax) - return timestepper.state.fields("f") + return timestepper.equation.fields("f") @pytest.mark.parametrize("DG", [True, False]) def test_scalar_diffusion(tmpdir, DG, tracer_setup): setup = tracer_setup(tmpdir, geometry="slice", blob=True) - state = setup.state + domain = setup.domain f_init = setup.f_init tmax = setup.tmax tol = 5.e-2 @@ -28,20 +28,20 @@ def test_scalar_diffusion(tmpdir, DG, tracer_setup): f_end_expr = (1/(1+4*tmax))*f_init**(1/(1+4*tmax)) if DG: - V = state.spaces("DG", "DG", 1) + V = domain.spaces("DG", "DG", degree=1) else: - V = state.spaces("theta", degree=1) + V = domain.spaces("theta", degree=1) mu = 5. diffusion_params = DiffusionParameters(kappa=kappa, mu=mu) - eqn = DiffusionEquation(state, V, "f", - diffusion_parameters=diffusion_params) + eqn = DiffusionEquation(domain, V, "f", diffusion_parameters=diffusion_params) + io = IO(domain, eqn, dt=setup.dt, output=setup.output) - diffusion_scheme = BackwardEuler(state) + diffusion_scheme = BackwardEuler(domain, io) - state.fields("f").interpolate(f_init) - f_end = run(eqn, diffusion_scheme, state, tmax) + eqn.fields("f").interpolate(f_init) + f_end = run(eqn, diffusion_scheme, io, tmax) assert errornorm(f_end_expr, f_end) < tol @@ -49,7 +49,7 @@ def test_scalar_diffusion(tmpdir, DG, tracer_setup): def test_vector_diffusion(tmpdir, DG, tracer_setup): setup = tracer_setup(tmpdir, geometry="slice", blob=True) - state = setup.state + domain = setup.domain f_init = setup.f_init tmax = setup.tmax tol = 3.e-2 @@ -59,24 +59,24 @@ def test_vector_diffusion(tmpdir, DG, tracer_setup): kappa = Constant([[kappa, 0.], [0., kappa]]) if DG: - V = VectorFunctionSpace(state.mesh, "DG", 1) + V = VectorFunctionSpace(domain.mesh, "DG", 1) else: - V = state.spaces("HDiv", "CG", 1) + V = domain.spaces("HDiv", "CG", 1) f_init = as_vector([f_init, 0.]) f_end_expr = as_vector([f_end_expr, 0.]) mu = 5. diffusion_params = DiffusionParameters(kappa=kappa, mu=mu) - eqn = DiffusionEquation(state, V, "f", - diffusion_parameters=diffusion_params) + eqn = DiffusionEquation(domain, V, "f", diffusion_parameters=diffusion_params) + io = IO(domain, eqn, dt=setup.dt, output=setup.output) if DG: - state.fields("f").interpolate(f_init) + eqn.fields("f").interpolate(f_init) else: - state.fields("f").project(f_init) + eqn.fields("f").project(f_init) - diffusion_scheme = BackwardEuler(state) + diffusion_scheme = BackwardEuler(domain, io) - f_end = run(eqn, diffusion_scheme, state, tmax) + f_end = run(eqn, diffusion_scheme, io, tmax) assert errornorm(f_end_expr, f_end) < tol diff --git a/integration-tests/model/test_checkpointing.py b/integration-tests/model/test_checkpointing.py index e4fc4ddb7..2087cac9c 100644 --- a/integration-tests/model/test_checkpointing.py +++ b/integration-tests/model/test_checkpointing.py @@ -20,43 +20,42 @@ def setup_checkpointing(dirname): # build volume mesh H = 1.0e4 # Height position of the model top mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) + domain = Domain(mesh, "CG", 1) - output = OutputParameters(dirname=dirname, dumpfreq=1, - chkptfreq=2, log_level='INFO') parameters = CompressibleParameters() + eqns = CompressibleEulerEquations(domain, parameters) - state = State(mesh, - dt=dt, - output=output, - parameters=parameters) + output = OutputParameters(dirname=dirname, dumpfreq=1, + chkptfreq=2, log_level='INFO') + io = IO(domain, eqns, dt=dt, output=output) - eqns = CompressibleEulerEquations(state, "CG", 1) + initialise_fields(eqns) # Set up transport schemes transported_fields = [] - transported_fields.append(SSPRK3(state, "u")) - transported_fields.append(SSPRK3(state, "rho")) - transported_fields.append(SSPRK3(state, "theta")) + transported_fields.append(SSPRK3(domain, io, "u")) + transported_fields.append(SSPRK3(domain, io, "rho")) + transported_fields.append(SSPRK3(domain, io, "theta")) # Set up linear solver - linear_solver = CompressibleSolver(state, eqns) + linear_solver = CompressibleSolver(eqns, io) # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, linear_solver=linear_solver) - return state, stepper, dt + return eqns, stepper, dt -def initialise_fields(state): +def initialise_fields(eqns): L = 1.e5 H = 1.0e4 # Height position of the model top # Initial conditions - u0 = state.fields("u") - rho0 = state.fields("rho") - theta0 = state.fields("theta") + u0 = eqns.fields("u") + rho0 = eqns.fields("rho") + theta0 = eqns.fields("theta") # spaces Vt = theta0.function_space() @@ -64,11 +63,11 @@ def initialise_fields(state): # Thermodynamic constants required for setting initial conditions # and reference profiles - g = state.parameters.g - N = state.parameters.N + g = eqns.parameters.g + N = eqns.parameters.N # N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) - x, z = SpatialCoordinate(state.mesh) + x, z = SpatialCoordinate(eqns.domain.mesh) Tsurf = 300. thetab = Tsurf*exp(N**2*z/g) @@ -76,7 +75,7 @@ def initialise_fields(state): rho_b = Function(Vr) # Calculate hydrostatic exner - compressible_hydrostatic_balance(state, theta_b, rho_b) + compressible_hydrostatic_balance(eqns, theta_b, rho_b) a = 5.0e3 deltaTheta = 1.0e-2 @@ -85,34 +84,32 @@ def initialise_fields(state): rho0.assign(rho_b) u0.project(as_vector([20.0, 0.0])) - state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) + eqns.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) def test_checkpointing(tmpdir): dirname_1 = str(tmpdir)+'/checkpointing_1' dirname_2 = str(tmpdir)+'/checkpointing_2' - state_1, stepper_1, dt = setup_checkpointing(dirname_1) - state_2, stepper_2, dt = setup_checkpointing(dirname_2) + eqns_1, stepper_1, dt = setup_checkpointing(dirname_1) + eqns_2, stepper_2, dt = setup_checkpointing(dirname_2) # ------------------------------------------------------------------------ # # Run for 4 time steps and store values # ------------------------------------------------------------------------ # - initialise_fields(state_1) stepper_1.run(t=0.0, tmax=4*dt) # ------------------------------------------------------------------------ # # Start again, run for 2 time steps, checkpoint and then run for 2 more # ------------------------------------------------------------------------ # - initialise_fields(state_2) stepper_2.run(t=0.0, tmax=2*dt) # Wipe fields, then pickup - state_2.fields('u').project(as_vector([-10.0, 0.0])) - state_2.fields('rho').interpolate(Constant(0.0)) - state_2.fields('theta').interpolate(Constant(0.0)) + eqns_2.fields('u').project(as_vector([-10.0, 0.0])) + eqns_2.fields('rho').interpolate(Constant(0.0)) + eqns_2.fields('theta').interpolate(Constant(0.0)) stepper_2.run(t=2*dt, tmax=4*dt, pickup=True) @@ -124,14 +121,14 @@ def test_checkpointing(tmpdir): # This is the best way to compare fields from different meshes for field_name in ['u', 'rho', 'theta']: with DumbCheckpoint(dirname_1+'/chkpt', mode=FILE_READ) as chkpt: - field_1 = Function(state_1.fields(field_name).function_space(), + field_1 = Function(eqns_1.fields(field_name).function_space(), name=field_name) chkpt.load(field_1) # These are preserved in the comments for when we can use CheckpointFile # mesh = chkpt.load_mesh(name='firedrake_default_extruded') # field_1 = chkpt.load_function(mesh, name=field_name) with DumbCheckpoint(dirname_2+'/chkpt', mode=FILE_READ) as chkpt: - field_2 = Function(state_1.fields(field_name).function_space(), + field_2 = Function(eqns_1.fields(field_name).function_space(), name=field_name) chkpt.load(field_2) # These are preserved in the comments for when we can use CheckpointFile diff --git a/integration-tests/model/test_passive_tracer.py b/integration-tests/model/test_passive_tracer.py index c5ba70a56..eae1e2df3 100644 --- a/integration-tests/model/test_passive_tracer.py +++ b/integration-tests/model/test_passive_tracer.py @@ -16,12 +16,12 @@ def run_tracer(setup): # Get initial conditions from shared config - state = setup.state - mesh = state.mesh - dt = state.dt - output = state.output + domain = setup.domain + mesh = domain.mesh + dt = setup.dt + output = setup.output - x = SpatialCoordinate(state.mesh) + x = SpatialCoordinate(mesh) H = 0.1 parameters = ShallowWaterParameters(H=H) Omega = parameters.Omega @@ -31,40 +31,42 @@ def run_tracer(setup): fexpr = 2*Omega*x[2]/R # Need to create a new state containing parameters - state = State(mesh, dt=dt, output=output, parameters=parameters) # Equations - eqns = LinearShallowWaterEquations(state, setup.family, - setup.degree, fexpr=fexpr) - tracer_eqn = AdvectionEquation(state, state.spaces("DG"), "tracer") + eqns = LinearShallowWaterEquations(domain, parameters, fexpr=fexpr) + tracer_eqn = AdvectionEquation(domain, domain.spaces("DG"), "tracer") + io = IO(domain, eqns, dt=dt, output=output) # Specify initial prognostic fields - u0 = state.fields("u") - D0 = state.fields("D") - tracer0 = state.fields("tracer", D0.function_space()) + u0 = eqns.fields("u") + D0 = eqns.fields("D") + tracer0 = tracer_eqn.fields("tracer", D0.function_space()) tracer_end = Function(D0.function_space()) # Expressions for initial fields corresponding to Williamson 2 test case Dexpr = H - ((R * Omega * umax)*(x[2]*x[2]/(R*R))) / g u0.project(setup.uexpr) D0.interpolate(Dexpr) + Dbar = Function(D0.function_space()).assign(H) tracer0.interpolate(setup.f_init) tracer_end.interpolate(setup.f_end) + eqns.set_reference_profiles([('D', Dbar)]) + # set up transport schemes - transport_schemes = [ForwardEuler(state, "D")] + transport_schemes = [ForwardEuler(domain, io, "D")] # Set up tracer transport - tracer_transport = [(tracer_eqn, SSPRK3(state))] + tracer_transport = [(tracer_eqn, SSPRK3(domain, io))] # build time stepper stepper = SemiImplicitQuasiNewton( - eqns, state, transport_schemes, + eqns, io, transport_schemes, auxiliary_equations_and_schemes=tracer_transport) stepper.run(t=0, tmax=setup.tmax) - error = norm(state.fields("tracer") - tracer_end) / norm(tracer_end) + error = norm(tracer_eqn.fields("tracer") - tracer_end) / norm(tracer_end) return error diff --git a/integration-tests/model/test_prescribed_transport.py b/integration-tests/model/test_prescribed_transport.py index a881c2f18..cb792c362 100644 --- a/integration-tests/model/test_prescribed_transport.py +++ b/integration-tests/model/test_prescribed_transport.py @@ -8,38 +8,38 @@ from firedrake import sin, cos, norm, pi, as_vector -def run(eqn, transport_scheme, state, tmax, f_end, prescribed_u): - timestepper = PrescribedTransport(eqn, transport_scheme, state, +def run(eqn, transport_scheme, io, tmax, f_end, prescribed_u): + timestepper = PrescribedTransport(eqn, transport_scheme, io, prescribed_transporting_velocity=prescribed_u) timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(eqn.fields("f") - f_end) / norm(f_end) def test_prescribed_transport_setup(tmpdir, tracer_setup): - # Make mesh and state using routine from conftest + # Make domain using routine from conftest geometry = "slice" setup = tracer_setup(tmpdir, geometry, degree=1) - state = setup.state - _, z = SpatialCoordinate(state.mesh) + domain = setup.domain + _, z = SpatialCoordinate(domain.mesh) - V = state.spaces("DG", "DG", 1) + V = domain.spaces("DG") # Make equation - eqn = AdvectionEquation(state, V, "f", - ufamily=setup.family, udegree=1) + eqn = AdvectionEquation(domain, V, "f") + io = IO(domain, eqn, dt=setup.dt, output=setup.output) # Initialise fields def u_evaluation(t): return as_vector([2.0*cos(2*pi*t/setup.tmax), sin(2*pi*t/setup.tmax)*sin(pi*z)]) - state.fields("f").interpolate(setup.f_init) - state.fields("u").project(u_evaluation(Constant(0.0))) + eqn.fields("f").interpolate(setup.f_init) + eqn.fields("u").project(u_evaluation(Constant(0.0))) - transport_scheme = SSPRK3(state) + transport_scheme = SSPRK3(domain, io) # Run and check error - error = run(eqn, transport_scheme, state, setup.tmax, + error = run(eqn, transport_scheme, io, setup.tmax, setup.f_init, u_evaluation) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py index c6172432e..f38f19bf4 100644 --- a/integration-tests/model/test_time_discretisation.py +++ b/integration-tests/model/test_time_discretisation.py @@ -3,10 +3,10 @@ import pytest -def run(eqn, transport_scheme, state, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, state) +def run(eqn, transport_scheme, io, tmax, f_end): + timestepper = PrescribedTransport(eqn, transport_scheme, io) timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(eqn.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("scheme", ["ssprk", "implicit_midpoint", @@ -14,23 +14,23 @@ def run(eqn, transport_scheme, state, tmax, f_end): def test_time_discretisation(tmpdir, scheme, tracer_setup): geometry = "sphere" setup = tracer_setup(tmpdir, geometry) - state = setup.state - V = state.spaces("DG", "DG", 1) + domain = setup.domain + V = domain.spaces("DG") - eqn = AdvectionEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) + eqn = AdvectionEquation(domain, V, "f") + io = IO(domain, eqn, dt=setup.dt, output=setup.output) - state.fields("f").interpolate(setup.f_init) - state.fields("u").project(setup.uexpr) + eqn.fields("f").interpolate(setup.f_init) + eqn.fields("u").project(setup.uexpr) if scheme == "ssprk": - transport_scheme = SSPRK3(state) + transport_scheme = SSPRK3(domain, io) elif scheme == "implicit_midpoint": - transport_scheme = ImplicitMidpoint(state) + transport_scheme = ImplicitMidpoint(domain, io) elif scheme == "RK4": - transport_scheme = RK4(state) + transport_scheme = RK4(domain, io) elif scheme == "Heun": - transport_scheme = Heun(state) + transport_scheme = Heun(domain, io) elif scheme == "BDF2": - transport_scheme = BDF2(state) - assert run(eqn, transport_scheme, state, setup.tmax, setup.f_end) < setup.tol + transport_scheme = BDF2(domain, io) + assert run(eqn, transport_scheme, io, setup.tmax, setup.f_end) < setup.tol diff --git a/integration-tests/physics/test_condensation.py b/integration-tests/physics/test_condensation.py index 9fdff1bb2..82a6f207e 100644 --- a/integration-tests/physics/test_condensation.py +++ b/integration-tests/physics/test_condensation.py @@ -24,33 +24,30 @@ def run_cond_evap(dirname, process): # make mesh m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) - x, z = SpatialCoordinate(mesh) - - dt = 2.0 - output = OutputParameters(dirname=dirname+"/cond_evap", - dumpfreq=1, - dumplist=['u']) - parameters = CompressibleParameters() + domain = Domain(mesh, "CG", 1) - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=[Sum('water_vapour', 'cloud_water')]) + x, z = SpatialCoordinate(mesh) # spaces - Vt = state.spaces("theta", degree=1) - Vr = state.spaces("DG", "DG", degree=1) + Vt = domain.spaces("theta", degree=1) + Vr = domain.spaces("DG", "DG", degree=1) # Set up equation -- use compressible to set up these spaces tracers = [WaterVapour(), CloudWater()] - eqn = CompressibleEulerEquations(state, "CG", 1, active_tracers=tracers) + parameters = CompressibleParameters() + eqn = CompressibleEulerEquations(domain, parameters, active_tracers=tracers) + + dt = 2.0 + output = OutputParameters(dirname=dirname+"/cond_evap", + dumpfreq=1, + dumplist=['u']) + io = IO(domain, eqn, dt=dt, output=output, diagnostic_fields=[Sum('water_vapour', 'cloud_water')]) # Declare prognostic fields - rho0 = state.fields("rho") - theta0 = state.fields("theta") - water_v0 = state.fields("water_vapour", Vt) - water_c0 = state.fields("cloud_water", Vt) + rho0 = eqn.fields("rho") + theta0 = eqn.fields("theta") + water_v0 = eqn.fields("water_vapour", Vt) + water_c0 = eqn.fields("cloud_water", Vt) # Set a background state with constant pressure and temperature pressure = Function(Vr).interpolate(Constant(100000.)) @@ -90,29 +87,29 @@ def run_cond_evap(dirname, process): eqn.residual = eqn.residual.label_map(lambda t: t.has_label(time_derivative), map_if_true=identity, map_if_false=drop) - physics_schemes = [(SaturationAdjustment(eqn, parameters), ForwardEuler(state))] + physics_schemes = [(SaturationAdjustment(eqn, parameters=parameters), ForwardEuler(domain, io))] # build time stepper - scheme = ForwardEuler(state) - stepper = SplitPhysicsTimestepper(eqn, scheme, state, + scheme = ForwardEuler(domain, io) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, physics_schemes=physics_schemes) stepper.run(t=0, tmax=dt) - return state, mv_true, mc_true, theta_d_true, mc_init + return eqn, mv_true, mc_true, theta_d_true, mc_init @pytest.mark.parametrize("process", ["evaporation", "condensation"]) def test_cond_evap(tmpdir, process): dirname = str(tmpdir) - state, mv_true, mc_true, theta_d_true, mc_init = run_cond_evap(dirname, process) + eqn, mv_true, mc_true, theta_d_true, mc_init = run_cond_evap(dirname, process) - water_v = state.fields('water_vapour') - water_c = state.fields('cloud_water') - theta_vd = state.fields('theta') + water_v = eqn.fields('water_vapour') + water_c = eqn.fields('cloud_water') + theta_vd = eqn.fields('theta') theta_d = Function(theta_vd.function_space()) - theta_d.interpolate(theta_vd/(1 + water_v * state.parameters.R_v / state.parameters.R_d)) + theta_d.interpolate(theta_vd/(1 + water_v * eqn.parameters.R_v / eqn.parameters.R_d)) # Check that water vapour is approximately equal to saturation amount assert norm(water_v - mv_true) / norm(mv_true) < 0.01, \ diff --git a/integration-tests/physics/test_instant_rain.py b/integration-tests/physics/test_instant_rain.py index 9d2cb0cf2..fc7820cfd 100644 --- a/integration-tests/physics/test_instant_rain.py +++ b/integration-tests/physics/test_instant_rain.py @@ -18,6 +18,7 @@ def run_instant_rain(dirname): L = 10 nx = 10 mesh = PeriodicSquareMesh(nx, nx, L) + domain = Domain(mesh, "BDM", 1) x, y = SpatialCoordinate(mesh) # parameters @@ -26,30 +27,21 @@ def run_instant_rain(dirname): fexpr = Constant(0) dt = 0.1 - output = OutputParameters(dirname=dirname+"/instant_rain", - dumpfreq=1, - dumplist=['vapour', "rain"]) - - parameters = ShallowWaterParameters(H=H, g=g) - - diagnostic_fields = [CourantNumber()] - - state = State(mesh, - dt=dt, - output=output, - diagnostic_fields=diagnostic_fields, - parameters=parameters) - vapour = WaterVapour(name="water_vapour", space='DG') rain = Rain(name="rain", space="DG", transport_eqn=TransportEquationType.no_transport) - VD = FunctionSpace(mesh, "DG", 1) - - eqns = ShallowWaterEquations(state, "BDM", 1, fexpr=fexpr, + parameters = ShallowWaterParameters(H=H, g=g) + eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, active_tracers=[vapour, rain]) - vapour0 = state.fields("water_vapour") + output = OutputParameters(dirname=dirname+"/instant_rain", + dumpfreq=1, + dumplist=['vapour', "rain"]) + diagnostic_fields = [CourantNumber(dt)] + io = IO(domain, eqns, dt=dt, output=output, diagnostic_fields=diagnostic_fields) + + vapour0 = eqns.fields("water_vapour") # set up vapour xc = L/2 @@ -60,6 +52,7 @@ def run_instant_rain(dirname): vapour0.interpolate(vapour_expr) + VD = FunctionSpace(mesh, "DG", 1) initial_vapour = Function(VD).interpolate(vapour_expr) # define saturation function @@ -71,20 +64,20 @@ def run_instant_rain(dirname): rain_true = Function(VD).interpolate(vapour0 - saturation) physics_schemes = [(InstantRain(eqns, saturation, rain_name="rain", - set_tau_to_dt=True), ForwardEuler(state))] + set_tau_to_dt=True), ForwardEuler(domain, io))] - stepper = PrescribedTransport(eqns, RK4(state), state, + stepper = PrescribedTransport(eqns, RK4(domain, io), io, physics_schemes=physics_schemes) stepper.run(t=0, tmax=5*dt) - return state, saturation, initial_vapour, vapour_true, rain_true + return eqns, saturation, initial_vapour, vapour_true, rain_true def test_instant_rain_setup(tmpdir): dirname = str(tmpdir) - state, saturation, initial_vapour, vapour_true, rain_true = run_instant_rain(dirname) - v = state.fields("water_vapour") - r = state.fields("rain") + eqns, saturation, initial_vapour, vapour_true, rain_true = run_instant_rain(dirname) + v = eqns.fields("water_vapour") + r = eqns.fields("rain") # check that the maximum of the vapour field is equal to the saturation assert v.dat.data.max() - saturation.values() < 0.001, "The maximum of the final vapour field should be equal to saturation" diff --git a/integration-tests/physics/test_precipitation.py b/integration-tests/physics/test_precipitation.py index d4d9ca679..ce9c5cc48 100644 --- a/integration-tests/physics/test_precipitation.py +++ b/integration-tests/physics/test_precipitation.py @@ -22,29 +22,23 @@ def setup_fallout(dirname): # make mesh m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) + domain = Domain(mesh, "CG", 1) x = SpatialCoordinate(mesh) + Vrho = domain.spaces("DG1_equispaced") + active_tracers = [Rain(space='DG1_equispaced')] + eqn = ForcedAdvectionEquation(domain, Vrho, "rho", active_tracers=active_tracers) + dt = 0.1 - output = OutputParameters(dirname=dirname+"/fallout", - dumpfreq=10, - dumplist=['rain']) - parameters = CompressibleParameters() + output = OutputParameters(dirname=dirname+"/fallout", dumpfreq=10, dumplist=['rain']) diagnostic_fields = [Precipitation()] - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) + io = IO(domain, eqn, dt=dt, output=output, diagnostic_fields=diagnostic_fields) - Vrho = state.spaces("DG1_equispaced") - active_tracers = [Rain(space='DG1_equispaced')] - eqn = ForcedAdvectionEquation(state, Vrho, "rho", ufamily="CG", udegree=1, - active_tracers=active_tracers) - scheme = ForwardEuler(state) - state.fields("rho").assign(1.) + scheme = ForwardEuler(domain, io) + eqn.fields("rho").assign(1.) - physics_schemes = [(Fallout(eqn, 'rain', state), SSPRK3(state, 'rain'))] - rain0 = state.fields("rain") + physics_schemes = [(Fallout(eqn, 'rain', domain), SSPRK3(domain, io, 'rain'))] + rain0 = eqn.fields("rain") # set up rain xc = L / 2 @@ -56,7 +50,7 @@ def setup_fallout(dirname): rain0.interpolate(rain_expr) # build time stepper - stepper = PrescribedTransport(eqn, scheme, state, + stepper = PrescribedTransport(eqn, scheme, io, physics_schemes=physics_schemes) return stepper, 10.0 diff --git a/integration-tests/transport/test_limiters.py b/integration-tests/transport/test_limiters.py index 1ca6dd1cb..34117717d 100644 --- a/integration-tests/transport/test_limiters.py +++ b/integration-tests/transport/test_limiters.py @@ -55,7 +55,7 @@ def setup_limiters(dirname, space): else: raise NotImplementedError - Vpsi = domain.spaces('CG', 'CG', degree, degree) + Vpsi = domain.spaces('CG', 'CG', degree+1) # set up the equation eqn = AdvectionEquation(domain, V, 'tracer') From 3bec4db61724c2e3e6cce9310fd645e8a124c590 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Tue, 13 Dec 2022 21:28:17 +0000 Subject: [PATCH 06/12] reorganise some of the new objects --- gusto/diagnostics.py | 9 +--- gusto/domain.py | 15 +++++- gusto/forcing.py | 4 +- gusto/function_spaces.py | 8 +-- gusto/io.py | 11 +--- gusto/linear_solvers.py | 15 +++--- gusto/time_discretisation.py | 30 +++++------ gusto/timeloop.py | 24 +++++---- .../balance/test_compressible_balance.py | 12 ++--- .../balance/test_saturated_balance.py | 20 +++---- .../balance/test_unsaturated_balance.py | 20 +++---- integration-tests/conftest.py | 14 ++--- integration-tests/diffusion/test_diffusion.py | 8 +-- .../equations/test_advection_diffusion.py | 26 ++++----- .../equations/test_dry_compressible.py | 47 ++++++++-------- .../equations/test_forced_advection.py | 30 +++++------ .../equations/test_incompressible.py | 45 ++++++++-------- .../equations/test_moist_compressible.py | 46 ++++++++-------- .../equations/test_sw_linear_triangle.py | 27 +++++----- .../equations/test_sw_triangle.py | 53 +++++++++---------- integration-tests/model/test_checkpointing.py | 12 ++--- .../model/test_passive_tracer.py | 7 ++- .../model/test_prescribed_transport.py | 2 +- .../model/test_time_discretisation.py | 12 ++--- .../physics/test_condensation.py | 13 +++-- .../physics/test_instant_rain.py | 12 ++--- .../physics/test_precipitation.py | 10 ++-- .../transport/test_dg_transport.py | 8 +-- .../transport/test_embedded_dg_advection.py | 4 +- integration-tests/transport/test_limiters.py | 12 ++--- .../transport/test_recovered_transport.py | 4 +- .../transport/test_subcycling.py | 4 +- .../transport/test_supg_transport.py | 12 ++--- .../transport/test_vector_recovered_space.py | 4 +- 34 files changed, 284 insertions(+), 296 deletions(-) diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index fc0d62e5b..df7b337dd 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -188,14 +188,6 @@ class CourantNumber(DiagnosticField): """Dimensionless Courant number diagnostic field.""" name = "CourantNumber" - def __init__(self, dt): - """ - Args: - dt (float): time increment over a model time step. - """ - super().__init__() - self.dt = dt - def setup(self, eqn): """ Sets up the :class:`Function` for the diagnostic field. @@ -205,6 +197,7 @@ def setup(self, eqn): """ if not self._initialised: super(CourantNumber, self).setup(eqn) + self.dt = eqn.domain.dt # set up area computation V = eqn.domain.spaces("DG0") test = TestFunction(V) diff --git a/gusto/domain.py b/gusto/domain.py index 346ef0d35..6a967d125 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -1,6 +1,7 @@ """ The Domain object that is provided in this module contains the model's mesh and -the set of compatible function spaces defined upon it. +the set of compatible function spaces defined upon it. It also contains the +model's time interval. """ from gusto.function_spaces import Spaces, check_degree_args @@ -19,11 +20,14 @@ class Domain(object): vertical degrees of the elements. Alternatively, if these degrees should be the same then this can be specified through the "degree" argument. """ - def __init__(self, mesh, family, degree=None, + def __init__(self, mesh, dt, family, degree=None, horizontal_degree=None, vertical_degree=None): """ Args: mesh (:class:`Mesh`): the model's mesh. + dt (:class:`Constant`): the time taken to perform a single model + step. If a float or int is passed, it will be cast to a + :class:`Constant`. family (str): the finite element space family used for the velocity field. This determines the other finite element spaces used via the de Rham complex. @@ -39,6 +43,13 @@ def __init__(self, mesh, family, degree=None, both "degree" and "horizontal_degree"). """ + if type(dt) is Constant: + self.dt = dt + elif type(dt) in (float, int): + self.dt = Constant(dt) + else: + raise TypeError(f'dt must be a Constant, float or int, not {type(dt)}') + check_degree_args('Domain', mesh, degree, horizontal_degree, vertical_degree) # Get degrees diff --git a/gusto/forcing.py b/gusto/forcing.py index ef97d2142..f09a77efc 100644 --- a/gusto/forcing.py +++ b/gusto/forcing.py @@ -21,7 +21,7 @@ class Forcing(object): discretisation. """ - def __init__(self, equation, io, alpha): + def __init__(self, equation, alpha): """ Args: equation (:class:`PrognosticEquationSet`): the prognostic equations @@ -33,7 +33,7 @@ def __init__(self, equation, io, alpha): self.field_name = equation.field_name implicit_terms = ["incompressibility", "sponge"] - dt = io.dt + dt = equation.domain.dt W = equation.function_space self.x0 = Function(W) diff --git a/gusto/function_spaces.py b/gusto/function_spaces.py index 956c71cc5..5ff3eab83 100644 --- a/gusto/function_spaces.py +++ b/gusto/function_spaces.py @@ -294,14 +294,16 @@ def check_degree_args(name, mesh, degree, horizontal_degree, vertical_degree): of a space. """ + extruded_mesh = hasattr(mesh, "_base_mesh") + # Checks on degree arguments if degree is None and horizontal_degree is None: raise ValueError(f'Either "degree" or "horizontal_degree" must be passed to {name}') - if mesh.extruded and degree is None and vertical_degree is None: + if extruded_mesh and degree is None and vertical_degree is None: raise ValueError(f'For extruded meshes, either degree or "vertical_degree" must be passed to {name}') if degree is not None and horizontal_degree is not None: raise ValueError(f'Cannot pass both "degree" and "horizontal_degree" to {name}') - if mesh.extruded and degree is not None and vertical_degree is not None: + if extruded_mesh and degree is not None and vertical_degree is not None: raise ValueError(f'Cannot pass both "degree" and "vertical_degree" to {name}') - if not mesh.extruded and vertical_degree is not None: + if not extruded_mesh and vertical_degree is not None: raise ValueError(f'Cannot pass "vertical_degree" to {name} if mesh is not extruded') diff --git a/gusto/io.py b/gusto/io.py index 7cb91d095..5c0d2c946 100644 --- a/gusto/io.py +++ b/gusto/io.py @@ -160,7 +160,7 @@ def dump(self, equation, t): class IO(object): """Controls the model's input, output and diagnostics.""" - def __init__(self, domain, equation, dt, + def __init__(self, domain, equation, output=None, parameters=None, diagnostics=None, @@ -170,9 +170,6 @@ def __init__(self, domain, equation, dt, domain (:class:`Domain`): the model's domain object, containing the mesh and the compatible function spaces. equation (:class:`PrognosticEquation`): the prognostic equation. - dt (:class:`Constant`): the time taken to perform a single model - step. If a float or int is passed, it will be cast to a - :class:`Constant`. output (:class:`OutputParameters`, optional): holds and describes the options for outputting. Defaults to None. diagnostics (:class:`Diagnostics`, optional): object holding and @@ -227,12 +224,6 @@ def __init__(self, domain, equation, dt, # Constant to hold current time self.t = Constant(0.0) - if type(dt) is Constant: - self.dt = dt - elif type(dt) in (float, int): - self.dt = Constant(dt) - else: - raise TypeError(f'dt must be a Constant, float or int, not {type(dt)}') def setup_diagnostics(self): """Concatenates the various types of diagnostic field.""" diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index c19d1ea5c..e4704d1dc 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -28,7 +28,7 @@ class TimesteppingSolver(object, metaclass=ABCMeta): """Base class for timestepping linear solvers for Gusto.""" - def __init__(self, equations, io, alpha=0.5, solver_parameters=None, + def __init__(self, equations, alpha=0.5, solver_parameters=None, overwrite_solver_parameters=False): """ Args: @@ -44,7 +44,7 @@ def __init__(self, equations, io, alpha=0.5, solver_parameters=None, passed in. Defaults to False. """ self.equations = equations - self.dt = io.dt + self.dt = equations.domain.dt self.alpha = alpha if solver_parameters is not None: @@ -121,7 +121,7 @@ class CompressibleSolver(TimesteppingSolver): 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}}} - def __init__(self, equations, io, alpha=0.5, + def __init__(self, equations, alpha=0.5, quadrature_degree=None, solver_parameters=None, overwrite_solver_parameters=False, moisture=None): """ @@ -142,6 +142,7 @@ def __init__(self, equations, io, alpha=0.5, moisture (list, optional): list of names of moisture fields. Defaults to None. """ + self.equations = equations self.moisture = moisture if quadrature_degree is not None: @@ -162,7 +163,7 @@ def __init__(self, equations, io, alpha=0.5, # Turn monitor on for the trace system self.solver_parameters["condensed_field"]["ksp_monitor_true_residual"] = None - super().__init__(equations, io, alpha, solver_parameters, + super().__init__(equations, alpha, solver_parameters, overwrite_solver_parameters) @timed_function("Gusto:SolverSetup") @@ -422,7 +423,7 @@ class IncompressibleSolver(TimesteppingSolver): @timed_function("Gusto:SolverSetup") def _setup_solver(self): - equation = self.equation # just cutting down line length a bit + equation = self.equations # just cutting down line length a bit dt = self.dt beta_ = dt*self.alpha Vu = equation.domain.spaces("HDiv") @@ -550,7 +551,7 @@ class LinearTimesteppingSolver(object): 'sub_pc_type': 'ilu'}} } - def __init__(self, equation, io, alpha): + def __init__(self, equation, alpha): """ Args: equation (:class:`PrognosticEquation`): the model's equation object. @@ -562,7 +563,7 @@ def __init__(self, equation, io, alpha): lambda t: Term(t.get(linearisation).form, t.labels), drop) - dt = io.dt + dt = equation.domain.dt W = equation.function_space beta = dt*alpha diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 1e23e914e..1e818d4dc 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -73,13 +73,12 @@ def new_apply(self, x_out, x_in): class TimeDiscretisation(object, metaclass=ABCMeta): """Base class for time discretisation schemes.""" - def __init__(self, domain, io, field_name=None, solver_parameters=None, + 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. - io (:class:`IO`): the model's object for controlling input/output. field_name (str, optional): name of the field to be evolved. Defaults to None. solver_parameters (dict, optional): dictionary of parameters to @@ -95,7 +94,7 @@ def __init__(self, domain, io, field_name=None, solver_parameters=None, self.field_name = field_name self.equation = None - self.dt = io.dt + self.dt = domain.dt self.limiter = limiter @@ -438,13 +437,12 @@ def apply(self, x_out, x_in): class ExplicitTimeDiscretisation(TimeDiscretisation): """Base class for explicit time discretisations.""" - def __init__(self, domain, io, field_name=None, subcycles=None, + def __init__(self, domain, field_name=None, subcycles=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. - io (:class:`IO`): the model's object for controlling input/output. field_name (str, optional): name of the field to be evolved. Defaults to None. subcycles (int, optional): the number of sub-steps to perform. @@ -458,7 +456,7 @@ def __init__(self, domain, io, field_name=None, subcycles=None, to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. """ - super().__init__(domain, io, field_name, + super().__init__(domain, field_name, solver_parameters=solver_parameters, limiter=limiter, options=options) @@ -780,13 +778,12 @@ class BackwardEuler(TimeDiscretisation): The backward Euler method for operator F is the most simple implicit scheme: y^(n+1) = y^n + dt*F[y^(n+1)]. """ - def __init__(self, domain, io, field_name=None, solver_parameters=None, + 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. - io (:class:`IO`): the model's object for controlling input/output. field_name (str, optional): name of the field to be evolved. Defaults to None. subcycles (int, optional): the number of sub-steps to perform. @@ -805,7 +802,7 @@ def __init__(self, domain, io, field_name=None, solver_parameters=None, """ if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") - super().__init__(domain=domain, io=io, field_name=field_name, + super().__init__(domain=domain, field_name=field_name, solver_parameters=solver_parameters, limiter=limiter, options=options) @@ -852,13 +849,12 @@ class ThetaMethod(TimeDiscretisation): for off-centring parameter theta. """ - def __init__(self, domain, io, field_name=None, theta=None, + def __init__(self, domain, field_name=None, theta=None, solver_parameters=None, options=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the mesh and the compatible function spaces. - io (:class:`IO`): the model's object for controlling input/output. field_name (str, optional): name of the field to be evolved. Defaults to None. theta (float, optional): the off-centring parameter. theta = 1 @@ -886,7 +882,7 @@ def __init__(self, domain, io, field_name=None, theta=None, 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} - super().__init__(domain, io, field_name, + super().__init__(domain, field_name, solver_parameters=solver_parameters, options=options) @@ -936,13 +932,12 @@ class ImplicitMidpoint(ThetaMethod): It is equivalent to the "theta" method with theta = 1/2. """ - def __init__(self, domain, io, field_name=None, solver_parameters=None, + 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. - io (:class:`IO`): the model's object for controlling input/output. field_name (str, optional): name of the field to be evolved. Defaults to None. solver_parameters (dict, optional): dictionary of parameters to @@ -952,7 +947,7 @@ def __init__(self, domain, io, field_name=None, solver_parameters=None, to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. """ - super().__init__(domain, io, field_name, theta=0.5, + super().__init__(domain, field_name, theta=0.5, solver_parameters=solver_parameters, options=options) @@ -960,13 +955,12 @@ def __init__(self, domain, io, field_name=None, solver_parameters=None, class MultilevelTimeDiscretisation(TimeDiscretisation): """Base class for multi-level timesteppers""" - def __init__(self, domain, io, field_name=None, solver_parameters=None, + 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. - io (:class:`IO`): the model's object for controlling input/output. field_name (str, optional): name of the field to be evolved. Defaults to None. solver_parameters (dict, optional): dictionary of parameters to @@ -980,7 +974,7 @@ def __init__(self, domain, io, field_name=None, solver_parameters=None, """ if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") - super().__init__(domain=domain, io=io, field_name=field_name, + super().__init__(domain=domain, field_name=field_name, solver_parameters=solver_parameters, limiter=limiter, options=options) self.initial_timesteps = 0 diff --git a/gusto/timeloop.py b/gusto/timeloop.py index e72924061..b0409ae49 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -1,7 +1,7 @@ """Classes for controlling the timestepping loop.""" from abc import ABCMeta, abstractmethod, abstractproperty -from firedrake import Function, Projector +from firedrake import Function, Projector, Constant from pyop2.profiling import timed_stage from gusto.configuration import logger from gusto.forcing import Forcing @@ -28,6 +28,8 @@ def __init__(self, equation, io): self.equation = equation self.io = io + self.dt = self.equation.domain.dt + self.t = Constant(0.0) self.setup_fields() self.setup_scheme() @@ -71,12 +73,12 @@ def run(self, t, tmax, pickup=False): with timed_stage("Dump output"): io.setup_dump(t, tmax, pickup) - io.t.assign(t) + self.t.assign(t) self.x.initialise(self.equation) - while float(io.t) < tmax - 0.5*float(io.dt): - logger.info(f'at start of timestep, t={float(io.t)}, dt={float(io.dt)}') + while float(self.t) < tmax - 0.5*float(self.dt): + logger.info(f'at start of timestep, t={float(self.t)}, dt={float(self.dt)}') self.x.update() @@ -85,15 +87,15 @@ def run(self, t, tmax, pickup=False): for field in self.x.np1: self.equation.fields(field.name()).assign(field) - io.t.assign(io.t + io.dt) + self.t.assign(self.t + self.dt) with timed_stage("Dump output"): - io.dump(float(io.t)) + io.dump(float(self.t)) if io.output.checkpoint: io.chkpt.close() - logger.info(f'TIMELOOP complete. t={float(io.t)}, tmax={tmax}') + logger.info(f'TIMELOOP complete. t={float(self.t)}, tmax={tmax}') class Timestepper(BaseTimestepper): @@ -157,6 +159,8 @@ def __init__(self, equation, scheme, io, physics_schemes=None): self.equation = equation self.scheme = scheme self.io = io + self.dt = self.equation.domain.dt + self.t = Constant(0.0) self.setup_fields() self.setup_scheme() @@ -295,10 +299,10 @@ def __init__(self, equation_set, io, transport_schemes, self.xrhs = Function(W) self.dy = Function(W) if linear_solver is None: - self.linear_solver = LinearTimesteppingSolver(equation_set, io, self.alpha) + self.linear_solver = LinearTimesteppingSolver(equation_set, self.alpha) else: self.linear_solver = linear_solver - self.forcing = Forcing(equation_set, io, self.alpha) + self.forcing = Forcing(equation_set, self.alpha) self.bcs = equation_set.bcs def _apply_bcs(self): @@ -445,7 +449,7 @@ def __init__(self, equation, scheme, io, physics_schemes=None, if prescribed_transporting_velocity is not None: self.velocity_projection = Projector( - prescribed_transporting_velocity(self.io.t), + prescribed_transporting_velocity(self.t), self.equation.fields('u')) else: self.velocity_projection = None diff --git a/integration-tests/balance/test_compressible_balance.py b/integration-tests/balance/test_compressible_balance.py index 3e62638a5..41db83aa3 100644 --- a/integration-tests/balance/test_compressible_balance.py +++ b/integration-tests/balance/test_compressible_balance.py @@ -24,12 +24,12 @@ def setup_balance(dirname): m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) - domain = Domain(mesh, "CG", 1) + domain = Domain(mesh, dt, "CG", 1) output = OutputParameters(dirname=dirname+'/dry_balance', dumpfreq=10, dumplist=['u']) parameters = CompressibleParameters() eqns = CompressibleEulerEquations(domain, parameters) - io = IO(domain, eqns, dt=dt, output=output) + io = IO(domain, eqns, output=output) # Initial conditions rho0 = eqns.fields("rho") @@ -46,12 +46,12 @@ def setup_balance(dirname): ('theta', theta0)]) # Set up transport schemes - transported_fields = [ImplicitMidpoint(domain, io, "u"), - SSPRK3(domain, io, "rho"), - SSPRK3(domain, io, "theta", options=EmbeddedDGOptions())] + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=EmbeddedDGOptions())] # Set up linear solver - linear_solver = CompressibleSolver(eqns, io) + linear_solver = CompressibleSolver(eqns) # build time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, diff --git a/integration-tests/balance/test_saturated_balance.py b/integration-tests/balance/test_saturated_balance.py index e40e02ef7..6c9b1f7a0 100644 --- a/integration-tests/balance/test_saturated_balance.py +++ b/integration-tests/balance/test_saturated_balance.py @@ -28,7 +28,7 @@ def setup_saturated(dirname, recovered): m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) - domain = Domain(mesh, "CG", degree) + domain = Domain(mesh, dt, "CG", degree) tracers = [WaterVapour(), CloudWater()] @@ -42,7 +42,7 @@ def setup_saturated(dirname, recovered): output = OutputParameters(dirname=dirname+'/saturated_balance', dumpfreq=1, dumplist=['u']) diagnostic_fields = [Theta_e()] - io = IO(domain, eqns, dt=dt, output=output, diagnostic_fields=diagnostic_fields) + io = IO(domain, eqns, output=output, diagnostic_fields=diagnostic_fields) # Initial conditions rho0 = eqns.fields("rho") @@ -92,20 +92,20 @@ def setup_saturated(dirname, recovered): wv_opts = EmbeddedDGOptions() wc_opts = EmbeddedDGOptions() - transported_fields = [SSPRK3(domain, io, 'rho', options=rho_opts), - SSPRK3(domain, io, 'theta', options=theta_opts), - SSPRK3(domain, io, 'water_vapour', options=wv_opts), - SSPRK3(domain, io, 'cloud_water', options=wc_opts)] + transported_fields = [SSPRK3(domain, 'rho', options=rho_opts), + SSPRK3(domain, 'theta', options=theta_opts), + SSPRK3(domain, 'water_vapour', options=wv_opts), + SSPRK3(domain, 'cloud_water', options=wc_opts)] if recovered: - transported_fields.append(SSPRK3(domain, io, 'u', options=u_opts)) + transported_fields.append(SSPRK3(domain, 'u', options=u_opts)) else: - transported_fields.append(ImplicitMidpoint(domain, io, 'u')) + transported_fields.append(ImplicitMidpoint(domain, 'u')) - linear_solver = CompressibleSolver(eqns, io, moisture=moisture) + linear_solver = CompressibleSolver(eqns, moisture=moisture) # add physics - physics_schemes = [(SaturationAdjustment(eqns), ForwardEuler(domain, io))] + physics_schemes = [(SaturationAdjustment(eqns), ForwardEuler(domain))] # build time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, diff --git a/integration-tests/balance/test_unsaturated_balance.py b/integration-tests/balance/test_unsaturated_balance.py index fa3715067..b9c59c209 100644 --- a/integration-tests/balance/test_unsaturated_balance.py +++ b/integration-tests/balance/test_unsaturated_balance.py @@ -28,7 +28,7 @@ def setup_unsaturated(dirname, recovered): m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) - domain = Domain(mesh, "CG", degree) + domain = Domain(mesh, dt, "CG", degree) tracers = [WaterVapour(), CloudWater()] @@ -42,7 +42,7 @@ def setup_unsaturated(dirname, recovered): output = OutputParameters(dirname=dirname+'/unsaturated_balance', dumpfreq=1) diagnostic_fields = [Theta_d(), RelativeHumidity()] - io = IO(domain, eqns, dt=dt, output=output, diagnostic_fields=diagnostic_fields) + io = IO(domain, eqns, output=output, diagnostic_fields=diagnostic_fields) # Initial conditions rho0 = eqns.fields("rho") @@ -82,19 +82,19 @@ def setup_unsaturated(dirname, recovered): rho_opts = None theta_opts = EmbeddedDGOptions() - transported_fields = [SSPRK3(domain, io, "rho", options=rho_opts), - SSPRK3(domain, io, "theta", options=theta_opts), - SSPRK3(domain, io, "water_vapour", options=theta_opts), - SSPRK3(domain, io, "cloud_water", options=theta_opts)] + transported_fields = [SSPRK3(domain, "rho", options=rho_opts), + SSPRK3(domain, "theta", options=theta_opts), + SSPRK3(domain, "water_vapour", options=theta_opts), + SSPRK3(domain, "cloud_water", options=theta_opts)] if recovered: - transported_fields.append(SSPRK3(domain, io, "u", options=u_opts)) + transported_fields.append(SSPRK3(domain, "u", options=u_opts)) else: - transported_fields.append(ImplicitMidpoint(domain, io, "u")) + transported_fields.append(ImplicitMidpoint(domain, "u")) - linear_solver = CompressibleSolver(eqns, io, moisture=moisture) + linear_solver = CompressibleSolver(eqns, moisture=moisture) # Set up physics - physics_schemes = [(SaturationAdjustment(eqns), ForwardEuler(domain, io))] + physics_schemes = [(SaturationAdjustment(eqns), ForwardEuler(domain))] # build time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py index 64ebc13d2..d31f16939 100644 --- a/integration-tests/conftest.py +++ b/integration-tests/conftest.py @@ -9,7 +9,7 @@ from collections import namedtuple import pytest -opts = ('domain', 'dt', 'tmax', 'output', 'f_init', 'f_end', 'degree', +opts = ('domain', 'tmax', 'output', 'f_init', 'f_end', 'degree', 'uexpr', 'umax', 'radius', 'tol') TracerSetup = namedtuple('TracerSetup', opts) TracerSetup.__new__.__defaults__ = (None,)*len(opts) @@ -29,7 +29,7 @@ def tracer_sphere(tmpdir, degree): dt = pi/3. * 0.02 output = OutputParameters(dirname=str(tmpdir), dumpfreq=15) - domain = Domain(mesh, family="BDM", degree=degree) + domain = Domain(mesh, dt, family="BDM", degree=degree) umax = 1.0 uexpr = as_vector([- umax * x[1] / radius, umax * x[0] / radius, 0.0]) @@ -40,7 +40,7 @@ def tracer_sphere(tmpdir, degree): tol = 0.05 - return TracerSetup(domain, dt, tmax, output, f_init, f_end, degree, + return TracerSetup(domain, tmax, output, f_init, f_end, degree, uexpr, umax, radius, tol) @@ -56,7 +56,7 @@ def tracer_slice(tmpdir, degree): dt = 0.01 tmax = 0.75 output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) - domain = Domain(mesh, family="CG", degree=degree) + domain = Domain(mesh, dt, family="CG", degree=degree) uexpr = as_vector([2.0, 0.0]) @@ -73,7 +73,7 @@ def tracer_slice(tmpdir, degree): tol = 0.12 - return TracerSetup(domain, dt, tmax, output, f_init, f_end, degree, uexpr, tol=tol) + return TracerSetup(domain, tmax, output, f_init, f_end, degree, uexpr, tol=tol) def tracer_blob_slice(tmpdir, degree): @@ -83,13 +83,13 @@ def tracer_blob_slice(tmpdir, degree): mesh = ExtrudedMesh(m, layers=10, layer_height=1.) output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) - domain = Domain(mesh, family="CG", degree=degree) + domain = Domain(mesh, dt, family="CG", degree=degree) tmax = 1. x = SpatialCoordinate(mesh) f_init = exp(-((x[0]-0.5*L)**2 + (x[1]-0.5*L)**2)) - return TracerSetup(domain, dt, tmax, output, f_init, degree=degree) + return TracerSetup(domain, tmax, output, f_init, degree=degree) @pytest.fixture() diff --git a/integration-tests/diffusion/test_diffusion.py b/integration-tests/diffusion/test_diffusion.py index 9560ac5dc..680330657 100644 --- a/integration-tests/diffusion/test_diffusion.py +++ b/integration-tests/diffusion/test_diffusion.py @@ -36,9 +36,9 @@ def test_scalar_diffusion(tmpdir, DG, tracer_setup): diffusion_params = DiffusionParameters(kappa=kappa, mu=mu) eqn = DiffusionEquation(domain, V, "f", diffusion_parameters=diffusion_params) - io = IO(domain, eqn, dt=setup.dt, output=setup.output) + io = IO(domain, eqn, output=setup.output) - diffusion_scheme = BackwardEuler(domain, io) + diffusion_scheme = BackwardEuler(domain) eqn.fields("f").interpolate(f_init) f_end = run(eqn, diffusion_scheme, io, tmax) @@ -69,14 +69,14 @@ def test_vector_diffusion(tmpdir, DG, tracer_setup): diffusion_params = DiffusionParameters(kappa=kappa, mu=mu) eqn = DiffusionEquation(domain, V, "f", diffusion_parameters=diffusion_params) - io = IO(domain, eqn, dt=setup.dt, output=setup.output) + io = IO(domain, eqn, output=setup.output) if DG: eqn.fields("f").interpolate(f_init) else: eqn.fields("f").project(f_init) - diffusion_scheme = BackwardEuler(domain, io) + diffusion_scheme = BackwardEuler(domain) f_end = run(eqn, diffusion_scheme, io, tmax) assert errornorm(f_end_expr, f_end) < tol diff --git a/integration-tests/equations/test_advection_diffusion.py b/integration-tests/equations/test_advection_diffusion.py index ad4201ca2..b1fb096ba 100644 --- a/integration-tests/equations/test_advection_diffusion.py +++ b/integration-tests/equations/test_advection_diffusion.py @@ -10,21 +10,23 @@ def run_advection_diffusion(tmpdir): - # Mesh, state and equation - L = 10 - mesh = PeriodicIntervalMesh(20, L) dt = 0.02 tmax = 1.0 + L = 10 + mesh = PeriodicIntervalMesh(20, L) + domain = Domain(mesh, dt, "CG", 1) diffusion_params = DiffusionParameters(kappa=0.75, mu=5) - output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) - state = State(mesh, dt=dt, output=output) - V = state.spaces("DG", "DG", 1) + V = domain.spaces("DG", "DG", 1) Vu = VectorFunctionSpace(mesh, "CG", 1) - equation = AdvectionDiffusionEquation(state, V, "f", Vu=Vu, + equation = AdvectionDiffusionEquation(domain, V, "f", Vu=Vu, diffusion_parameters=diffusion_params) + output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) + io = IO(domain, equation, output=output) + + # Initial conditions x = SpatialCoordinate(mesh) xc_init = 0.25*L @@ -45,15 +47,15 @@ def run_advection_diffusion(tmpdir): f_init_expr = f_init*exp(-(x_init / f_width_init)**2) f_end_expr = f_end*exp(-(x_end / f_width_end)**2) - state.fields('f').interpolate(f_init_expr) - state.fields('u').interpolate(as_vector([Constant(umax)])) - f_end = state.fields('f_end', V).interpolate(f_end_expr) + equation.fields('f').interpolate(f_init_expr) + equation.fields('u').interpolate(as_vector([Constant(umax)])) + f_end = equation.fields('f_end', V).interpolate(f_end_expr) # Time stepper - timestepper = PrescribedTransport(equation, SSPRK3(state), state) + timestepper = PrescribedTransport(equation, SSPRK3(domain), io) timestepper.run(0, tmax=tmax) - error = norm(state.fields('f') - f_end) / norm(f_end) + error = norm(equation.fields('f') - f_end) / norm(f_end) return error diff --git a/integration-tests/equations/test_dry_compressible.py b/integration-tests/equations/test_dry_compressible.py index 272e7d796..cdeb221ba 100644 --- a/integration-tests/equations/test_dry_compressible.py +++ b/integration-tests/equations/test_dry_compressible.py @@ -20,23 +20,20 @@ def run_dry_compressible(tmpdir): Lz = 1000.0 m = PeriodicIntervalMesh(ncols, Lx) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=Lz/nlayers) + domain = Domain(mesh, dt, "CG", 1) - output = OutputParameters(dirname=tmpdir+"/dry_compressible", - dumpfreq=2, chkptfreq=2) parameters = CompressibleParameters() R_d = parameters.R_d g = parameters.g + eqn = CompressibleEulerEquations(domain, parameters) - state = State(mesh, - dt=dt, - output=output, - parameters=parameters) - - eqn = CompressibleEulerEquations(state, "CG", 1) + output = OutputParameters(dirname=tmpdir+"/dry_compressible", + dumpfreq=2, chkptfreq=2) + io = IO(domain, eqn, output=output) # Initial conditions - rho0 = state.fields("rho") - theta0 = state.fields("theta") + rho0 = eqn.fields("rho") + theta0 = eqn.fields("theta") # Approximate hydrostatic balance x, z = SpatialCoordinate(mesh) @@ -46,8 +43,7 @@ def run_dry_compressible(tmpdir): theta0.interpolate(tde.theta(parameters, T, p)) rho0.interpolate(p / (R_d * T)) - state.set_reference_profiles([('rho', rho0), - ('theta', theta0)]) + eqn.set_reference_profiles([('rho', rho0), ('theta', theta0)]) # Add perturbation r = sqrt((x-Lx/2)**2 + (z-Lz/2)**2) @@ -55,41 +51,42 @@ def run_dry_compressible(tmpdir): theta0.interpolate(theta0 + theta_pert) # Set up transport schemes - transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "rho"), - SSPRK3(state, "theta")] + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta")] # Set up linear solver for the timestepping scheme - linear_solver = CompressibleSolver(state, eqn) + linear_solver = CompressibleSolver(eqn) # build time stepper - stepper = SemiImplicitQuasiNewton(eqn, state, transported_fields, + stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, linear_solver=linear_solver) # Run stepper.run(t=0, tmax=tmax) - # State for checking checkpoints + # IO for checking checkpoints checkpoint_name = 'dry_compressible_chkpt' new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') + check_eqn = CompressibleEulerEquations(domain, parameters) check_output = OutputParameters(dirname=tmpdir+"/dry_compressible", checkpoint_pickup_filename=new_path) - check_state = State(mesh, dt=dt, output=check_output, parameters=parameters) - check_eqn = CompressibleEulerEquations(check_state, "CG", 1) - check_stepper = SemiImplicitQuasiNewton(check_eqn, check_state, []) + check_io = IO(domain, check_eqn, output=check_output) + check_eqn.set_reference_profiles([]) + check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, []) check_stepper.run(t=0, tmax=0, pickup=True) - return state, check_state + return eqn, check_eqn def test_dry_compressible(tmpdir): dirname = str(tmpdir) - state, check_state = run_dry_compressible(dirname) + eqn, check_eqn = run_dry_compressible(dirname) for variable in ['u', 'rho', 'theta']: - new_variable = state.fields(variable) - check_variable = check_state.fields(variable) + new_variable = eqn.fields(variable) + check_variable = check_eqn.fields(variable) error = norm(new_variable - check_variable) / norm(check_variable) # Slack values chosen to be robust to different platforms diff --git a/integration-tests/equations/test_forced_advection.py b/integration-tests/equations/test_forced_advection.py index 9bada6329..420a08726 100644 --- a/integration-tests/equations/test_forced_advection.py +++ b/integration-tests/equations/test_forced_advection.py @@ -15,7 +15,6 @@ def run_forced_advection(tmpdir): - # mesh, state and equation Lx = 100 delta_x = 2.0 nx = int(Lx/delta_x) @@ -23,15 +22,7 @@ def run_forced_advection(tmpdir): x = SpatialCoordinate(mesh)[0] dt = 0.2 - output = OutputParameters(dirname=str(tmpdir), dumpfreq=1) - diagnostic_fields = [CourantNumber()] - - state = State(mesh, - dt=dt, - output=output, - parameters=None, - diagnostics=None, - diagnostic_fields=diagnostic_fields) + domain = Domain(mesh, dt, "CG", 1) VD = FunctionSpace(mesh, "DG", 1) Vu = VectorFunctionSpace(mesh, "CG", 1) @@ -54,17 +45,22 @@ def run_forced_advection(tmpdir): rain = Rain(space='tracer', transport_eqn=TransportEquationType.no_transport) - meqn = ForcedAdvectionEquation(state, VD, field_name="water_vapour", Vu=Vu, + meqn = ForcedAdvectionEquation(domain, VD, field_name="water_vapour", Vu=Vu, active_tracers=[rain]) physics_schemes = [(InstantRain(meqn, msat, rain_name="rain", - set_tau_to_dt=True), ForwardEuler(state))] + set_tau_to_dt=True), ForwardEuler(domain))] + + output = OutputParameters(dirname=str(tmpdir), dumpfreq=1) + diagnostic_fields = [CourantNumber()] + io = IO(domain, meqn, output=output, diagnostic_fields=diagnostic_fields) + - state.fields("u").project(as_vector([u_max])) - qv = state.fields("water_vapour") + meqn.fields("u").project(as_vector([u_max])) + qv = meqn.fields("water_vapour") qv.project(mexpr) # exact rainfall profile (analytically) - r_exact = state.fields("r_exact", VD) + r_exact = meqn.fields("r_exact", VD) lim1 = Lx/(2*pi) * acos((C0 + K0 - Csat)/Ksat) lim2 = Lx/2 coord = (Ksat*cos(2*pi*x/Lx) + Csat - C0)/K0 @@ -73,12 +69,12 @@ def run_forced_advection(tmpdir): r_exact.interpolate(r_expr) # build time stepper - stepper = PrescribedTransport(meqn, RK4(state), state, + stepper = PrescribedTransport(meqn, RK4(domain), io, physics_schemes=physics_schemes) stepper.run(0, tmax=tmax) - error = errornorm(r_exact, state.fields("rain")) + error = errornorm(r_exact, meqn.fields("rain")) return error diff --git a/integration-tests/equations/test_incompressible.py b/integration-tests/equations/test_incompressible.py index 0ae835a72..c1be5e5c5 100644 --- a/integration-tests/equations/test_incompressible.py +++ b/integration-tests/equations/test_incompressible.py @@ -19,21 +19,18 @@ def run_incompressible(tmpdir): Lz = 1000.0 m = PeriodicIntervalMesh(ncols, Lx) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=Lz/nlayers) + domain = Domain(mesh, dt, "CG", 1) - output = OutputParameters(dirname=tmpdir+"/incompressible", - dumpfreq=2, chkptfreq=2) parameters = CompressibleParameters() + eqn = IncompressibleBoussinesqEquations(domain, parameters) - state = State(mesh, - dt=dt, - output=output, - parameters=parameters) - - eqns = IncompressibleBoussinesqEquations(state, "CG", 1) + output = OutputParameters(dirname=tmpdir+"/incompressible", + dumpfreq=2, chkptfreq=2) + io = IO(domain, eqn, output=output) # Initial conditions - p0 = state.fields("p") - b0 = state.fields("b") + p0 = eqn.fields("p") + b0 = eqn.fields("b") # z.grad(bref) = N**2 x, z = SpatialCoordinate(mesh) @@ -41,9 +38,8 @@ def run_incompressible(tmpdir): bref = z*(N**2) b_b = Function(b0.function_space()).interpolate(bref) - incompressible_hydrostatic_balance(state, b_b, p0) - state.initialise([('p', p0), - ('b', b_b)]) + incompressible_hydrostatic_balance(eqn, b_b, p0) + eqn.set_reference_profiles([('p', p0), ('b', b_b)]) # Add perturbation r = sqrt((x-Lx/2)**2 + (z-Lz/2)**2) @@ -52,14 +48,14 @@ def run_incompressible(tmpdir): # Set up transport schemes b_opts = SUPGOptions() - transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "b", options=b_opts)] + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "b", options=b_opts)] # Set up linear solver for the timestepping scheme - linear_solver = IncompressibleSolver(state, eqns) + linear_solver = IncompressibleSolver(eqn) # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, + stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, linear_solver=linear_solver) # Run @@ -68,24 +64,25 @@ def run_incompressible(tmpdir): # State for checking checkpoints checkpoint_name = 'incompressible_chkpt' new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') + check_eqn = IncompressibleBoussinesqEquations(domain, parameters) + check_eqn.set_reference_profiles([]) check_output = OutputParameters(dirname=tmpdir+"/incompressible", checkpoint_pickup_filename=new_path) - check_state = State(mesh, dt=dt, output=check_output) - check_eqn = IncompressibleBoussinesqEquations(check_state, "CG", 1) - check_stepper = SemiImplicitQuasiNewton(check_eqn, check_state, []) + check_io = IO(domain, check_eqn, output=check_output) + check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, []) check_stepper.run(t=0, tmax=0, pickup=True) - return state, check_state + return eqn, check_eqn def test_incompressible(tmpdir): dirname = str(tmpdir) - state, check_state = run_incompressible(dirname) + eqn, check_eqn = run_incompressible(dirname) for variable in ['u', 'b', 'p']: - new_variable = state.fields(variable) - check_variable = check_state.fields(variable) + new_variable = eqn.fields(variable) + check_variable = check_eqn.fields(variable) error = norm(new_variable - check_variable) / norm(check_variable) # Slack values chosen to be robust to different platforms diff --git a/integration-tests/equations/test_moist_compressible.py b/integration-tests/equations/test_moist_compressible.py index e7c636bc3..26c58fdef 100644 --- a/integration-tests/equations/test_moist_compressible.py +++ b/integration-tests/equations/test_moist_compressible.py @@ -20,27 +20,24 @@ def run_moist_compressible(tmpdir): Lz = 1000.0 m = PeriodicIntervalMesh(ncols, Lx) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=Lz/nlayers) + domain = Domain(mesh, dt, "CG", 1) - output = OutputParameters(dirname=tmpdir+"/moist_compressible", - dumpfreq=2, chkptfreq=2) parameters = CompressibleParameters() R_d = parameters.R_d R_v = parameters.R_v g = parameters.g tracers = [WaterVapour(name='vapour_mixing_ratio'), CloudWater(name='cloud_liquid_mixing_ratio')] + eqn = CompressibleEulerEquations(domain, parameters, active_tracers=tracers) - state = State(mesh, - dt=dt, - output=output, - parameters=parameters) - - eqn = CompressibleEulerEquations(state, "CG", 1, active_tracers=tracers) + output = OutputParameters(dirname=tmpdir+"/moist_compressible", + dumpfreq=2, chkptfreq=2) + io = IO(domain, eqn, output=output) # Initial conditions - rho0 = state.fields("rho") - theta0 = state.fields("theta") - m_v0 = state.fields("vapour_mixing_ratio") + rho0 = eqn.fields("rho") + theta0 = eqn.fields("theta") + m_v0 = eqn.fields("vapour_mixing_ratio") # Approximate hydrostatic balance x, z = SpatialCoordinate(mesh) @@ -52,7 +49,7 @@ def run_moist_compressible(tmpdir): theta0.interpolate(tde.theta(parameters, T_vd, p)) rho0.interpolate(p / (R_d * T)) - state.set_reference_profiles([('rho', rho0), + eqn.set_reference_profiles([('rho', rho0), ('theta', theta0)]) # Add perturbation @@ -61,15 +58,15 @@ def run_moist_compressible(tmpdir): theta0.interpolate(theta0 + theta_pert) # Set up transport schemes - transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "rho"), - SSPRK3(state, "theta")] + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta")] # Set up linear solver for the timestepping scheme - linear_solver = CompressibleSolver(state, eqn, moisture=['vapour_mixing_ratio']) + linear_solver = CompressibleSolver(eqn, moisture=['vapour_mixing_ratio']) # build time stepper - stepper = SemiImplicitQuasiNewton(eqn, state, transported_fields, + stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, linear_solver=linear_solver) # Run @@ -78,24 +75,25 @@ def run_moist_compressible(tmpdir): # State for checking checkpoints checkpoint_name = 'moist_compressible_chkpt' new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') + check_eqn = CompressibleEulerEquations(domain, parameters, active_tracers=tracers) + check_eqn.set_reference_profiles([]) check_output = OutputParameters(dirname=tmpdir+"/moist_compressible", checkpoint_pickup_filename=new_path) - check_state = State(mesh, dt=dt, output=check_output, parameters=parameters) - check_eqn = CompressibleEulerEquations(check_state, "CG", 1, active_tracers=tracers) - check_stepper = SemiImplicitQuasiNewton(check_eqn, check_state, []) + check_io = IO(domain, check_eqn, output=check_output) + check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, []) check_stepper.run(t=0, tmax=0, pickup=True) - return state, check_state + return eqn, check_eqn def test_moist_compressible(tmpdir): dirname = str(tmpdir) - state, check_state = run_moist_compressible(dirname) + eqn, check_eqn = run_moist_compressible(dirname) for variable in ['u', 'rho', 'theta', 'vapour_mixing_ratio']: - new_variable = state.fields(variable) - check_variable = check_state.fields(variable) + new_variable = eqn.fields(variable) + check_variable = check_eqn.fields(variable) error = norm(new_variable - check_variable) / norm(check_variable) # Slack values chosen to be robust to different platforms diff --git a/integration-tests/equations/test_sw_linear_triangle.py b/integration-tests/equations/test_sw_linear_triangle.py index 57ffe22be..6ad9f18b4 100644 --- a/integration-tests/equations/test_sw_linear_triangle.py +++ b/integration-tests/equations/test_sw_linear_triangle.py @@ -11,6 +11,8 @@ def setup_sw(dirname): + + dt = 3600. refinements = 3 # number of horizontal cells = 20*(4^refinements) R = 6371220. @@ -22,25 +24,21 @@ def setup_sw(dirname): x = SpatialCoordinate(mesh) mesh.init_cell_orientations(x) - dt = 3600. - output = OutputParameters(dirname=dirname+"/sw_linear_w2", steady_state_error_fields=['u', 'D'], dumpfreq=12) - parameters = ShallowWaterParameters(H=H) - - state = State(mesh, - dt=dt, - output=output, - parameters=parameters) + domain = Domain(mesh, dt, "BDM", degree=1) # Coriolis + parameters = ShallowWaterParameters(H=H) Omega = parameters.Omega fexpr = 2*Omega*x[2]/R + eqns = LinearShallowWaterEquations(domain, parameters, fexpr=fexpr) - eqns = LinearShallowWaterEquations(state, "BDM", 1, fexpr=fexpr) + output = OutputParameters(dirname=dirname+"/sw_linear_w2", steady_state_error_fields=['u', 'D'], dumpfreq=12) + io = IO(domain, eqns, output=output) # interpolate initial conditions # Initial/current conditions - u0 = state.fields("u") - D0 = state.fields("D") + u0 = eqns.fields("u") + D0 = eqns.fields("D") u_max = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s) uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) g = parameters.g @@ -48,10 +46,13 @@ def setup_sw(dirname): u0.project(uexpr) D0.interpolate(Dexpr) - transport_schemes = [ForwardEuler(state, "D")] + Dbar = Function(D0.function_space()).assign(H) + eqns.set_reference_profiles([('D', Dbar)]) + + transport_schemes = [ForwardEuler(domain, "D")] # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transport_schemes) + stepper = SemiImplicitQuasiNewton(eqns, io, transport_schemes) return stepper, 2*day diff --git a/integration-tests/equations/test_sw_triangle.py b/integration-tests/equations/test_sw_triangle.py index 86a20256e..406427ac0 100644 --- a/integration-tests/equations/test_sw_triangle.py +++ b/integration-tests/equations/test_sw_triangle.py @@ -26,11 +26,16 @@ def setup_sw(dirname, dt, u_transport_option): mesh = IcosahedralSphereMesh(radius=R, refinement_level=refinements) + domain = Domain(mesh, dt, family="BDM", degree=1) x = SpatialCoordinate(mesh) mesh.init_cell_orientations(x) - output = OutputParameters(dirname=dirname+"/sw", dumplist_latlon=['D', 'D_error'], steady_state_error_fields=['D', 'u']) parameters = ShallowWaterParameters(H=H) + Omega = parameters.Omega + fexpr = 2*Omega*x[2]/R + eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, + u_transport_option=u_transport_option) + diagnostic_fields = [RelativeVorticity(), AbsoluteVorticity(), PotentialVorticity(), ShallowWaterPotentialEnstrophy('RelativeVorticity'), @@ -49,22 +54,12 @@ def setup_sw(dirname, dt, u_transport_option): MeridionalComponent('u'), ZonalComponent('u'), RadialComponent('u')] - - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) - - Omega = parameters.Omega - fexpr = 2*Omega*x[2]/R - eqns = ShallowWaterEquations(state, family="BDM", degree=1, - fexpr=fexpr, - u_transport_option=u_transport_option) + output = OutputParameters(dirname=dirname+"/sw", dumplist_latlon=['D', 'D_error'], steady_state_error_fields=['D', 'u']) + io = IO(domain, eqns, output=output, diagnostic_fields=diagnostic_fields) # interpolate initial conditions - u0 = state.fields("u") - D0 = state.fields("D") + u0 = eqns.fields("u") + D0 = eqns.fields("D") uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) g = parameters.g Dexpr = H - ((R * Omega * u_max + u_max*u_max/2.0)*(x[2]*x[2]/(R*R)))/g @@ -72,17 +67,21 @@ def setup_sw(dirname, dt, u_transport_option): u0.project(uexpr) D0.interpolate(Dexpr) - vspace = FunctionSpace(state.mesh, "CG", 3) + Dbar = Function(D0.function_space()).assign(H) + eqns.set_reference_profiles([('D', Dbar)]) + + vspace = FunctionSpace(domain.mesh, "CG", 3) vexpr = (2*u_max/R)*x[2]/R - f = state.fields("coriolis") - vrel_analytical = state.fields("AnalyticalRelativeVorticity", vspace) + # TODO: these fields should not be in eqns + f = eqns.fields("coriolis") + vrel_analytical = eqns.fields("AnalyticalRelativeVorticity", vspace) vrel_analytical.interpolate(vexpr) - vabs_analytical = state.fields("AnalyticalAbsoluteVorticity", vspace) + vabs_analytical = eqns.fields("AnalyticalAbsoluteVorticity", vspace) vabs_analytical.interpolate(vexpr + f) - pv_analytical = state.fields("AnalyticalPotentialVorticity", vspace) + pv_analytical = eqns.fields("AnalyticalPotentialVorticity", vspace) pv_analytical.interpolate((vexpr+f)/D0) - return state, eqns + return domain, eqns, io def check_results(dirname): @@ -139,12 +138,12 @@ def test_sw_setup(tmpdir, u_transport_option): dirname = str(tmpdir) dt = 1500 - state, eqns = setup_sw(dirname, dt, u_transport_option) + domain, eqns, io = setup_sw(dirname, dt, u_transport_option) transported_fields = [] - transported_fields.append((ImplicitMidpoint(state, "u"))) - transported_fields.append((SSPRK3(state, "D"))) - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields) + transported_fields.append((ImplicitMidpoint(domain, "u"))) + transported_fields.append((SSPRK3(domain, "D"))) + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields) stepper.run(t=0, tmax=0.25*day) check_results(dirname) @@ -157,9 +156,9 @@ def test_sw_ssprk3(tmpdir, u_transport_option): dirname = str(tmpdir) dt = 100 - state, eqns = setup_sw(dirname, dt, u_transport_option) + domain, eqns, io = setup_sw(dirname, dt, u_transport_option) - stepper = Timestepper(eqns, SSPRK3(state), state) + stepper = Timestepper(eqns, SSPRK3(domain), io) stepper.run(t=0, tmax=0.01*day) check_results(dirname) diff --git a/integration-tests/model/test_checkpointing.py b/integration-tests/model/test_checkpointing.py index 2087cac9c..694aebe32 100644 --- a/integration-tests/model/test_checkpointing.py +++ b/integration-tests/model/test_checkpointing.py @@ -20,25 +20,25 @@ def setup_checkpointing(dirname): # build volume mesh H = 1.0e4 # Height position of the model top mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) - domain = Domain(mesh, "CG", 1) + domain = Domain(mesh, dt, "CG", 1) parameters = CompressibleParameters() eqns = CompressibleEulerEquations(domain, parameters) output = OutputParameters(dirname=dirname, dumpfreq=1, chkptfreq=2, log_level='INFO') - io = IO(domain, eqns, dt=dt, output=output) + io = IO(domain, eqns, output=output) initialise_fields(eqns) # Set up transport schemes transported_fields = [] - transported_fields.append(SSPRK3(domain, io, "u")) - transported_fields.append(SSPRK3(domain, io, "rho")) - transported_fields.append(SSPRK3(domain, io, "theta")) + transported_fields.append(SSPRK3(domain, "u")) + transported_fields.append(SSPRK3(domain, "rho")) + transported_fields.append(SSPRK3(domain, "theta")) # Set up linear solver - linear_solver = CompressibleSolver(eqns, io) + linear_solver = CompressibleSolver(eqns) # build time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, diff --git a/integration-tests/model/test_passive_tracer.py b/integration-tests/model/test_passive_tracer.py index eae1e2df3..08eb82e27 100644 --- a/integration-tests/model/test_passive_tracer.py +++ b/integration-tests/model/test_passive_tracer.py @@ -18,7 +18,6 @@ def run_tracer(setup): # Get initial conditions from shared config domain = setup.domain mesh = domain.mesh - dt = setup.dt output = setup.output x = SpatialCoordinate(mesh) @@ -35,7 +34,7 @@ def run_tracer(setup): # Equations eqns = LinearShallowWaterEquations(domain, parameters, fexpr=fexpr) tracer_eqn = AdvectionEquation(domain, domain.spaces("DG"), "tracer") - io = IO(domain, eqns, dt=dt, output=output) + io = IO(domain, eqns, output=output) # Specify initial prognostic fields u0 = eqns.fields("u") @@ -54,10 +53,10 @@ def run_tracer(setup): eqns.set_reference_profiles([('D', Dbar)]) # set up transport schemes - transport_schemes = [ForwardEuler(domain, io, "D")] + transport_schemes = [ForwardEuler(domain, "D")] # Set up tracer transport - tracer_transport = [(tracer_eqn, SSPRK3(domain, io))] + tracer_transport = [(tracer_eqn, SSPRK3(domain))] # build time stepper stepper = SemiImplicitQuasiNewton( diff --git a/integration-tests/model/test_prescribed_transport.py b/integration-tests/model/test_prescribed_transport.py index cb792c362..cbe216b33 100644 --- a/integration-tests/model/test_prescribed_transport.py +++ b/integration-tests/model/test_prescribed_transport.py @@ -26,7 +26,7 @@ def test_prescribed_transport_setup(tmpdir, tracer_setup): V = domain.spaces("DG") # Make equation eqn = AdvectionEquation(domain, V, "f") - io = IO(domain, eqn, dt=setup.dt, output=setup.output) + io = IO(domain, eqn, output=setup.output) # Initialise fields def u_evaluation(t): diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py index f38f19bf4..0ba151ef7 100644 --- a/integration-tests/model/test_time_discretisation.py +++ b/integration-tests/model/test_time_discretisation.py @@ -18,19 +18,19 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup): V = domain.spaces("DG") eqn = AdvectionEquation(domain, V, "f") - io = IO(domain, eqn, dt=setup.dt, output=setup.output) + io = IO(domain, eqn, output=setup.output) eqn.fields("f").interpolate(setup.f_init) eqn.fields("u").project(setup.uexpr) if scheme == "ssprk": - transport_scheme = SSPRK3(domain, io) + transport_scheme = SSPRK3(domain) elif scheme == "implicit_midpoint": - transport_scheme = ImplicitMidpoint(domain, io) + transport_scheme = ImplicitMidpoint(domain) elif scheme == "RK4": - transport_scheme = RK4(domain, io) + transport_scheme = RK4(domain) elif scheme == "Heun": - transport_scheme = Heun(domain, io) + transport_scheme = Heun(domain) elif scheme == "BDF2": - transport_scheme = BDF2(domain, io) + transport_scheme = BDF2(domain) assert run(eqn, transport_scheme, io, setup.tmax, setup.f_end) < setup.tol diff --git a/integration-tests/physics/test_condensation.py b/integration-tests/physics/test_condensation.py index 82a6f207e..d48b514df 100644 --- a/integration-tests/physics/test_condensation.py +++ b/integration-tests/physics/test_condensation.py @@ -15,6 +15,9 @@ def run_cond_evap(dirname, process): + + dt = 2.0 + # declare grid shape, with length L and height H L = 1000. H = 1000. @@ -24,7 +27,8 @@ def run_cond_evap(dirname, process): # make mesh m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) - domain = Domain(mesh, "CG", 1) + + domain = Domain(mesh, dt, "CG", 1) x, z = SpatialCoordinate(mesh) @@ -37,11 +41,10 @@ def run_cond_evap(dirname, process): parameters = CompressibleParameters() eqn = CompressibleEulerEquations(domain, parameters, active_tracers=tracers) - dt = 2.0 output = OutputParameters(dirname=dirname+"/cond_evap", dumpfreq=1, dumplist=['u']) - io = IO(domain, eqn, dt=dt, output=output, diagnostic_fields=[Sum('water_vapour', 'cloud_water')]) + io = IO(domain, eqn, output=output, diagnostic_fields=[Sum('water_vapour', 'cloud_water')]) # Declare prognostic fields rho0 = eqn.fields("rho") @@ -87,10 +90,10 @@ def run_cond_evap(dirname, process): eqn.residual = eqn.residual.label_map(lambda t: t.has_label(time_derivative), map_if_true=identity, map_if_false=drop) - physics_schemes = [(SaturationAdjustment(eqn, parameters=parameters), ForwardEuler(domain, io))] + physics_schemes = [(SaturationAdjustment(eqn, parameters=parameters), ForwardEuler(domain))] # build time stepper - scheme = ForwardEuler(domain, io) + scheme = ForwardEuler(domain) stepper = SplitPhysicsTimestepper(eqn, scheme, io, physics_schemes=physics_schemes) diff --git a/integration-tests/physics/test_instant_rain.py b/integration-tests/physics/test_instant_rain.py index fc7820cfd..c35cca902 100644 --- a/integration-tests/physics/test_instant_rain.py +++ b/integration-tests/physics/test_instant_rain.py @@ -18,14 +18,14 @@ def run_instant_rain(dirname): L = 10 nx = 10 mesh = PeriodicSquareMesh(nx, nx, L) - domain = Domain(mesh, "BDM", 1) + dt = 0.1 + domain = Domain(mesh, dt, "BDM", 1) x, y = SpatialCoordinate(mesh) # parameters H = 30 g = 10 fexpr = Constant(0) - dt = 0.1 vapour = WaterVapour(name="water_vapour", space='DG') rain = Rain(name="rain", space="DG", @@ -38,8 +38,8 @@ def run_instant_rain(dirname): output = OutputParameters(dirname=dirname+"/instant_rain", dumpfreq=1, dumplist=['vapour', "rain"]) - diagnostic_fields = [CourantNumber(dt)] - io = IO(domain, eqns, dt=dt, output=output, diagnostic_fields=diagnostic_fields) + diagnostic_fields = [CourantNumber()] + io = IO(domain, eqns, output=output, diagnostic_fields=diagnostic_fields) vapour0 = eqns.fields("water_vapour") @@ -64,9 +64,9 @@ def run_instant_rain(dirname): rain_true = Function(VD).interpolate(vapour0 - saturation) physics_schemes = [(InstantRain(eqns, saturation, rain_name="rain", - set_tau_to_dt=True), ForwardEuler(domain, io))] + set_tau_to_dt=True), ForwardEuler(domain))] - stepper = PrescribedTransport(eqns, RK4(domain, io), io, + stepper = PrescribedTransport(eqns, RK4(domain), io, physics_schemes=physics_schemes) stepper.run(t=0, tmax=5*dt) diff --git a/integration-tests/physics/test_precipitation.py b/integration-tests/physics/test_precipitation.py index ce9c5cc48..995908069 100644 --- a/integration-tests/physics/test_precipitation.py +++ b/integration-tests/physics/test_precipitation.py @@ -14,6 +14,7 @@ def setup_fallout(dirname): # declare grid shape, with length L and height H + dt = 0.1 L = 10. H = 10. nlayers = 10 @@ -22,22 +23,21 @@ def setup_fallout(dirname): # make mesh m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) - domain = Domain(mesh, "CG", 1) + domain = Domain(mesh, dt, "CG", 1) x = SpatialCoordinate(mesh) Vrho = domain.spaces("DG1_equispaced") active_tracers = [Rain(space='DG1_equispaced')] eqn = ForcedAdvectionEquation(domain, Vrho, "rho", active_tracers=active_tracers) - dt = 0.1 output = OutputParameters(dirname=dirname+"/fallout", dumpfreq=10, dumplist=['rain']) diagnostic_fields = [Precipitation()] - io = IO(domain, eqn, dt=dt, output=output, diagnostic_fields=diagnostic_fields) + io = IO(domain, eqn, output=output, diagnostic_fields=diagnostic_fields) - scheme = ForwardEuler(domain, io) + scheme = ForwardEuler(domain) eqn.fields("rho").assign(1.) - physics_schemes = [(Fallout(eqn, 'rain', domain), SSPRK3(domain, io, 'rain'))] + physics_schemes = [(Fallout(eqn, 'rain', domain), SSPRK3(domain, 'rain'))] rain0 = eqn.fields("rain") # set up rain diff --git a/integration-tests/transport/test_dg_transport.py b/integration-tests/transport/test_dg_transport.py index 457296896..4d4003102 100644 --- a/integration-tests/transport/test_dg_transport.py +++ b/integration-tests/transport/test_dg_transport.py @@ -26,11 +26,11 @@ def test_dg_transport_scalar(tmpdir, geometry, equation_form, tracer_setup): else: eqn = ContinuityEquation(domain, V, "f") - io = IO(domain, eqn, dt=setup.dt, output=setup.output) + io = IO(domain, eqn, output=setup.output) eqn.fields("f").interpolate(setup.f_init) eqn.fields("u").project(setup.uexpr) - transport_scheme = SSPRK3(domain, io) + transport_scheme = SSPRK3(domain) error = run(eqn, transport_scheme, io, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' @@ -49,10 +49,10 @@ def test_dg_transport_vector(tmpdir, geometry, equation_form, tracer_setup): else: eqn = ContinuityEquation(domain, V, "f") - io = IO(domain, eqn, dt=setup.dt, output=setup.output) + io = IO(domain, eqn, output=setup.output) eqn.fields("f").interpolate(f_init) eqn.fields("u").project(setup.uexpr) - transport_schemes = SSPRK3(domain, io) + transport_schemes = SSPRK3(domain) f_end = as_vector([setup.f_end]*gdim) error = run(eqn, transport_schemes, io, setup.tmax, f_end) assert error < setup.tol, \ diff --git a/integration-tests/transport/test_embedded_dg_advection.py b/integration-tests/transport/test_embedded_dg_advection.py index 376da9ba4..c6238a811 100644 --- a/integration-tests/transport/test_embedded_dg_advection.py +++ b/integration-tests/transport/test_embedded_dg_advection.py @@ -33,11 +33,11 @@ def test_embedded_dg_advection_scalar(tmpdir, ibp, equation_form, space, else: eqn = ContinuityEquation(domain, V, "f", ibp=ibp) - io = IO(domain, eqn, dt=setup.dt, output=setup.output) + io = IO(domain, eqn, output=setup.output) eqn.fields("f").interpolate(setup.f_init) eqn.fields("u").project(setup.uexpr) - transport_schemes = SSPRK3(domain, io, options=opts) + transport_schemes = SSPRK3(domain, options=opts) error = run(eqn, transport_schemes, io, setup.tmax, setup.f_end) assert error < setup.tol, \ diff --git a/integration-tests/transport/test_limiters.py b/integration-tests/transport/test_limiters.py index 34117717d..2cc02d645 100644 --- a/integration-tests/transport/test_limiters.py +++ b/integration-tests/transport/test_limiters.py @@ -36,7 +36,7 @@ def setup_limiters(dirname, space): degree = 0 if space in ['DG0', 'Vtheta_degree_0'] else 1 - domain = Domain(mesh, family="CG", degree=degree) + domain = Domain(mesh, dt, family="CG", degree=degree) if space == 'DG0': V = domain.spaces('DG') @@ -60,7 +60,7 @@ def setup_limiters(dirname, space): # set up the equation eqn = AdvectionEquation(domain, V, 'tracer') - io = IO(domain, eqn, dt=dt, output=output) + io = IO(domain, eqn, output=output) # ------------------------------------------------------------------------ # # Initial condition @@ -160,18 +160,18 @@ def setup_limiters(dirname, space): recovered_space=VCG1, project_low_method='recover', boundary_method=BoundaryMethod.taylor) - transport_schemes = SSPRK3(domain, io, options=opts, + transport_schemes = SSPRK3(domain, options=opts, limiter=VertexBasedLimiter(VDG1)) elif space == 'DG1': - transport_schemes = SSPRK3(domain, io, limiter=DG1Limiter(V)) + transport_schemes = SSPRK3(domain, limiter=DG1Limiter(V)) elif space == 'DG1_equispaced': - transport_schemes = SSPRK3(domain, io, limiter=VertexBasedLimiter(V)) + transport_schemes = SSPRK3(domain, limiter=VertexBasedLimiter(V)) elif space == 'Vtheta_degree_1': opts = EmbeddedDGOptions() - transport_schemes = SSPRK3(domain, io, options=opts, limiter=ThetaLimiter(V)) + transport_schemes = SSPRK3(domain, options=opts, limiter=ThetaLimiter(V)) else: raise NotImplementedError diff --git a/integration-tests/transport/test_recovered_transport.py b/integration-tests/transport/test_recovered_transport.py index 057e93a06..6584e210c 100644 --- a/integration-tests/transport/test_recovered_transport.py +++ b/integration-tests/transport/test_recovered_transport.py @@ -31,7 +31,7 @@ def test_recovered_space_setup(tmpdir, geometry, tracer_setup): # Make equation eqn = ContinuityEquation(domain, VDG0, "f") - io = IO(domain, eqn, dt=setup.dt, output=setup.output) + io = IO(domain, eqn, output=setup.output) # Initialise fields eqn.fields("f").interpolate(setup.f_init) @@ -42,7 +42,7 @@ def test_recovered_space_setup(tmpdir, geometry, tracer_setup): recovered_space=VCG1, boundary_method=BoundaryMethod.taylor) - transport_scheme = SSPRK3(domain, io, options=recovery_opts) + transport_scheme = SSPRK3(domain, options=recovery_opts) # Run and check error error = run(eqn, transport_scheme, io, setup.tmax, setup.f_end) diff --git a/integration-tests/transport/test_subcycling.py b/integration-tests/transport/test_subcycling.py index 963d3ac9a..4b496314c 100644 --- a/integration-tests/transport/test_subcycling.py +++ b/integration-tests/transport/test_subcycling.py @@ -25,12 +25,12 @@ def test_subcyling(tmpdir, equation_form, tracer_setup): else: eqn = ContinuityEquation(domain, V, "f") - io = IO(domain, eqn, dt=setup.dt, output=setup.output) + io = IO(domain, eqn, output=setup.output) eqn.fields("f").interpolate(setup.f_init) eqn.fields("u").project(setup.uexpr) - transport_scheme = SSPRK3(domain, io, subcycles=2) + transport_scheme = SSPRK3(domain, subcycles=2) error = run(eqn, transport_scheme, io, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_supg_transport.py b/integration-tests/transport/test_supg_transport.py index 8608e4c8f..f88190ec0 100644 --- a/integration-tests/transport/test_supg_transport.py +++ b/integration-tests/transport/test_supg_transport.py @@ -37,15 +37,15 @@ def test_supg_transport_scalar(tmpdir, equation_form, scheme, space, else: eqn = ContinuityEquation(domain, V, "f") - io = IO(domain, eqn, dt=setup.dt, output=setup.output) + io = IO(domain, eqn, output=setup.output) eqn.fields("f").interpolate(setup.f_init) eqn.fields("u").project(setup.uexpr) if scheme == "ssprk": - transport_scheme = SSPRK3(domain, io, options=opts) + transport_scheme = SSPRK3(domain, options=opts) elif scheme == "implicit_midpoint": - transport_scheme = ImplicitMidpoint(domain, io, options=opts) + transport_scheme = ImplicitMidpoint(domain, options=opts) error = run(eqn, transport_scheme, io, setup.tmax, setup.f_end) assert error < setup.tol, \ @@ -77,7 +77,7 @@ def test_supg_transport_vector(tmpdir, equation_form, scheme, space, else: eqn = ContinuityEquation(domain, V, "f") - io = IO(domain, eqn, dt=setup.dt, output=setup.output) + io = IO(domain, eqn, output=setup.output) f = eqn.fields("f") if space == "CG": @@ -86,9 +86,9 @@ def test_supg_transport_vector(tmpdir, equation_form, scheme, space, f.project(f_init) eqn.fields("u").project(setup.uexpr) if scheme == "ssprk": - transport_scheme = SSPRK3(domain, io, options=opts) + transport_scheme = SSPRK3(domain, options=opts) elif scheme == "implicit_midpoint": - transport_scheme = ImplicitMidpoint(domain, io, options=opts) + transport_scheme = ImplicitMidpoint(domain, options=opts) f_end = as_vector([setup.f_end]*gdim) error = run(eqn, transport_scheme, io, setup.tmax, f_end) diff --git a/integration-tests/transport/test_vector_recovered_space.py b/integration-tests/transport/test_vector_recovered_space.py index ee0f19023..8c0f2b88b 100644 --- a/integration-tests/transport/test_vector_recovered_space.py +++ b/integration-tests/transport/test_vector_recovered_space.py @@ -41,14 +41,14 @@ def test_vector_recovered_space_setup(tmpdir, geometry, tracer_setup): # Make equation eqn = AdvectionEquation(domain, Vu, "f") - io = IO(domain, eqn, dt=setup.dt, output=setup.output) + io = IO(domain, eqn, output=setup.output) # Initialise fields f_init = as_vector([setup.f_init]*gdim) eqn.fields("f").project(f_init) eqn.fields("u").project(setup.uexpr) - transport_scheme = SSPRK3(domain, io, options=rec_opts) + transport_scheme = SSPRK3(domain, options=rec_opts) f_end = as_vector([setup.f_end]*gdim) From 28853fca688a07a0c2d81b56d5b6a35fff7c53c5 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Wed, 14 Dec 2022 20:59:34 +0000 Subject: [PATCH 07/12] fix problems with tests --- gusto/equations.py | 7 +++++-- gusto/io.py | 2 +- gusto/time_discretisation.py | 2 +- integration-tests/equations/test_moist_compressible.py | 3 +-- integration-tests/model/test_prescribed_transport.py | 2 +- 5 files changed, 9 insertions(+), 7 deletions(-) diff --git a/gusto/equations.py b/gusto/equations.py index 4318f800d..b7a8f60cc 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -48,7 +48,7 @@ def __init__(self, domain, function_space, field_name): if len(function_space) > 1: assert hasattr(self, "field_names") self.fields(field_name, function_space, - subfield_names=self.field_names, pickup=True) + subfield_names=self.field_names, pickup=True) for fname in self.field_names: self.bcs[fname] = [] else: @@ -223,7 +223,10 @@ def __init__(self, field_names, domain, linearisation_map=None, self.reference_profiles_initialised = False # Build finite element spaces - self.spaces = domain.compatible_spaces + # TODO: this implies order of spaces matches order of variables + # we should not assume this and should instead specify which variable + # is in which space + self.spaces = [space for space in domain.compatible_spaces] # Add active tracers to the list of prognostics if active_tracers is None: diff --git a/gusto/io.py b/gusto/io.py index 5c0d2c946..d6212c5ec 100644 --- a/gusto/io.py +++ b/gusto/io.py @@ -232,7 +232,7 @@ def setup_diagnostics(self): self.diagnostic_fields.append(f) for name in self.output.steady_state_error_fields: - f = SteadyStateError(self, name) + f = SteadyStateError(self.equation, name) self.diagnostic_fields.append(f) fields = set([f.name() for f in self.equation.fields]) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 1e818d4dc..326c20efe 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -130,7 +130,7 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): self.equation = equation self.residual = equation.residual - if self.field_name is not None: + if self.field_name is not None and hasattr(equation, "field_names"): self.idx = equation.field_names.index(self.field_name) self.fs = equation.fields(self.field_name).function_space() self.residual = self.residual.label_map( diff --git a/integration-tests/equations/test_moist_compressible.py b/integration-tests/equations/test_moist_compressible.py index 26c58fdef..2b9e71ddc 100644 --- a/integration-tests/equations/test_moist_compressible.py +++ b/integration-tests/equations/test_moist_compressible.py @@ -49,8 +49,7 @@ def run_moist_compressible(tmpdir): theta0.interpolate(tde.theta(parameters, T_vd, p)) rho0.interpolate(p / (R_d * T)) - eqn.set_reference_profiles([('rho', rho0), - ('theta', theta0)]) + eqn.set_reference_profiles([('rho', rho0), ('theta', theta0)]) # Add perturbation r = sqrt((x-Lx/2)**2 + (z-Lz/2)**2) diff --git a/integration-tests/model/test_prescribed_transport.py b/integration-tests/model/test_prescribed_transport.py index cbe216b33..754f6880b 100644 --- a/integration-tests/model/test_prescribed_transport.py +++ b/integration-tests/model/test_prescribed_transport.py @@ -36,7 +36,7 @@ def u_evaluation(t): eqn.fields("f").interpolate(setup.f_init) eqn.fields("u").project(u_evaluation(Constant(0.0))) - transport_scheme = SSPRK3(domain, io) + transport_scheme = SSPRK3(domain) # Run and check error error = run(eqn, transport_scheme, io, setup.tmax, From 6390badc53889ca42b713aa9eb32c0de06e84529 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Wed, 21 Dec 2022 13:17:36 +0000 Subject: [PATCH 08/12] PR #320: move fields into stepper, reorganise diagnostics and kill state --- .../compressible/dcmip_3_1_meanflow_quads.py | 72 +- examples/compressible/dry_bryan_fritsch.py | 106 +- examples/compressible/moist_bryan_fritsch.py | 113 +- examples/compressible/mountain_hydrostatic.py | 128 +- .../compressible/mountain_nonhydrostatic.py | 89 +- examples/compressible/robert_bubble.py | 75 +- .../skamarock_klemp_hydrostatic.py | 85 +- .../compressible/skamarock_klemp_nonlinear.py | 89 +- examples/compressible/straka_bubble.py | 91 +- examples/compressible/unsaturated_bubble.py | 161 +- .../skamarock_klemp_incompressible.py | 74 +- examples/shallow_water/linear_williamson_2.py | 53 +- examples/shallow_water/williamson_2.py | 56 +- examples/shallow_water/williamson_5.py | 63 +- gusto/configuration.py | 4 - gusto/diagnostics.py | 1410 +++++++++-------- gusto/domain.py | 3 +- gusto/equations.py | 117 +- gusto/fields.py | 165 +- gusto/function_spaces.py | 10 +- gusto/initialisation_tools.py | 24 +- gusto/io.py | 124 +- gusto/linear_solvers.py | 60 +- gusto/physics.py | 48 +- gusto/state.py | 807 ---------- gusto/time_discretisation.py | 14 +- gusto/timeloop.py | 126 +- .../balance/test_compressible_balance.py | 44 +- .../balance/test_saturated_balance.py | 69 +- .../balance/test_unsaturated_balance.py | 60 +- integration-tests/conftest.py | 11 +- integration-tests/diffusion/test_diffusion.py | 26 +- .../equations/test_advection_diffusion.py | 31 +- .../equations/test_dry_compressible.py | 62 +- .../equations/test_forced_advection.py | 46 +- .../equations/test_incompressible.py | 57 +- .../equations/test_moist_compressible.py | 68 +- .../equations/test_sw_linear_triangle.py | 36 +- .../equations/test_sw_triangle.py | 56 +- integration-tests/model/test_checkpointing.py | 32 +- .../model/test_passive_tracer.py | 45 +- .../model/test_prescribed_transport.py | 20 +- .../model/test_time_discretisation.py | 18 +- .../physics/test_condensation.py | 56 +- .../physics/test_instant_rain.py | 45 +- .../physics/test_precipitation.py | 29 +- .../transport/test_dg_transport.py | 31 +- .../transport/test_embedded_dg_advection.py | 16 +- integration-tests/transport/test_limiters.py | 79 +- .../transport/test_recovered_transport.py | 19 +- .../transport/test_subcycling.py | 17 +- .../transport/test_supg_transport.py | 36 +- .../transport/test_vector_recovered_space.py | 16 +- 53 files changed, 2527 insertions(+), 2665 deletions(-) delete mode 100644 gusto/state.py diff --git a/examples/compressible/dcmip_3_1_meanflow_quads.py b/examples/compressible/dcmip_3_1_meanflow_quads.py index a0f93c33b..7b7caa20d 100644 --- a/examples/compressible/dcmip_3_1_meanflow_quads.py +++ b/examples/compressible/dcmip_3_1_meanflow_quads.py @@ -11,6 +11,9 @@ from firedrake import exp, acos, cos, sin, pi, sqrt, asin, atan_2 import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # nlayers = 10 # Number of vertical layers refinements = 3 # Number of horiz. refinements @@ -24,7 +27,6 @@ tmax = 3600.0 dumpfreq = int(tmax / (4*dt)) - parameters = CompressibleParameters() a_ref = 6.37122e6 # Radius of the Earth (m) X = 125.0 # Reduced-size Earth reduction factor @@ -43,18 +45,24 @@ phi_c = 0.0 # Latitudinal centerpoint of Theta' (equator) deltaTheta = 1.0 # Maximum amplitude of Theta' (K) L_z = 20000.0 # Vertical wave length of the Theta' perturb. +z_top = 1.0e4 # Height position of the model top + +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # +# Domain # Cubed-sphere horizontal mesh m = CubedSphereMesh(radius=a, refinement_level=refinements, degree=2) - # Build volume mesh -z_top = 1.0e4 # Height position of the model top mesh = ExtrudedMesh(m, layers=nlayers, layer_height=z_top/nlayers, extrusion_type="radial") +domain = Domain(mesh, dt, "RTCF", 1) x = SpatialCoordinate(mesh) + # Create polar coordinates: # Since we use a CG1 field, this is constant on layers W_Q1 = FunctionSpace(mesh, "CG", 1) @@ -64,29 +72,41 @@ lat = Function(W_Q1).interpolate(lat_expr) lon = Function(W_Q1).interpolate(atan_2(x[1], x[0])) -dirname = 'dcmip_3_1_meanflow' +# Equation +eqns = CompressibleEulerEquations(domain, parameters) +# I/O +dirname = 'dcmip_3_1_meanflow' output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, - perturbation_fields=['theta', 'rho'], log_level='INFO') +diagnostic_fields = [Perturbation('theta'), Perturbation('rho')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -state = State(mesh, - dt=dt, - output=output, - parameters=parameters) +# Transport schemes +transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho", subcycles=2), + SSPRK3(domain, "theta", options=SUPGOptions(), subcycles=2)] -eqns = CompressibleEulerEquations(state, "RTCF", 1) +# Linear solver +linear_solver = CompressibleSolver(eqns) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) + +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields.u -theta0 = state.fields.theta -rho0 = state.fields.rho +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields.u +theta0 = stepper.fields.theta +rho0 = stepper.fields.rho # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") # Initial conditions with u0 uexpr = as_vector([-u_max*x[1]/a, u_max*x[0]/a, 0.0]) @@ -122,7 +142,7 @@ theta0.interpolate(theta_b) # Compute the balanced density -compressible_hydrostatic_balance(state, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, top=False, @@ -131,20 +151,12 @@ theta0 += theta_b rho0.assign(rho_b) -state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)]) -state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) +stepper.initialise([('u', u0), ('rho', rho0), ('theta', theta0)]) +stepper.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) -# Set up transport schemes -transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "rho", subcycles=2), - SSPRK3(state, "theta", options=SUPGOptions(), subcycles=2)] - -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns) - -# Build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # # Run! stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/dry_bryan_fritsch.py b/examples/compressible/dry_bryan_fritsch.py index f7775bb21..842064165 100644 --- a/examples/compressible/dry_bryan_fritsch.py +++ b/examples/compressible/dry_bryan_fritsch.py @@ -13,7 +13,14 @@ FunctionSpace, VectorFunctionSpace) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 1.0 +L = 10000. +H = 10000. + if '--running-tests' in sys.argv: deltax = 1000. tmax = 5. @@ -24,43 +31,70 @@ degree = 0 dirname = 'dry_bryan_fritsch' -# make mesh -L = 10000. -H = 10000. +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain nlayers = int(H/deltax) ncolumns = int(L/deltax) m = IntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) +domain = Domain(mesh, dt, "CG", degree) +# Equation +params = CompressibleParameters() +u_transport_option = "vector_advection_form" +eqns = CompressibleEulerEquations(domain, params, + u_transport_option=u_transport_option, + no_normal_flow_bc_ids=[1, 2]) +# I/O output = OutputParameters(dirname=dirname, dumpfreq=int(tmax / (5*dt)), dumplist=['u'], - perturbation_fields=['theta'], log_level='INFO') +diagnostic_fields = [Perturbation('theta')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -params = CompressibleParameters() +# Transport schemes -- set up options for using recovery wrapper +VDG1 = domain.spaces("DG1_equispaced") +VCG1 = FunctionSpace(mesh, "CG", 1) +Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) +Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) -state = State(mesh, - dt=dt, - output=output, - parameters=params) +u_opts = RecoveryOptions(embedding_space=Vu_DG1, + recovered_space=Vu_CG1, + boundary_method=BoundaryMethod.taylor) +rho_opts = RecoveryOptions(embedding_space=VDG1, + recovered_space=VCG1, + boundary_method=BoundaryMethod.taylor) +theta_opts = RecoveryOptions(embedding_space=VDG1, + recovered_space=VCG1) -u_transport_option = "vector_advection_form" +transported_fields = [SSPRK3(domain, "rho", options=rho_opts), + SSPRK3(domain, "theta", options=theta_opts), + SSPRK3(domain, "u", options=u_opts)] -eqns = CompressibleEulerEquations(state, "CG", degree, - u_transport_option=u_transport_option, - no_normal_flow_bc_ids=[1, 2]) +# Linear solver +linear_solver = CompressibleSolver(eqns) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) + +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") x, z = SpatialCoordinate(mesh) # Define constant theta_e and water_t @@ -68,7 +102,7 @@ theta_b = Function(Vt).interpolate(Constant(Tsurf)) # Calculate hydrostatic fields -compressible_hydrostatic_balance(state, theta_b, rho0, solve_for_rho=True) +compressible_hydrostatic_balance(eqns, theta_b, rho0, solve_for_rho=True) # make mean fields rho_b = Function(Vr).assign(rho0) @@ -95,33 +129,11 @@ rho_solver = LinearVariationalSolver(rho_problem) rho_solver.solve() -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - -# Set up transport schemes -VDG1 = state.spaces("DG1_equispaced") -VCG1 = FunctionSpace(mesh, "CG", 1) -Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) -Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) - -u_opts = RecoveryOptions(embedding_space=Vu_DG1, - recovered_space=Vu_CG1, - boundary_method=BoundaryMethod.taylor) -rho_opts = RecoveryOptions(embedding_space=VDG1, - recovered_space=VCG1, - boundary_method=BoundaryMethod.taylor) -theta_opts = RecoveryOptions(embedding_space=VDG1, - recovered_space=VCG1) - -transported_fields = [SSPRK3(state, "rho", options=rho_opts), - SSPRK3(state, "theta", options=theta_opts), - SSPRK3(state, "u", options=u_opts)] +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns) - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/moist_bryan_fritsch.py b/examples/compressible/moist_bryan_fritsch.py index dd5977846..7e3e65493 100644 --- a/examples/compressible/moist_bryan_fritsch.py +++ b/examples/compressible/moist_bryan_fritsch.py @@ -15,7 +15,14 @@ LinearVariationalProblem, LinearVariationalSolver) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 1.0 +L = 10000. +H = 10000. + if '--running-tests' in sys.argv: deltax = 1000. tmax = 5. @@ -23,47 +30,65 @@ deltax = 200 tmax = 1000. -L = 10000. -H = 10000. +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain nlayers = int(H/deltax) ncolumns = int(L/deltax) m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) degree = 1 +domain = Domain(mesh, dt, 'CG', degree) -dirname = 'moist_bryan_fritsch' +# Equation +params = CompressibleParameters() +tracers = [WaterVapour(), CloudWater()] +eqns = CompressibleEulerEquations(domain, params, active_tracers=tracers) +# I/O +dirname = 'moist_bryan_fritsch' output = OutputParameters(dirname=dirname, dumpfreq=int(tmax / (5*dt)), dumplist=['u'], - perturbation_fields=[], log_level='INFO') +diagnostic_fields = [Theta_e(eqns)] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -params = CompressibleParameters() -diagnostic_fields = [Theta_e()] -tracers = [WaterVapour(), CloudWater()] +# Transport schemes +transported_fields = [SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=EmbeddedDGOptions()), + SSPRK3(domain, "water_vapour", options=EmbeddedDGOptions()), + SSPRK3(domain, "cloud_water", options=EmbeddedDGOptions()), + ImplicitMidpoint(domain, "u")] + +# Linear solver +linear_solver = CompressibleSolver(eqns) -state = State(mesh, - dt=dt, - output=output, - parameters=params, - diagnostic_fields=diagnostic_fields) +# Physics schemes (condensation/evaporation) +physics_schemes = [(SaturationAdjustment(eqns), ForwardEuler(domain))] -eqns = CompressibleEulerEquations(state, "CG", degree, active_tracers=tracers) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver, + physics_schemes=physics_schemes) +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") -water_v0 = state.fields("water_vapour") -water_c0 = state.fields("cloud_water") -moisture = ["water_vapour", "cloud_water"] +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") +water_v0 = stepper.fields("water_vapour") +water_c0 = stepper.fields("cloud_water") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") x, z = SpatialCoordinate(mesh) quadrature_degree = (4, 4) dxp = dx(degree=(quadrature_degree)) @@ -75,15 +100,15 @@ water_t = Function(Vt).assign(total_water) # Calculate hydrostatic fields -saturated_hydrostatic_balance(state, theta_e, water_t) +saturated_hydrostatic_balance(eqns, stepper.fields, theta_e, water_t) # make mean fields theta_b = Function(Vt).assign(theta0) rho_b = Function(Vr).assign(rho0) water_vb = Function(Vt).assign(water_v0) water_cb = Function(Vt).assign(water_t - water_vb) -exner_b = thermodynamics.exner_pressure(state.parameters, rho_b, theta_b) -Tb = thermodynamics.T(state.parameters, theta_b, exner_b, r_v=water_vb) +exner_b = thermodynamics.exner_pressure(eqns.parameters, rho_b, theta_b) +Tb = thermodynamics.T(eqns.parameters, theta_b, exner_b, r_v=water_vb) # define perturbation xc = L / 2 @@ -115,10 +140,10 @@ rho_recoverer = Recoverer(rho0, rho_averaged) rho_recoverer.project() -exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) -p = thermodynamics.p(state.parameters, exner) -T = thermodynamics.T(state.parameters, theta0, exner, r_v=w_v) -w_sat = thermodynamics.r_sat(state.parameters, T, p) +exner = thermodynamics.exner_pressure(eqns.parameters, rho_averaged, theta0) +p = thermodynamics.p(eqns.parameters, exner) +T = thermodynamics.T(eqns.parameters, theta0, exner, r_v=w_v) +w_sat = thermodynamics.r_sat(eqns.parameters, T, p) w_functional = (phi * w_v * dxp - phi * w_sat * dxp) w_problem = NonlinearVariationalProblem(w_functional, w_v) @@ -128,29 +153,13 @@ water_v0.assign(w_v) water_c0.assign(water_t - water_v0) -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b), - ('water_vapour', water_vb)]) - -rho_opts = None -theta_opts = EmbeddedDGOptions() -u_transport = ImplicitMidpoint(state, "u") +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b), + ('water_vapour', water_vb), + ('cloud_water', water_cb)]) -transported_fields = [SSPRK3(state, "rho", options=rho_opts), - SSPRK3(state, "theta", options=theta_opts), - SSPRK3(state, "water_vapour", options=theta_opts), - SSPRK3(state, "cloud_water", options=theta_opts), - u_transport] - -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns, moisture=moisture) - -# define condensation -physics_schemes = [(SaturationAdjustment(eqns, params), ForwardEuler(state))] - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver, - physics_schemes=physics_schemes) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/mountain_hydrostatic.py b/examples/compressible/mountain_hydrostatic.py index 944e101b7..d97744bd3 100644 --- a/examples/compressible/mountain_hydrostatic.py +++ b/examples/compressible/mountain_hydrostatic.py @@ -9,7 +9,13 @@ exp, pi, cos, Function, conditional, Mesh, op2, sqrt) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 5.0 +L = 240000. # Domain length +H = 50000. # Height position of the model top if '--running-tests' in sys.argv: tmax = dt @@ -21,13 +27,15 @@ dumpfreq = int(tmax / (5*dt)) +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain nlayers = res*20 # horizontal layers columns = res*12 # number of columns -L = 240000. m = PeriodicIntervalMesh(columns, L) -# build volume mesh -H = 50000. # Height position of the model top ext_mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) Vc = VectorFunctionSpace(ext_mesh, "DG", 2) coord = SpatialCoordinate(ext_mesh) @@ -38,42 +46,74 @@ hm = 1. zs = hm*a**2/((x-xc)**2 + a**2) -dirname = 'hydrostatic_mountain' zh = 5000. xexpr = as_vector([x, conditional(z < zh, z + cos(0.5*pi*z/zh)**6*zs, z)]) new_coords = Function(Vc).interpolate(xexpr) mesh = Mesh(new_coords) +domain = Domain(mesh, dt, "CG", 1) + +# Equation +parameters = CompressibleParameters(g=9.80665, cp=1004.) +sponge = SpongeLayerParameters(H=H, z_level=H-20000, mubar=0.3/dt) +eqns = CompressibleEulerEquations(domain, parameters, sponge=sponge) +# I/O +dirname = 'hydrostatic_mountain' output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, dumplist=['u'], - perturbation_fields=['theta', 'rho'], log_level='INFO') +diagnostic_fields = [CourantNumber(), VelocityZ(), HydrostaticImbalance(eqns), + Perturbation('theta'), Perturbation('rho')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -parameters = CompressibleParameters(g=9.80665, cp=1004.) -diagnostic_fields = [CourantNumber(), VelocityZ(), HydrostaticImbalance()] +# Transport schemes +theta_opts = SUPGOptions() +transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=theta_opts)] -state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) +# Linear solver +params = {'mat_type': 'matfree', + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'firedrake.SCPC', + # Velocity mass operator is singular in the hydrostatic case. + # So for reconstruction, we eliminate rho into u + 'pc_sc_eliminate_fields': '1, 0', + 'condensed_field': {'ksp_type': 'fgmres', + 'ksp_rtol': 1.0e-8, + 'ksp_atol': 1.0e-8, + 'ksp_max_it': 100, + 'pc_type': 'gamg', + 'pc_gamg_sym_graph': True, + 'mg_levels': {'ksp_type': 'gmres', + 'ksp_max_it': 5, + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu'}}} -# sponge function -sponge = SpongeLayerParameters(H=H, z_level=H-20000, mubar=0.3/dt) +alpha = 0.51 # off-centering parameter +linear_solver = CompressibleSolver(eqns, alpha, solver_parameters=params, + overwrite_solver_parameters=True) -eqns = CompressibleEulerEquations(state, "CG", 1, sponge=sponge) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver, + alpha=alpha) +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") # Thermodynamic constants required for setting initial conditions # and reference profiles @@ -107,7 +147,7 @@ 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}} -compressible_hydrostatic_balance(state, theta_b, rho_b, exner, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner, top=True, exner_boundary=0.5, params=exner_params) @@ -123,13 +163,13 @@ def minimum(f): p0 = minimum(exner) -compressible_hydrostatic_balance(state, theta_b, rho_b, exner, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner, top=True, params=exner_params) p1 = minimum(exner) alpha = 2.*(p1-p0) beta = p1-alpha exner_top = (1.-beta)/alpha -compressible_hydrostatic_balance(state, theta_b, rho_b, exner, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner, top=True, exner_boundary=exner_top, solve_for_rho=True, params=exner_params) @@ -138,41 +178,11 @@ def minimum(f): u0.project(as_vector([20.0, 0.0])) remove_initial_w(u0) -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - -# Set up transport schemes -theta_opts = SUPGOptions() -transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "rho"), - SSPRK3(state, "theta", options=theta_opts)] - -# Set up linear solver -params = {'mat_type': 'matfree', - 'ksp_type': 'preonly', - 'pc_type': 'python', - 'pc_python_type': 'firedrake.SCPC', - # Velocity mass operator is singular in the hydrostatic case. - # So for reconstruction, we eliminate rho into u - 'pc_sc_eliminate_fields': '1, 0', - 'condensed_field': {'ksp_type': 'fgmres', - 'ksp_rtol': 1.0e-8, - 'ksp_atol': 1.0e-8, - 'ksp_max_it': 100, - 'pc_type': 'gamg', - 'pc_gamg_sym_graph': True, - 'mg_levels': {'ksp_type': 'gmres', - 'ksp_max_it': 5, - 'pc_type': 'bjacobi', - 'sub_pc_type': 'ilu'}}} - -alpha = 0.51 # off-centering parameter -linear_solver = CompressibleSolver(state, eqns, alpha, solver_parameters=params, - overwrite_solver_parameters=True) +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver, - alpha=alpha) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/mountain_nonhydrostatic.py b/examples/compressible/mountain_nonhydrostatic.py index 56509715b..b8e6656be 100644 --- a/examples/compressible/mountain_nonhydrostatic.py +++ b/examples/compressible/mountain_nonhydrostatic.py @@ -9,7 +9,14 @@ exp, pi, cos, Function, conditional, Mesh, op2) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 5.0 +L = 144000. # Domain length +H = 35000. # Height position of the model top + if '--running-tests' in sys.argv: tmax = dt dumpfreq = 1 @@ -21,11 +28,12 @@ nlayers = 70 # horizontal layers columns = 180 # number of columns -L = 144000. -m = PeriodicIntervalMesh(columns, L) +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # -# build volume mesh -H = 35000. # Height position of the model top +# Domain +m = PeriodicIntervalMesh(columns, L) ext_mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) Vc = VectorFunctionSpace(ext_mesh, "DG", 2) coord = SpatialCoordinate(ext_mesh) @@ -36,42 +44,52 @@ hm = 1. zs = hm*a**2/((x-xc)**2 + a**2) -dirname = 'nonhydrostatic_mountain' zh = 5000. xexpr = as_vector([x, conditional(z < zh, z + cos(0.5*pi*z/zh)**6*zs, z)]) new_coords = Function(Vc).interpolate(xexpr) mesh = Mesh(new_coords) +domain = Domain(mesh, dt, "CG", 1) +# Equation +parameters = CompressibleParameters(g=9.80665, cp=1004.) +sponge = SpongeLayerParameters(H=H, z_level=H-10000, mubar=0.15/dt) +eqns = CompressibleEulerEquations(domain, parameters, sponge=sponge) + +# I/O +dirname = 'nonhydrostatic_mountain' output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, dumplist=['u'], - perturbation_fields=['theta', 'rho'], log_level='INFO') +diagnostic_fields = [CourantNumber(), VelocityZ(), Perturbation('theta'), Perturbation('rho')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -parameters = CompressibleParameters(g=9.80665, cp=1004.) -diagnostic_fields = [CourantNumber(), VelocityZ()] - -state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) +# Transport schemes +theta_opts = SUPGOptions() +transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=theta_opts)] -# sponge function -sponge = SpongeLayerParameters(H=H, z_level=H-10000, mubar=0.15/dt) +# Linear solver +linear_solver = CompressibleSolver(eqns) -eqns = CompressibleEulerEquations(state, "CG", 1, sponge=sponge) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") # Thermodynamic constants required for setting initial conditions # and reference profiles @@ -103,7 +121,7 @@ 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}} -compressible_hydrostatic_balance(state, theta_b, rho_b, exner, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner, top=True, exner_boundary=0.5, params=exner_params) @@ -119,13 +137,13 @@ def minimum(f): p0 = minimum(exner) -compressible_hydrostatic_balance(state, theta_b, rho_b, exner, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner, top=True, params=exner_params) p1 = minimum(exner) alpha = 2.*(p1-p0) beta = p1-alpha exner_top = (1.-beta)/alpha -compressible_hydrostatic_balance(state, theta_b, rho_b, exner, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner, top=True, exner_boundary=exner_top, solve_for_rho=True, params=exner_params) @@ -134,20 +152,11 @@ def minimum(f): u0.project(as_vector([10.0, 0.0])) remove_initial_w(u0) -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) -# Set up transport schemes -theta_opts = SUPGOptions() -transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "rho"), - SSPRK3(state, "theta", options=theta_opts)] - -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns) - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/robert_bubble.py b/examples/compressible/robert_bubble.py index 1c4cf397a..52c305a65 100644 --- a/examples/compressible/robert_bubble.py +++ b/examples/compressible/robert_bubble.py @@ -9,6 +9,10 @@ Constant, pi, cos, Function, sqrt, conditional) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 1. L = 1000. H = 1000. @@ -24,37 +28,54 @@ nlayers = int(H/10.) ncolumns = int(L/10.) +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) +domain = Domain(mesh, dt, "CG", 1) -dirname = 'robert_bubble' +# Equation +parameters = CompressibleParameters() +eqns = CompressibleEulerEquations(domain, parameters) +# I/O +dirname = 'robert_bubble' output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, dumplist=['u'], - perturbation_fields=['theta', 'rho'], log_level='INFO') +diagnostic_fields = [CourantNumber(), Perturbation('theta'), Perturbation('rho')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -parameters = CompressibleParameters() -diagnostic_fields = [CourantNumber()] +# Transport schemes +theta_opts = EmbeddedDGOptions() +transported_fields = [] +transported_fields.append(ImplicitMidpoint(domain, "u")) +transported_fields.append(SSPRK3(domain, "rho")) +transported_fields.append(SSPRK3(domain, "theta", options=theta_opts)) -state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) +# Linear solver +linear_solver = CompressibleSolver(eqns) -eqns = CompressibleEulerEquations(state, "CG", 1) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") # Isentropic background state Tsurf = Constant(300.) @@ -63,7 +84,7 @@ rho_b = Function(Vr) # Calculate hydrostatic exner -compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True) +compressible_hydrostatic_balance(eqns, theta_b, rho_b, solve_for_rho=True) x = SpatialCoordinate(mesh) xc = 500. @@ -75,21 +96,11 @@ theta0.interpolate(theta_b + theta_pert) rho0.interpolate(rho_b) -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) -# Set up transport schemes -theta_opts = EmbeddedDGOptions() -transported_fields = [] -transported_fields.append(ImplicitMidpoint(state, "u")) -transported_fields.append(SSPRK3(state, "rho")) -transported_fields.append(SSPRK3(state, "theta", options=theta_opts)) - -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns) - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/skamarock_klemp_hydrostatic.py b/examples/compressible/skamarock_klemp_hydrostatic.py index 9f3ba1770..7e2e12255 100644 --- a/examples/compressible/skamarock_klemp_hydrostatic.py +++ b/examples/compressible/skamarock_klemp_hydrostatic.py @@ -10,6 +10,10 @@ ExtrudedMesh, exp, sin, Function, pi) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 25. if '--running-tests' in sys.argv: nlayers = 5 # horizontal layers @@ -22,45 +26,59 @@ tmax = 60000.0 dumpfreq = int(tmax / (2*dt)) +L = 6.0e6 # Length of domain +H = 1.0e4 # Height position of the model top -L = 6.0e6 -m = PeriodicRectangleMesh(columns, 1, L, 1.e4, quadrilateral=True) +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # -# build volume mesh -H = 1.0e4 # Height position of the model top +# Domain -- 3D volume mesh +m = PeriodicRectangleMesh(columns, 1, L, 1.e4, quadrilateral=True) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) +domain = Domain(mesh, dt, "RTCF", 1) -dirname = 'skamarock_klemp_hydrostatic' +# Equation +parameters = CompressibleParameters() +Omega = as_vector((0., 0., 0.5e-4)) +balanced_pg = as_vector((0., -1.0e-4*20, 0.)) +eqns = CompressibleEulerEquations(domain, parameters, Omega=Omega, + extra_terms=[("u", balanced_pg)]) +# I/O +dirname = 'skamarock_klemp_hydrostatic' output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, dumplist=['u'], - perturbation_fields=['theta', 'rho'], log_level='INFO') +diagnostic_fields = [CourantNumber(), Perturbation('theta'), Perturbation('rho')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -parameters = CompressibleParameters() -diagnostic_fields = [CourantNumber()] +# Transport schemes +transported_fields = [] +transported_fields.append(ImplicitMidpoint(domain, "u")) +transported_fields.append(SSPRK3(domain, "rho")) +transported_fields.append(SSPRK3(domain, "theta", options=SUPGOptions())) -state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) +# Linear solver +linear_solver = CompressibleSolver(eqns) -Omega = as_vector((0., 0., 0.5e-4)) -balanced_pg = as_vector((0., -1.0e-4*20, 0.)) -eqns = CompressibleEulerEquations(state, "RTCF", 1, Omega=Omega, - extra_terms=[("u", balanced_pg)]) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") # Thermodynamic constants required for setting initial conditions # and reference profiles @@ -85,26 +103,17 @@ theta_pert = deltaTheta*sin(pi*z/H)/(1 + (x - L/2)**2/a**2) theta0.interpolate(theta_b + theta_pert) -compressible_hydrostatic_balance(state, theta_b, rho_b, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, solve_for_rho=True) rho0.assign(rho_b) u0.project(as_vector([20.0, 0.0, 0.0])) -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - -# Set up transport schemes -transported_fields = [] -transported_fields.append(ImplicitMidpoint(state, "u")) -transported_fields.append(SSPRK3(state, "rho")) -transported_fields.append(SSPRK3(state, "theta", options=SUPGOptions())) +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns) - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/skamarock_klemp_nonlinear.py b/examples/compressible/skamarock_klemp_nonlinear.py index 984fd4e06..15d947a68 100644 --- a/examples/compressible/skamarock_klemp_nonlinear.py +++ b/examples/compressible/skamarock_klemp_nonlinear.py @@ -13,7 +13,14 @@ import numpy as np import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 6. +L = 3.0e5 # Domain length +H = 1.0e4 # Height position of the model top + if '--running-tests' in sys.argv: nlayers = 5 columns = 30 @@ -25,53 +32,61 @@ tmax = 3600. dumpfreq = int(tmax / (2*dt)) +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # -L = 3.0e5 +# Domain -- 3D volume mesh m = PeriodicIntervalMesh(columns, L) - -# build volume mesh -H = 1.0e4 # Height position of the model top mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) +domain = Domain(mesh, dt, "CG", 1) +# Equation +Tsurf = 300. +parameters = CompressibleParameters() +eqns = CompressibleEulerEquations(domain, parameters) + +# I/O points_x = np.linspace(0., L, 100) points_z = [H/2.] points = np.array([p for p in itertools.product(points_x, points_z)]) - dirname = 'skamarock_klemp_nonlinear' - output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, pddumpfreq=dumpfreq, dumplist=['u'], - perturbation_fields=['theta', 'rho'], point_data=[('theta_perturbation', points)], log_level='INFO') +diagnostic_fields = [CourantNumber(), Gradient("u"), Perturbation('theta'), + Gradient("theta_perturbation"), Perturbation('rho'), + RichardsonNumber("theta", parameters.g/Tsurf), Gradient("theta")] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -parameters = CompressibleParameters() -g = parameters.g -Tsurf = 300. - -diagnostic_fields = [CourantNumber(), Gradient("u"), - Gradient("theta_perturbation"), - RichardsonNumber("theta", g/Tsurf), Gradient("theta")] +# Transport schemes +theta_opts = SUPGOptions() +transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=theta_opts)] -state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) +# Linear solver +linear_solver = CompressibleSolver(eqns) -eqns = CompressibleEulerEquations(state, "CG", 1) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") # Thermodynamic constants required for setting initial conditions # and reference profiles @@ -81,14 +96,13 @@ x, z = SpatialCoordinate(mesh) # N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) -Tsurf = 300. thetab = Tsurf*exp(N**2*z/g) theta_b = Function(Vt).interpolate(thetab) rho_b = Function(Vr) # Calculate hydrostatic exner -compressible_hydrostatic_balance(state, theta_b, rho_b) +compressible_hydrostatic_balance(eqns, theta_b, rho_b) a = 5.0e3 deltaTheta = 1.0e-2 @@ -97,20 +111,11 @@ rho0.assign(rho_b) u0.project(as_vector([20.0, 0.0])) -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - -# Set up transport schemes -theta_opts = SUPGOptions() -transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "rho"), - SSPRK3(state, "theta", options=theta_opts)] +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns) - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/straka_bubble.py b/examples/compressible/straka_bubble.py index 14e7cbda4..972ab491e 100644 --- a/examples/compressible/straka_bubble.py +++ b/examples/compressible/straka_bubble.py @@ -10,6 +10,10 @@ conditional) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + if '--running-tests' in sys.argv: res_dt = {800.: 4.} tmax = 4. @@ -27,45 +31,65 @@ for delta, dt in res_dt.items(): - dirname = "straka_dx%s_dt%s" % (delta, dt) + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain nlayers = int(H/delta) # horizontal layers columns = int(L/delta) # number of columns - m = PeriodicIntervalMesh(columns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) + domain = Domain(mesh, dt, "CG", 1) + # Equation + parameters = CompressibleParameters() + diffusion_options = [ + ("u", DiffusionParameters(kappa=75., mu=10./delta)), + ("theta", DiffusionParameters(kappa=75., mu=10./delta))] + eqns = CompressibleEulerEquations(domain, parameters, + diffusion_options=diffusion_options) + + # I/O + dirname = "straka_dx%s_dt%s" % (delta, dt) dumpfreq = int(tmax / (ndumps*dt)) output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, dumplist=['u'], - perturbation_fields=['theta', 'rho'], log_level='INFO') + diagnostic_fields = [CourantNumber(), Perturbation('theta'), Perturbation('rho')] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) - parameters = CompressibleParameters() - diagnostic_fields = [CourantNumber()] + # Transport schemes + theta_opts = SUPGOptions() + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=theta_opts)] - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) + # Linear solver + linear_solver = CompressibleSolver(eqns) - diffusion_options = [ - ("u", DiffusionParameters(kappa=75., mu=10./delta)), - ("theta", DiffusionParameters(kappa=75., mu=10./delta))] + # Diffusion schemes + diffusion_schemes = [BackwardEuler(domain, "u"), + BackwardEuler(domain, "theta")] - eqns = CompressibleEulerEquations(state, "CG", 1, - diffusion_options=diffusion_options) + # Time stepper + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver, + diffusion_schemes=diffusion_schemes) + # ------------------------------------------------------------------------ # # Initial conditions - u0 = state.fields("u") - rho0 = state.fields("rho") - theta0 = state.fields("theta") + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") # spaces - Vu = state.spaces("HDiv") - Vt = state.spaces("theta") - Vr = state.spaces("DG") + Vu = domain.spaces("HDiv") + Vt = domain.spaces("theta") + Vr = domain.spaces("DG") # Isentropic background state Tsurf = Constant(300.) @@ -75,7 +99,7 @@ exner = Function(Vr) # Calculate hydrostatic exner - compressible_hydrostatic_balance(state, theta_b, rho_b, exner0=exner, + compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner0=exner, solve_for_rho=True) x = SpatialCoordinate(mesh) @@ -89,24 +113,11 @@ theta0.interpolate(theta_b + T_pert*exner) rho0.assign(rho_b) - state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) + stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) - # Set up transport schemes - theta_opts = SUPGOptions() - transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "rho"), - SSPRK3(state, "theta", options=theta_opts)] - - # Set up linear solver - linear_solver = CompressibleSolver(state, eqns) - - diffusion_schemes = [BackwardEuler(state, "u"), - BackwardEuler(state, "theta")] - - # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver, - diffusion_schemes=diffusion_schemes) + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/unsaturated_bubble.py b/examples/compressible/unsaturated_bubble.py index 3de8dddec..b07eaa34e 100644 --- a/examples/compressible/unsaturated_bubble.py +++ b/examples/compressible/unsaturated_bubble.py @@ -14,6 +14,10 @@ from firedrake.slope_limiter.vertex_based_limiter import VertexBasedLimiter import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 1.0 if '--running-tests' in sys.argv: deltax = 240. @@ -28,63 +32,101 @@ h = 2400. nlayers = int(h/deltax) ncolumns = int(L/deltax) +degree = 0 +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=h/nlayers) -degree = 0 +domain = Domain(mesh, dt, "CG", degree) +# Equation +params = CompressibleParameters() +tracers = [WaterVapour(), CloudWater(), Rain()] +eqns = CompressibleEulerEquations(domain, params, + active_tracers=tracers) + +# I/O dirname = 'unsaturated_bubble' output = OutputParameters(dirname=dirname, dumpfreq=tdump, - perturbation_fields=['theta', 'water_vapour', 'rho'], log_level='INFO') -params = CompressibleParameters() -diagnostic_fields = [RelativeHumidity()] -tracers = [WaterVapour(), CloudWater(), Rain()] +diagnostic_fields = [RelativeHumidity(eqns), Perturbation('theta'), + Perturbation('water_vapour'), Perturbation('rho')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -state = State(mesh, - dt=dt, - output=output, - parameters=params, - diagnostic_fields=diagnostic_fields) +# Transport schemes -- specify options for using recovery wrapper +VDG1 = domain.spaces("DG1_equispaced") +VCG1 = FunctionSpace(mesh, "CG", 1) +Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) +Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) -eqns = CompressibleEulerEquations(state, "CG", degree, - active_tracers=tracers) +u_opts = RecoveryOptions(embedding_space=Vu_DG1, + recovered_space=Vu_CG1, + boundary_method=BoundaryMethod.taylor) +rho_opts = RecoveryOptions(embedding_space=VDG1, + recovered_space=VCG1, + boundary_method=BoundaryMethod.taylor) +theta_opts = RecoveryOptions(embedding_space=VDG1, recovered_space=VCG1) +limiter = VertexBasedLimiter(VDG1) +transported_fields = [SSPRK3(domain, "u", options=u_opts), + SSPRK3(domain, "rho", options=rho_opts), + SSPRK3(domain, "theta", options=theta_opts), + SSPRK3(domain, "water_vapour", options=theta_opts, limiter=limiter), + SSPRK3(domain, "cloud_water", options=theta_opts, limiter=limiter), + SSPRK3(domain, "rain", options=theta_opts, limiter=limiter)] + +# Linear solver +linear_solver = CompressibleSolver(eqns) + +# Physics schemes +# NB: to use wrapper options with Fallout, need to pass field name to time discretisation +physics_schemes = [(Fallout(eqns, 'rain', domain), SSPRK3(domain, field_name='rain', options=theta_opts, limiter=limiter)), + (Coalescence(eqns), ForwardEuler(domain)), + (EvaporationOfRain(eqns), ForwardEuler(domain)), + (SaturationAdjustment(eqns), ForwardEuler(domain))] + +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver, + physics_schemes=physics_schemes) + +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") -water_v0 = state.fields("water_vapour") -water_c0 = state.fields("cloud_water") -rain0 = state.fields("rain", theta0.function_space()) -moisture = ["water_vapour", "cloud_water", "rain"] +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") +water_v0 = stepper.fields("water_vapour") +water_c0 = stepper.fields("cloud_water") +rain0 = stepper.fields("rain") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") x, z = SpatialCoordinate(mesh) quadrature_degree = (4, 4) dxp = dx(degree=(quadrature_degree)) -VDG1 = state.spaces("DG1_equispaced") -VCG1 = FunctionSpace(mesh, "CG", 1) -Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) -Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) physics_boundary_method = BoundaryMethod.extruded # Define constant theta_e and water_t Tsurf = 283.0 psurf = 85000. -exner_surf = (psurf / state.parameters.p_0) ** state.parameters.kappa +exner_surf = (psurf / eqns.parameters.p_0) ** eqns.parameters.kappa humidity = 0.2 S = 1.3e-5 -theta_surf = thermodynamics.theta(state.parameters, Tsurf, psurf) +theta_surf = thermodynamics.theta(eqns.parameters, Tsurf, psurf) theta_d = Function(Vt).interpolate(theta_surf * exp(S*z)) H = Function(Vt).assign(humidity) # Calculate hydrostatic fields -unsaturated_hydrostatic_balance(state, theta_d, H, +unsaturated_hydrostatic_balance(eqns, stepper.fields, theta_d, H, exner_boundary=Constant(exner_surf)) # make mean fields @@ -116,21 +158,21 @@ w_h = Function(Vt) delta = 1.0 -R_d = state.parameters.R_d -R_v = state.parameters.R_v +R_d = eqns.parameters.R_d +R_v = eqns.parameters.R_v epsilon = R_d / R_v # make expressions for determining water_v0 -exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) -p = thermodynamics.p(state.parameters, exner) -T = thermodynamics.T(state.parameters, theta0, exner, water_v0) -r_v_expr = thermodynamics.r_v(state.parameters, H, T, p) +exner = thermodynamics.exner_pressure(eqns.parameters, rho_averaged, theta0) +p = thermodynamics.p(eqns.parameters, exner) +T = thermodynamics.T(eqns.parameters, theta0, exner, water_v0) +r_v_expr = thermodynamics.r_v(eqns.parameters, H, T, p) # make expressions to evaluate residual -exner_ev = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) -p_ev = thermodynamics.p(state.parameters, exner_ev) -T_ev = thermodynamics.T(state.parameters, theta0, exner_ev, water_v0) -RH_ev = thermodynamics.RH(state.parameters, water_v0, T_ev, p_ev) +exner_ev = thermodynamics.exner_pressure(eqns.parameters, rho_averaged, theta0) +p_ev = thermodynamics.p(eqns.parameters, exner_ev) +T_ev = thermodynamics.T(eqns.parameters, theta0, exner_ev, water_v0) +RH_ev = thermodynamics.RH(eqns.parameters, water_v0, T_ev, p_ev) RH = Function(Vt) # set-up rho problem to keep exner constant @@ -175,40 +217,11 @@ raise RuntimeError('Hydrostatic balance solve has not converged within %i' % i, 'iterations') # initialise fields -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b), - ('water_vapour', water_vb)]) - -# Set up transport schemes -u_opts = RecoveryOptions(embedding_space=Vu_DG1, - recovered_space=Vu_CG1, - boundary_method=BoundaryMethod.taylor) -rho_opts = RecoveryOptions(embedding_space=VDG1, - recovered_space=VCG1, - boundary_method=BoundaryMethod.taylor) -theta_opts = RecoveryOptions(embedding_space=VDG1, recovered_space=VCG1) -limiter = VertexBasedLimiter(VDG1) - -transported_fields = [SSPRK3(state, "u", options=u_opts), - SSPRK3(state, "rho", options=rho_opts), - SSPRK3(state, "theta", options=theta_opts), - SSPRK3(state, "water_vapour", options=theta_opts, limiter=limiter), - SSPRK3(state, "cloud_water", options=theta_opts, limiter=limiter), - SSPRK3(state, "rain", options=theta_opts, limiter=limiter)] - -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns, moisture=moisture) - -# define physics schemes -# NB: to use wrapper options with Fallout, need to pass field name to time discretisation -physics_schemes = [(Fallout(eqns, 'rain', state), SSPRK3(state, field_name='rain', options=theta_opts, limiter=limiter)), - (Coalescence(eqns), ForwardEuler(state)), - (EvaporationOfRain(eqns, params), ForwardEuler(state)), - (SaturationAdjustment(eqns, params), ForwardEuler(state))] - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver, - physics_schemes=physics_schemes) +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b), + ('water_vapour', water_vb)]) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/incompressible/skamarock_klemp_incompressible.py b/examples/incompressible/skamarock_klemp_incompressible.py index 8b05c1ded..e2296967e 100644 --- a/examples/incompressible/skamarock_klemp_incompressible.py +++ b/examples/incompressible/skamarock_klemp_incompressible.py @@ -10,7 +10,14 @@ sin, SpatialCoordinate, Function, pi) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 6. +L = 3.0e5 # Domain length +H = 1.0e4 # Height position of the model top + if '--running-tests' in sys.argv: tmax = dt dumpfreq = 1 @@ -23,38 +30,47 @@ columns = 300 # number of columns nlayers = 10 # horizontal layers -# set up mesh -L = 3.0e5 +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain m = PeriodicIntervalMesh(columns, L) -H = 1.0e4 # Height position of the model top mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) +domain = Domain(mesh, dt, 'CG', 1) -# output parameters +# Equation +parameters = CompressibleParameters() +eqns = IncompressibleBoussinesqEquations(domain, parameters) + +# I/O output = OutputParameters(dirname='skamarock_klemp_incompressible', dumpfreq=dumpfreq, dumplist=['u'], - perturbation_fields=['b'], log_level='INFO') - -# physical parameters -parameters = CompressibleParameters() - # list of diagnostic fields, each defined in a class in diagnostics.py -diagnostic_fields = [CourantNumber(), Divergence()] +diagnostic_fields = [CourantNumber(), Divergence(), Perturbation('b')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -# setup state -state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) +# Transport schemes +b_opts = SUPGOptions() +transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "b", options=b_opts)] -eqns = IncompressibleBoussinesqEquations(state, "CG", 1) +# Linear solver +linear_solver = IncompressibleSolver(eqns) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) + +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -b0 = state.fields("b") -p0 = state.fields("p") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +b0 = stepper.fields("b") +p0 = stepper.fields("p") # spaces Vb = b0.function_space() @@ -75,25 +91,17 @@ # interpolate the expression to the function b0.interpolate(b_b + b_pert) -incompressible_hydrostatic_balance(state, b_b, p0) +incompressible_hydrostatic_balance(eqns, b_b, p0) uinit = (as_vector([20.0, 0.0])) u0.project(uinit) # set the background buoyancy -state.set_reference_profiles([('b', b_b)]) +stepper.set_reference_profiles([('b', b_b)]) -# Set up transport schemes -b_opts = SUPGOptions() -transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "b", options=b_opts)] - -# Set up linear solver for the timestepping scheme -linear_solver = IncompressibleSolver(state, eqns) - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # # Run! stepper.run(t=0, tmax=tmax) diff --git a/examples/shallow_water/linear_williamson_2.py b/examples/shallow_water/linear_williamson_2.py index 88c4ba768..75930ac65 100644 --- a/examples/shallow_water/linear_williamson_2.py +++ b/examples/shallow_water/linear_williamson_2.py @@ -9,6 +9,10 @@ from firedrake import IcosahedralSphereMesh, SpatialCoordinate, as_vector, pi import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 3600. day = 24.*60.*60. if '--running-tests' in sys.argv: @@ -23,32 +27,43 @@ R = 6371220. H = 2000. +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain mesh = IcosahedralSphereMesh(radius=R, refinement_level=refinements, degree=3) x = SpatialCoordinate(mesh) mesh.init_cell_orientations(x) +domain = Domain(mesh, dt, 'BDM', 1) + +# Equation +parameters = ShallowWaterParameters(H=H) +Omega = parameters.Omega +x = SpatialCoordinate(mesh) +fexpr = 2*Omega*x[2]/R +eqns = LinearShallowWaterEquations(domain, parameters, fexpr=fexpr) +# I/O output = OutputParameters(dirname='linear_williamson_2', dumpfreq=dumpfreq, - steady_state_error_fields=['u', 'D'], log_level='INFO') -parameters = ShallowWaterParameters(H=H) +diagnostic_fields = [SteadyStateError('u'), SteadyStateError('D')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -state = State(mesh, - dt=dt, - output=output, - parameters=parameters) +# Transport schemes +transport_schemes = [ForwardEuler(domain, "D")] -# Coriolis expression -Omega = parameters.Omega -x = SpatialCoordinate(mesh) -fexpr = 2*Omega*x[2]/R -eqns = LinearShallowWaterEquations(state, "BDM", 1, fexpr=fexpr) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transport_schemes) + +# ---------------------------------------------------------------------------- # +# Initial conditions +# ---------------------------------------------------------------------------- # -# interpolate initial conditions -# Initial/current conditions -u0 = state.fields("u") -D0 = state.fields("D") +u0 = stepper.fields("u") +D0 = stepper.fields("D") u_max = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s) uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) g = parameters.g @@ -56,9 +71,11 @@ u0.project(uexpr) D0.interpolate(Dexpr) -transport_schemes = [ForwardEuler(state, "D")] +Dbar = Function(D0.function_space()).assign(H) +stepper.set_reference_profiles([('D', Dbar)]) -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transport_schemes) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/shallow_water/williamson_2.py b/examples/shallow_water/williamson_2.py index 9ff00b3c4..1053798d5 100644 --- a/examples/shallow_water/williamson_2.py +++ b/examples/shallow_water/williamson_2.py @@ -10,6 +10,10 @@ from firedrake import IcosahedralSphereMesh, SpatialCoordinate, as_vector, pi import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + day = 24.*60.*60. if '--running-tests' in sys.argv: ref_dt = {3: 3000.} @@ -30,38 +34,51 @@ for ref_level, dt in ref_dt.items(): - dirname = "williamson_2_ref%s_dt%s" % (ref_level, dt) + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain mesh = IcosahedralSphereMesh(radius=R, refinement_level=ref_level, degree=1) x = SpatialCoordinate(mesh) global_normal = x mesh.init_cell_orientations(x) + domain = Domain(mesh, dt, 'BDM', 1) + + # Equation + Omega = parameters.Omega + fexpr = 2*Omega*x[2]/R + eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr) + # I/O + dirname = "williamson_2_ref%s_dt%s" % (ref_level, dt) dumpfreq = int(tmax / (ndumps*dt)) output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, dumplist_latlon=['D', 'D_error'], - steady_state_error_fields=['D', 'u'], log_level='INFO') diagnostic_fields = [RelativeVorticity(), PotentialVorticity(), ShallowWaterKineticEnergy(), - ShallowWaterPotentialEnergy(), - ShallowWaterPotentialEnstrophy()] + ShallowWaterPotentialEnergy(parameters), + ShallowWaterPotentialEnstrophy(), + SteadyStateError('u'), SteadyStateError('D')] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) + # Transport schemes + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "D", subcycles=2)] - Omega = parameters.Omega - fexpr = 2*Omega*x[2]/R - eqns = ShallowWaterEquations(state, "BDM", 1, fexpr=fexpr) + # Time stepper + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # - # interpolate initial conditions - u0 = state.fields("u") - D0 = state.fields("D") + u0 = stepper.fields("u") + D0 = stepper.fields("D") x = SpatialCoordinate(mesh) u_max = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s) uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) @@ -71,10 +88,11 @@ u0.project(uexpr) D0.interpolate(Dexpr) - transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "D", subcycles=2)] + Dbar = Function(D0.function_space()).assign(H) + stepper.set_reference_profiles([('D', Dbar)]) - # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields) + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # stepper.run(t=0, tmax=tmax) diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index 8bfc1d144..db586d6db 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -10,6 +10,10 @@ as_vector, pi, sqrt, Min) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + day = 24.*60.*60. if '--running-tests' in sys.argv: ref_dt = {3: 3000.} @@ -30,26 +34,18 @@ for ref_level, dt in ref_dt.items(): - dirname = "williamson_5_ref%s_dt%s" % (ref_level, dt) + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain mesh = IcosahedralSphereMesh(radius=R, refinement_level=ref_level, degree=1) x = SpatialCoordinate(mesh) mesh.init_cell_orientations(x) + domain = Domain(mesh, dt, 'BDM', 1) - dumpfreq = int(tmax / (ndumps*dt)) - output = OutputParameters(dirname=dirname, - dumplist_latlon=['D'], - dumpfreq=dumpfreq, - log_level='INFO') - - diagnostic_fields = [Sum('D', 'topography')] - - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) - + # Equation Omega = parameters.Omega fexpr = 2*Omega*x[2]/R theta, lamda = latlon_coords(mesh) @@ -62,11 +58,31 @@ rsq = Min(R0sq, lsq+thsq) r = sqrt(rsq) bexpr = 2000 * (1 - r/R0) - eqns = ShallowWaterEquations(state, "BDM", 1, fexpr=fexpr, bexpr=bexpr) + eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, bexpr=bexpr) + + # I/O + dirname = "williamson_5_ref%s_dt%s" % (ref_level, dt) + dumpfreq = int(tmax / (ndumps*dt)) + output = OutputParameters(dirname=dirname, + dumplist_latlon=['D'], + dumpfreq=dumpfreq, + log_level='INFO') + diagnostic_fields = [Sum('D', 'topography')] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "D")] + + # Time stepper + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # - # interpolate initial conditions - u0 = state.fields('u') - D0 = state.fields('D') + u0 = stepper.fields('u') + D0 = stepper.fields('D') u_max = 20. # Maximum amplitude of the zonal wind (m/s) uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) g = parameters.g @@ -76,10 +92,11 @@ u0.project(uexpr) D0.interpolate(Dexpr) - transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "D")] + Dbar = Function(D0.function_space()).assign(H) + stepper.set_reference_profiles([('D', Dbar)]) - # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields) + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # stepper.run(t=0, tmax=tmax) diff --git a/gusto/configuration.py b/gusto/configuration.py index 6b97638ac..dc6efa5e7 100644 --- a/gusto/configuration.py +++ b/gusto/configuration.py @@ -116,10 +116,6 @@ class OutputParameters(Configuration): #: TODO: Should the output fields be interpolated or projected to #: a linear space? Default is interpolation. project_fields = False - #: List of fields to dump error fields for steady state simulation - steady_state_error_fields = [] - #: List of fields for computing perturbations from the initial state - perturbation_fields = [] #: List of ordered pairs (name, points) where name is the field # name and points is the points at which to dump them point_data = [] diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index df7b337dd..7e12a8655 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -1,14 +1,18 @@ """Common diagnostic fields.""" from firedrake import op2, assemble, dot, dx, FunctionSpace, Function, sqrt, \ - TestFunction, TrialFunction, Constant, grad, inner, \ + TestFunction, TrialFunction, Constant, grad, inner, curl, \ LinearVariationalProblem, LinearVariationalSolver, FacetNormal, \ - ds, ds_b, ds_v, ds_t, dS_v, div, avg, jump, DirichletBC, \ - TensorFunctionSpace, SpatialCoordinate, VectorFunctionSpace, as_vector + ds_b, ds_v, ds_t, dS_v, div, avg, jump, \ + TensorFunctionSpace, SpatialCoordinate, as_vector, \ + Projector, Interpolator +from firedrake.assign import Assigner from abc import ABCMeta, abstractmethod, abstractproperty -from gusto import thermodynamics +import gusto.thermodynamics as tde from gusto.recovery import Recoverer, BoundaryMethod +from gusto.equations import CompressibleEulerEquations +from gusto.active_tracers import TracerVariableType, Phases import numpy as np __all__ = ["Diagnostics", "CourantNumber", "VelocityX", "VelocityZ", "VelocityY", "Gradient", @@ -130,304 +134,325 @@ def total(f): class DiagnosticField(object, metaclass=ABCMeta): """Base object to represent diagnostic fields for outputting.""" - def __init__(self, required_fields=()): + 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 + + # 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 - def setup(self, eqn, space=None): + @abstractmethod + def setup(self, domain, state_fields, space=None): """ Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. + 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 space is None: - space = eqn.domain.spaces("DG0", "DG", 0) + if self.space is None: + if space is None: + space = domain.spaces("DG0", "DG", 0) + self.space = space + else: + space = self.space - # TODO: would it be better for these diagnostics to be stored in the IO? - self.field = eqn.fields(self.name, space, pickup=False) - self._initialised = True + self.field = state_fields(self.name, space=space, dump=True, pickup=False) - @abstractmethod - def compute(self, eqn): - """ - Compute the diagnostic field from the current state. + if self.method != 'solve': + assert self.expr is not None, \ + f"The expression for diagnostic {self.name} has not been specified" - Args: - eqn (:class:`PrognosticEquation`): the model's equation. - """ - pass + # 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) - def __call__(self, eqn): - """ - Compute the diagnostic field from the current state. + self._initialised = True - Args: - eqn (:class:`PrognosticEquation`): the model's equation. - """ - return self.compute(eqn) + def compute(self): + """Compute the diagnostic field from the current state.""" + + 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 setup(self, eqn): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - if not self._initialised: - super(CourantNumber, self).setup(eqn) - self.dt = eqn.domain.dt - # set up area computation - V = eqn.domain.spaces("DG0") - test = TestFunction(V) - self.area = Function(V) - assemble(test*dx, tensor=self.area) - def compute(self, eqn): - """ - Compute and return the diagnostic field from the current state. + # set up area computation + V = domain.spaces("DG0", "DG", 0) + test = TestFunction(V) + self.area = Function(V) + assemble(test*dx, tensor=self.area) + u = state_fields("u") - Args: - eqn (:class:`PrognosticEquation`): the model's equation. + self.expr = sqrt(dot(u, u))/sqrt(self.area)*domain.dt - Returns: - :class:`Function`: the diagnostic field. - """ - # TODO: update dt here? - u = eqn.fields("u") - return self.field.project(sqrt(dot(u, u))/sqrt(self.area)*self.dt) + super(CourantNumber, self).setup(domain, state_fields) +# TODO: unify all component diagnostics class VelocityX(DiagnosticField): """The geocentric Cartesian X component of the velocity field.""" name = "VelocityX" - def setup(self, eqn): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - """ - if not self._initialised: - space = eqn.domain.spaces("CG1", "CG", 1) - super(VelocityX, self).setup(eqn, space=space) - - def compute(self, eqn): - """ - Compute and return the diagnostic field from the current state. - - Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - u = eqn.fields("u") - uh = u[0] - return self.field.interpolate(uh) + u = state_fields("u") + self.expr = u[0] + super(VelocityX, self).setup(domain, state_fields) class VelocityZ(DiagnosticField): """The geocentric Cartesian Z component of the velocity field.""" name = "VelocityZ" - def setup(self, eqn): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - """ - if not self._initialised: - space = eqn.domain.spaces("CG1", "CG", 1) - super(VelocityZ, self).setup(eqn, space=space) - - def compute(self, eqn): - """ - Compute and return the diagnostic field from the current state. - - Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - u = eqn.fields("u") - w = u[u.geometric_dimension() - 1] - return self.field.interpolate(w) + u = state_fields("u") + self.expr = u[u.geometric_dimension() - 1] + super(VelocityZ, self).setup(domain, state_fields) class VelocityY(DiagnosticField): """The geocentric Cartesian Y component of the velocity field.""" name = "VelocityY" - def setup(self, eqn): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - if not self._initialised: - space = eqn.domain.spaces("CG1", "CG", 1) - super(VelocityY, self).setup(eqn, space=space) - - def compute(self, eqn): - """ - Compute and return the diagnostic field from the current state. - - Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. - """ - u = eqn.fields("u") - v = u[1] - return self.field.interpolate(v) + u = state_fields("u") + self.expr = u[1] + super(VelocityY, self).setup(domain, state_fields) class Gradient(DiagnosticField): """Diagnostic for computing the gradient of fields.""" - def __init__(self, name): + 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'. """ - super().__init__() 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, eqn): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - if not self._initialised: - mesh_dim = eqn.domain.mesh.geometric_dimension() - try: - field_dim = eqn.fields(self.fname).ufl_shape[0] - except IndexError: - field_dim = 1 - shape = (mesh_dim, ) * field_dim - space = TensorFunctionSpace(eqn.domain.mesh, "CG", 1, shape=shape) - super().setup(eqn, space=space) - - f = eqn.fields(self.fname) - test = TestFunction(space) - trial = TrialFunction(space) - n = FacetNormal(eqn.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.solver = LinearVariationalSolver(prob) + f = state_fields(self.fname) - def compute(self, eqn): - """ - Compute and return the diagnostic field from the current state. + 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) - Args: - eqn (:class:`PrognosticEquation`): the model's equation. + if self.method != 'solve': + self.expr = grad(f) - Returns: - :class:`Function`: the diagnostic field. - """ - self.solver.solve() - return self.field + 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.""" - name = "Divergence" - - def setup(self, eqn): + def __init__(self, name='u', space=None, method='interpolate'): """ - Sets up the :class:`Function` for the diagnostic field. - Args: - eqn (:class:`PrognosticEquation`): the model's equation. + 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'. """ - if not self._initialised: - space = eqn.domain.spaces("DG1", "DG", 1) - super(Divergence, self).setup(eqn, space=space) + self.fname = name + super().__init__(space=space, method=method, required_fields=(self.fname,)) - def compute(self, eqn): + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.fname+"_divergence" + + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - u = eqn.fields("u") - return self.field.interpolate(div(u)) + f = state_fields(self.fname) + self.expr = div(f) + space = domain.spaces("DG") + super(Divergence, self).setup(domain, state_fields, space=space) class SphericalComponent(DiagnosticField): """Base diagnostic for computing spherical-polar components of fields.""" - def __init__(self, name): + 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'. """ - super().__init__() self.fname = name + super().__init__(space=space, method=method, required_fields=(name,)) - def setup(self, eqn): + # TODO: these routines must be moved to somewhere more available generally + # (e.g. initialisation tools?) + def _spherical_polar_unit_vectors(self, domain): """ - Sets up the :class:`Function` for the diagnostic field. + Generate ufl expressions for the spherical polar unit vectors. Args: - eqn (:class:`PrognosticEquation`): the model's equation. + domain (:class:`Domain`): the model's domain, containing its mesh. + + Returns: + tuple of (:class:`ufl.Expr`): the zonal, meridional and radial unit + vectors. """ - if not self._initialised: - # check geometric dimension is 3D - if eqn.domain.mesh.geometric_dimension() != 3: - raise ValueError('Spherical components only work when the geometric dimension is 3!') - space = FunctionSpace(eqn.domain.mesh, "CG", 1) - super().setup(eqn, space=space) - - V = VectorFunctionSpace(eqn.domain.mesh, "CG", 1) - self.x, self.y, self.z = SpatialCoordinate(eqn.domain.mesh) - self.x_hat = Function(V).interpolate(Constant(as_vector([1.0, 0.0, 0.0]))) - self.y_hat = Function(V).interpolate(Constant(as_vector([0.0, 1.0, 0.0]))) - self.z_hat = Function(V).interpolate(Constant(as_vector([0.0, 0.0, 1.0]))) - self.R = sqrt(self.x**2 + self.y**2) # distance from z axis - self.r = sqrt(self.x**2 + self.y**2 + self.z**2) # distance from origin - self.f = eqn.fields(self.fname) - if np.prod(self.f.ufl_shape) != 3: + x, y, z = SpatialCoordinate(domain.mesh) + x_hat = Constant(as_vector([1.0, 0.0, 0.0])) + y_hat = Constant(as_vector([0.0, 1.0, 0.0])) + z_hat = Constant(as_vector([0.0, 0.0, 1.0])) + R = sqrt(x**2 + y**2) # distance from z axis + r = sqrt(x**2 + y**2 + z**2) # distance from origin + + lambda_hat = (-x*z/R * x_hat - y*z/R * y_hat + R * z_hat) / r + phi_hat = (x * y_hat - y * x_hat) / R + r_hat = (x * x_hat + y * y_hat + z * z_hat) / r + + return lambda_hat, phi_hat, r_hat + + 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.') @@ -438,20 +463,19 @@ def name(self): """Gives the name of this diagnostic field.""" return self.fname+"_meridional" - def compute(self, eqn): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - lambda_hat = (-self.x * self.z * self.x_hat / self.R - - self.y * self.z * self.y_hat / self.R - + self.R * self.z_hat) / self.r - return self.field.project(dot(self.f, lambda_hat)) + f = state_fields(self.fname) + self._check_args(domain, f) + lambda_hat, _, _ = self._spherical_polar_unit_vectors(domain) + self.expr = dot(f, lambda_hat) + super().setup(domain, state_fields) class ZonalComponent(SphericalComponent): @@ -461,18 +485,19 @@ def name(self): """Gives the name of this diagnostic field.""" return self.fname+"_zonal" - def compute(self, eqn): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - phi_hat = (self.x * self.y_hat - self.y * self.x_hat) / self.R - return self.field.project(dot(self.f, phi_hat)) + f = state_fields(self.fname) + self._check_args(domain, f) + _, phi_hat, _ = self._spherical_polar_unit_vectors(domain) + self.expr = dot(f, phi_hat) + super().setup(domain, state_fields) class RadialComponent(SphericalComponent): @@ -482,67 +507,65 @@ def name(self): """Gives the name of this diagnostic field.""" return self.fname+"_radial" - def compute(self, eqn): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - r_hat = (self.x * self.x_hat + self.y * self.y_hat + self.z * self.z_hat) / self.r - return self.field.project(dot(self.f, r_hat)) + f = state_fields(self.fname) + self._check_args(domain, f) + _, _, r_hat = self._spherical_polar_unit_vectors(domain) + self.expr = dot(f, r_hat) + super().setup(domain, state_fields) class RichardsonNumber(DiagnosticField): """Dimensionless Richardson number diagnostic field.""" name = "RichardsonNumber" - def __init__(self, density_field, factor=1.): + def __init__(self, density_field, factor=1., space=None, method='interpolate'): u""" Args: - density_field (:class:`Function`): the density field. + 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. - """ - super().__init__(required_fields=(density_field, "u_gradient")) + 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, eqn): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ rho_grad = self.density_field+"_gradient" - super().setup(eqn) - self.grad_density = eqn.fields(rho_grad) - self.gradu = eqn.fields("u_gradient") - - def compute(self, eqn): - """ - Compute and return the diagnostic field from the current state. + grad_density = state_fields(rho_grad) + gradu = state_fields("u_gradient") - Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. - """ denom = 0. - z_dim = eqn.domain.mesh.geometric_dimension() - 1 - u_dim = eqn.fields("u").ufl_shape[0] + z_dim = domain.mesh.geometric_dimension() - 1 + u_dim = state_fields("u").ufl_shape[0] for i in range(u_dim-1): - denom += self.gradu[i, z_dim]**2 - Nsq = self.factor*self.grad_density[z_dim] - self.field.interpolate(Nsq/denom) - return self.field + denom += gradu[i, z_dim]**2 + Nsq = self.factor*grad_density[z_dim] + self.expr = Nsq/denom + 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): @@ -568,70 +591,123 @@ class KineticEnergy(Energy): """Diagnostic kinetic energy density.""" name = "KineticEnergy" - def compute(self, eqn): + def __init__(self, space=None, method='interpolate'): """ - Compute and return the diagnostic field from the current state. - Args: - eqn (:class:`PrognosticEquation`): the model's 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'. + """ + super().__init__(space=space, method=method, required_fields=("u")) - Returns: - :class:`Function`: the diagnostic field. + def setup(self, domain, state_fields): """ - u = eqn.fields("u") - energy = self.kinetic(u) - return self.field.interpolate(energy) + 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 compute(self, eqn): + def __init__(self, space=None, method='interpolate'): """ - Compute and return the diagnostic field from the current state. - Args: - eqn (:class:`PrognosticEquation`): the model's 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'. + """ + super().__init__(space=space, method=method, required_fields=("D", "u")) - Returns: - :class:`Function`: the diagnostic field. + def setup(self, domain, state_fields): """ - u = eqn.fields("u") - D = eqn.fields("D") - energy = self.kinetic(u, D) - return self.field.interpolate(energy) + 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 compute(self, eqn): + def __init__(self, parameters, space=None, method='interpolate'): """ - Compute and return the diagnostic field from the current state. - Args: - eqn (:class:`PrognosticEquation`): the model's equation. + 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")) - Returns: - :class:`Function`: the diagnostic field. + def setup(self, domain, state_fields): """ - g = eqn.parameters.g - D = eqn.fields("D") - energy = 0.5*g*D**2 - return self.field.interpolate(energy) + 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"): + 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". - """ - super().__init__() + 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 @@ -640,38 +716,48 @@ def name(self): base_name = "SWPotentialEnstrophy" return "_from_".join((base_name, self.base_field_name)) - def compute(self, eqn): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ if self.base_field_name == "PotentialVorticity": - pv = eqn.fields("PotentialVorticity") - D = eqn.fields("D") - enstrophy = 0.5*pv**2*D + pv = state_fields("PotentialVorticity") + D = state_fields("D") + self.expr = 0.5*pv**2*D elif self.base_field_name == "RelativeVorticity": - zeta = eqn.fields("RelativeVorticity") - D = eqn.fields("D") - f = eqn.fields("coriolis") - enstrophy = 0.5*(zeta + f)**2/D + 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 = eqn.fields("AbsoluteVorticity") - D = eqn.fields("D") - enstrophy = 0.5*(zeta_abs)**2/D + zeta_abs = state_fields("AbsoluteVorticity") + D = state_fields("D") + self.expr = 0.5*(zeta_abs)**2/D else: - raise ValueError("Don't know how to compute enstrophy with base_field_name=%s; base_field_name should be %s %s or %s." % (self.base_field_name, "RelativeVorticity", "AbsoluteVorticity", "PotentialVorticity")) - return self.field.interpolate(enstrophy) + 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 compute(self, eqn): """ Compute and return the diagnostic field from the current state. @@ -690,161 +776,150 @@ def compute(self, eqn): class Exner(DiagnosticField): """The diagnostic Exner pressure field.""" - def __init__(self, reference=False): + 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. - """ - super(Exner, self).__init__() + 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 = "rhobar" - self.theta_name = "thetabar" + 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 "Exnerbar" + return "Exner_bar" else: return "Exner" - def setup(self, eqn): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - """ - if not self._initialised: - space = eqn.domain.spaces("CG1", "CG", 1) - super(Exner, self).setup(eqn, space=space) - - def compute(self, eqn): - """ - Compute and return the diagnostic field from the current state. - - Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - rho = eqn.fields(self.rho_name) - theta = eqn.fields(self.theta_name) - exner = thermodynamics.exner_pressure(eqn.parameters, rho, theta) - return self.field.interpolate(exner) + 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 Sum(DiagnosticField): """Base diagnostic for computing the sum of two fields.""" - def __init__(self, field1, field2): + def __init__(self, field_name1, field_name2): """ Args: - field1 (:class:`Function`): one field to be added. - field2 (:class:`Function`): the other field to be added. + 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__(required_fields=(field1, field2)) - self.field1 = field1 - self.field2 = field2 + 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.field1+"_plus_"+self.field2 + return self.field_name1+"_plus_"+self.field_name2 - def setup(self, eqn): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - if not self._initialised: - space = eqn.fields(self.field1).function_space() - super(Sum, self).setup(eqn, space=space) - - def compute(self, eqn): - """ - Compute and return the diagnostic field from the current state. - - Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. - """ - field1 = eqn.fields(self.field1) - field2 = eqn.fields(self.field2) - return self.field.assign(field1 + field2) + field1 = state_fields(self.field_name1) + field2 = state_fields(self.field_name2) + space = field1.function_space() + self.expr = field1 + field2 + super(Sum, self).setup(domain, state_fields, space=space) class Difference(DiagnosticField): """Base diagnostic for calculating the difference between two fields.""" - def __init__(self, field1, field2): + def __init__(self, field_name1, field_name2): """ Args: - field1 (:class:`Function`): the field to be subtracted from. - field2 (:class:`Function`): the field to be subtracted. + 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__(required_fields=(field1, field2)) - self.field1 = field1 - self.field2 = field2 + 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.field1+"_minus_"+self.field2 + return self.field_name1+"_minus_"+self.field_name2 - def setup(self, eqn): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - if not self._initialised: - space = eqn.fields(self.field1).function_space() - super(Difference, self).setup(eqn, space=space) - def compute(self, eqn): - """ - Compute and return the diagnostic field from the current state. - - Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. - """ - field1 = eqn.fields(self.field1) - field2 = eqn.fields(self.field2) - return self.field.assign(field1 - field2) + field1 = state_fields(self.field_name1) + field2 = state_fields(self.field_name2) + self.expr = field1 - field2 + space = field1.function_space() + super(Difference, self).setup(domain, state_fields, space=space) class SteadyStateError(Difference): """Base diagnostic for computing the steady-state error in a field.""" - def __init__(self, eqn, name): + def __init__(self, name): """ Args: - eqn (:class:`PrognosticEquation`): the model's equation. - name (str): name of the field to take the perturbation of. + 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)) + + def setup(self, domain, state_fields): """ - DiagnosticField.__init__(self) - self.field1 = name - self.field2 = name+'_init' - field1 = eqn.fields(name) - field2 = eqn.fields(self.field2, field1.function_space()) + 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. + """ + # Create and store initial field + field1 = state_fields(self.field_name1) + field2 = state_fields(self.field_name2, space=field1.function_space(), dump=False) + + # TODO: when checkpointing, the initial field should either be picked up + # or computed again (picking up can be easily specified if we change the line above) field2.assign(field1) + super(SteadyStateError, self).setup(domain, state_fields) + @property def name(self): """Gives the name of this diagnostic field.""" - return self.field1+"_error" + return self.field_name1+"_error" class Perturbation(Difference): @@ -854,466 +929,549 @@ def __init__(self, name): Args: name (str): name of the field to take the perturbation of. """ - self.field1 = name - self.field2 = name+'bar' - DiagnosticField.__init__(self, required_fields=(self.field1, self.field2)) + field_name1 = name + field_name2 = name+'_bar' + Difference.__init__(self, field_name1, field_name2) @property def name(self): """Gives the name of this diagnostic field.""" - return self.field1+"_perturbation" + return self.field_name1+"_perturbation" +# TODO: unify thermodynamic diagnostics class ThermodynamicDiagnostic(DiagnosticField): """Base thermodynamic diagnostic field, computing many common fields.""" - def setup(self, eqn): - """ - Sets up the :class:`Function` for the diagnostic field. - - Args: - eqn (:class:`PrognosticEquationSet`): the model's equation. - """ - if not self._initialised: - domain = eqn.domain - space = domain.spaces('theta') - h_deg = space.ufl_element().degree()[0] - v_deg = space.ufl_element().degree()[1]-1 - boundary_method = BoundaryMethod.extruded if (v_deg == 0 and h_deg == 0) else None - super().setup(eqn, space=space) - - # now let's attach all of our fields - self.u = eqn.fields("u") - self.rho = eqn.fields("rho") - self.theta = eqn.fields("theta") - self.rho_averaged = Function(space) - self.recoverer = Recoverer(self.rho, self.rho_averaged, boundary_method=boundary_method) - try: - self.r_v = eqn.fields("water_vapour") - except NotImplementedError: - self.r_v = Constant(0.0) - try: - self.r_c = eqn.fields("cloud_water") - except NotImplementedError: - self.r_c = Constant(0.0) - try: - self.rain = eqn.fields("rain") - except NotImplementedError: - self.rain = Constant(0.0) - - # now let's store the most common expressions - self.exner = thermodynamics.exner_pressure(eqn.parameters, self.rho_averaged, self.theta) - self.T = thermodynamics.T(eqn.parameters, self.theta, self.exner, r_v=self.r_v) - self.p = thermodynamics.p(eqn.parameters, self.exner) - self.r_l = self.r_c + self.rain - self.r_t = self.r_v + self.r_c + self.rain + 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 compute(self, eqn): + def _setup_thermodynamics(self, domain, state_fields): """ - Compute thermodynamic auxiliary fields commonly used by diagnostics. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - """ + 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 compute(self, eqn): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(eqn) - - return self.field.interpolate(thermodynamics.theta_e(eqn.parameters, self.T, self.p, self.r_v, self.r_t)) + 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 compute(self, eqn): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(eqn) - - return self.field.interpolate(thermodynamics.internal_energy(eqn.parameters, self.rho_averaged, self.T, r_v=self.r_v, r_l=self.r_l)) + 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, eqn): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().setup(eqn) - self.x = SpatialCoordinate(eqn.domain.mesh) - - def compute(self, eqn): - """ - Compute and return the diagnostic field from the current state. - - Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. - """ - super().compute(eqn) - - return self.field.interpolate(self.rho_averaged * (1 + self.r_t) * eqn.parameters.g * dot(self.x, eqn.domain.k)) + x = SpatialCoordinate(domain.mesh) + self.expr = self.rho_averaged * (1 + self.r_t) * self.parameters.g * dot(x, domain.k) + 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 compute(self, eqn): + 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): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(eqn) - - return self.field.project(0.5 * self.rho_averaged * (1 + self.r_t) * dot(self.u, self.u)) + u = state_fields('u') + 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 compute(self, eqn): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(eqn) - - return self.field.interpolate(thermodynamics.T_dew(eqn.parameters, self.p, self.r_v)) + 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 compute(self, eqn): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(eqn) - - return self.field.assign(self.T) + 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 compute(self, eqn): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(eqn) - - return self.field.interpolate(self.theta / (1 + self.r_v * eqn.parameters.R_v / eqn.parameters.R_d)) + 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 compute(self, eqn): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(eqn) - - return self.field.interpolate(thermodynamics.RH(eqn.parameters, self.r_v, self.T, self.p)) + 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 compute(self, eqn): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(eqn) - - return self.field.assign(self.p) + 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 compute(self, eqn): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(eqn) - - return self.field.assign(self.exner) + 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 setup(self, eqn): + 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: - eqn (:class:`PrognosticEquation`): the model's equation. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - if not self._initialised: - Vu = eqn.domain.spaces("HDiv") - space = FunctionSpace(eqn.domain.mesh, Vu.ufl_element()._elements[-1]) - super().setup(eqn, space=space) - rho = eqn.fields("rho") - rhobar = eqn.fields("rhobar") - theta = eqn.fields("theta") - thetabar = eqn.fields("thetabar") - exner = thermodynamics.exner_pressure(eqn.parameters, rho, theta) - exnerbar = thermodynamics.exner_pressure(eqn.parameters, rhobar, thetabar) - - cp = Constant(eqn.parameters.cp) - n = FacetNormal(eqn.domain.mesh) - - F = TrialFunction(space) - w = TestFunction(space) - 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 = [DirichletBC(space, 0.0, "bottom"), - DirichletBC(space, 0.0, "top")] - - imbalanceproblem = LinearVariationalProblem(a, L, self.field, bcs=bcs) - self.imbalance_solver = LinearVariationalSolver(imbalanceproblem) + 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) - def compute(self, eqn): - """ - Compute and return the diagnostic field from the current state. + cp = Constant(self.parameters.cp) + n = FacetNormal(domain.mesh) - Args: - eqn (:class:`PrognosticEquation`): the model's equation. + # 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) - Returns: - :class:`Function`: the diagnostic field. + 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() - return self.field[1] + super().compute() class Precipitation(DiagnosticField): """The total precipitation falling through the domain's bottom surface.""" name = "Precipitation" - def setup(self, eqn): - """ - Sets up the :class:`Function` for the diagnostic field. + def __init__(self): + self.solve_implemented = True + required_fields = ('rain', 'rainfall_velocity', 'rho') + super().__init__(method='solve', required_fields=required_fields) - Args: - eqn (:class:`PrognosticEquation`): the model's equation. + def setup(self, domain, state_fields): """ - if not self._initialised: - space = eqn.domain.spaces("DG0", "DG", 0) - super().setup(eqn, space=space) - - rain = eqn.fields('rain') - rho = eqn.fields('rho') - v = eqn.fields('rainfall_velocity') - self.phi = TestFunction(space) - flux = TrialFunction(space) - n = FacetNormal(eqn.domain.mesh) - un = 0.5 * (dot(v, n) + abs(dot(v, n))) - self.flux = Function(space) - - a = self.phi * flux * dx - L = self.phi * rain * un * rho - if space.extruded: - L = L * (ds_b + ds_t + ds_v) - else: - L = L * ds - - # setup solver - problem = LinearVariationalProblem(a, L, self.flux) - self.solver = LinearVariationalSolver(problem) - - def compute(self, eqn): - """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. - """ + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + space = domain.spaces("DG0", "DG", 0) + assert space.extruded, 'Cannot compute precipitation on a non-extruded mesh' + rain = state_fields('rain') + rho = state_fields('rho') + v = state_fields('rainfall_velocity') + # Set up problem + self.phi = TestFunction(space) + flux = TrialFunction(space) + n = FacetNormal(domain.mesh) + un = 0.5 * (dot(v, n) + abs(dot(v, n))) + self.flux = Function(space) + + a = self.phi * flux * dx + L = self.phi * rain * un * rho * (ds_b + ds_t + ds_v) + + # setup solver + problem = LinearVariationalProblem(a, L, self.flux) + self.solver = LinearVariationalSolver(problem) + self.space = space + self.field = state_fields(self.name, space=space, dump=True, pickup=False) + # TODO: might we want to pick up this field? Otherwise initialise to zero + self.field.assign(0.0) + + def compute(self): + """Compute the diagnostic field from the current state.""" self.solver.solve() self.field.assign(self.field + assemble(self.flux * self.phi * dx)) - return self.field class Vorticity(DiagnosticField): """Base diagnostic field class for shallow-water vorticity variables.""" - def setup(self, eqn, vorticity_type=None): + def setup(self, domain, state_fields, vorticity_type=None): """ Sets up the :class:`Function` for the diagnostic field. Args: - eqn (:class:`PrognosticEquation`): the model's equation. + 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. """ - if not self._initialised: - vorticity_types = ["relative", "absolute", "potential"] - if vorticity_type not in vorticity_types: - raise ValueError("vorticity type must be one of %s, not %s" % (vorticity_types, vorticity_type)) - try: - space = eqn.domain.spaces("CG") - except ValueError: - dgspace = eqn.domain.spaces("DG") - cg_degree = dgspace.ufl_element().degree() + 2 - space = FunctionSpace(eqn.domain.mesh, "CG", cg_degree) - super().setup(eqn, space=space) - u = eqn.fields("u") + + 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}") + try: + space = domain.spaces("CG") + except ValueError: + dgspace = domain.spaces("DG") + # TODO: should this be degree + 1? + cg_degree = dgspace.ufl_element().degree() + 2 + space = FunctionSpace(domain.mesh, "CG", cg_degree) + + 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": - D = eqn.fields("D") a = q*gamma*D*dx else: a = q*gamma*dx - L = (- inner(eqn.domain.perp(grad(gamma)), u))*dx + L = (- inner(domain.perp(grad(gamma)), u))*dx if vorticity_type != "relative": - f = eqn.fields("coriolis") + f = state_fields("coriolis") L += gamma*f*dx problem = LinearVariationalProblem(a, L, self.field) - self.solver = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"}) - - def compute(self, eqn): - """ - Compute and return the diagnostic field from the current state. - - Args: - eqn (:class:`PrognosticEquation`): the model's equation. - - Returns: - :class:`Function`: the diagnostic field. - """ - self.solver.solve() - return 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 setup(self, eqn): + 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: - eqn (:class:`PrognosticEquation`): the model's equation. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().setup(eqn, vorticity_type="potential") + super().setup(domain, state_fields, vorticity_type="potential") class AbsoluteVorticity(Vorticity): u"""Diagnostic field for absolute vorticity, ζ=∇×(u+f)""" name = "AbsoluteVorticity" - def setup(self, eqn): + 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: - eqn (:class:`PrognosticEquation`): the model's equation. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().setup(eqn, vorticity_type="absolute") + super().setup(domain, state_fields, vorticity_type="absolute") class RelativeVorticity(Vorticity): u"""Diagnostic field for relative vorticity, ζ=∇×u""" name = "RelativeVorticity" - def setup(self, eqn): + 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: - eqn (:class:`PrognosticEquation`): the model's equation. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().setup(eqn, vorticity_type="relative") + super().setup(domain, state_fields, vorticity_type="relative") diff --git a/gusto/domain.py b/gusto/domain.py index 6a967d125..301b4c81e 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -8,6 +8,7 @@ from firedrake import (Constant, SpatialCoordinate, sqrt, CellNormal, cross, as_vector, inner, interpolate) + class Domain(object): """ The Domain holds the model's mesh and its compatible function spaces. @@ -64,7 +65,7 @@ def __init__(self, mesh, dt, family, degree=None, # Figure out if we're on a sphere # TODO: could we run on other domains that could confuse this? - if hasattr(mesh, "_base_mesh"): + if hasattr(mesh, "_base_mesh") and hasattr(mesh._base_mesh, 'geometric_dimension'): self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3 and mesh._base_mesh.topological_dimension() == 2) else: self.on_sphere = (mesh.geometric_dimension() == 3 and mesh.topological_dimension() == 2) diff --git a/gusto/equations.py b/gusto/equations.py index b7a8f60cc..9f10e20a7 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -6,7 +6,7 @@ TrialFunction, FacetNormal, jump, avg, dS_v, DirichletBC, conditional, SpatialCoordinate, split, Constant, action) -from gusto.fields import StateFields +from gusto.fields import PrescribedFields from gusto.fml.form_manipulation_labelling import Term, all_terms, identity, drop from gusto.labels import (subject, time_derivative, transport, prognostic, transporting_velocity, replace_subject, linearisation, @@ -41,18 +41,15 @@ def __init__(self, domain, function_space, field_name): self.domain = domain self.function_space = function_space + self.X = Function(function_space) self.field_name = field_name - self.fields = StateFields() # TODO: should there be an argument here? self.bcs = {} + self.prescribed_fields = PrescribedFields() if len(function_space) > 1: assert hasattr(self, "field_names") - self.fields(field_name, function_space, - subfield_names=self.field_names, pickup=True) for fname in self.field_names: self.bcs[fname] = [] - else: - self.fields(field_name, function_space) self.bcs[field_name] = [] @@ -75,12 +72,12 @@ def __init__(self, domain, function_space, field_name, Vu=None, **kwargs): """ super().__init__(domain, function_space, field_name) - if not hasattr(self.fields, "u"): - if Vu is not None: - V = domain.spaces("HDiv", V=Vu) - else: - V = domain.spaces("HDiv") - self.fields("u", V) + if Vu is not None: + V = domain.spaces("HDiv", V=Vu) + else: + V = domain.spaces("HDiv") + self.prescribed_fields("u", V) + test = TestFunction(function_space) q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) @@ -108,12 +105,12 @@ def __init__(self, domain, function_space, field_name, Vu=None, **kwargs): """ super().__init__(domain, function_space, field_name) - if not hasattr(self.fields, "u"): - if Vu is not None: - V = domain.spaces("HDiv", V=Vu) - else: - V = domain.spaces("HDiv") - self.fields("u", V) + if Vu is not None: + V = domain.spaces("HDiv", V=Vu) + else: + V = domain.spaces("HDiv") + self.prescribed_fields("u", V) + test = TestFunction(function_space) q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) @@ -172,12 +169,12 @@ def __init__(self, domain, function_space, field_name, Vu=None, super().__init__(domain, function_space, field_name) - if not hasattr(self.fields, "u"): - if Vu is not None: - V = domain.spaces("HDiv", V=Vu) - else: - V = domain.spaces("HDiv") - self.fields("u", V) + if Vu is not None: + V = domain.spaces("HDiv", V=Vu) + else: + V = domain.spaces("HDiv") + self.prescribed_fields("u", V) + test = TestFunction(function_space) q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) @@ -220,7 +217,6 @@ def __init__(self, field_names, domain, linearisation_map=None, self.field_names = field_names self.active_tracers = active_tracers self.linearisation_map = lambda t: False if linearisation_map is None else linearisation_map(t) - self.reference_profiles_initialised = False # Build finite element spaces # TODO: this implies order of spaces matches order of variables @@ -243,7 +239,6 @@ def __init__(self, field_names, domain, linearisation_map=None, # Set up test functions, trials and prognostics self.tests = TestFunctions(W) self.trials = TrialFunction(W) - self.X = Function(W) self.X_ref = Function(W) # Set up no-normal-flow boundary conditions @@ -344,36 +339,6 @@ def linearise_equation_set(self): self.residual = self.residual.label_map( all_terms, replace_trial_function(self.X)) - # ======================================================================== # - # Reference Profile Routines - # ======================================================================== # - - def set_reference_profiles(self, reference_profiles): - """ - Initialise the equation'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. - """ - # TODO: come back and consider all aspects of this - for name, profile in reference_profiles: - if name+'bar' in self.fields: - # For reference profiles already added to state, allow - # interpolation from expressions - ref = self.fields(name+'bar') - elif isinstance(profile, Function): - # Need to add reference profile to state so profile must be - # a Function - ref = self.fields(name+'bar', space=profile.function_space(), dump=False) - else: - raise ValueError(f'When initialising reference profile {name}' - + ' the passed profile must be a Function') - ref.interpolate(profile) - - self.reference_profiles_initialised = True - # ======================================================================== # # Boundary Condition Routines # ======================================================================== # @@ -534,12 +499,11 @@ def __init__(self, domain, function_space, field_name, Vu=None, full_field_name = "_".join(self.field_names) PrognosticEquation.__init__(self, domain, W, full_field_name) - if not hasattr(self.fields, "u"): - if Vu is not None: - V = domain.spaces("HDiv", V=Vu) - else: - V = domain.spaces("HDiv") - self.fields("u", V) + if Vu is not None: + V = domain.spaces("HDiv", V=Vu) + else: + V = domain.spaces("HDiv") + self.prescribed_fields("u", V) self.tests = TestFunctions(W) self.X = Function(W) @@ -679,10 +643,14 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # -------------------------------------------------------------------- # # Extra Terms (Coriolis and Topography) # -------------------------------------------------------------------- # + # 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.fields("coriolis", space=V) - f.interpolate(fexpr) + # TODO: link this to state fields + 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 @@ -692,20 +660,13 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, residual += coriolis_form if bexpr is not None: - b = self.fields("topography", domain.spaces("DG")) - b.interpolate(bexpr) + b = self.prescribed_fields("topography", domain.spaces("DG")).interpolate(bexpr) topography_form = subject(prognostic(-g*div(w)*b*dx, "u"), self.X) residual += topography_form # -------------------------------------------------------------------- # # Linearise equations # -------------------------------------------------------------------- # - u_ref, D_ref = self.X_ref.split()[0:2] - # Linearise about D = H - # TODO: add interface to update linearisation state - D_ref.assign(Constant(H)) - u_ref.assign(Constant(0.0)) - # Add linearisations to equations self.residual = self.generate_linear_terms(residual, self.linearisation_map) @@ -858,8 +819,7 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, w, phi, gamma = self.tests[0:3] u, rho, theta = split(self.X)[0:3] u_trial = split(self.trials)[0] - rhobar = self.fields("rhobar", space=domain.spaces("DG"), dump=False) - thetabar = self.fields("thetabar", space=domain.spaces("theta"), dump=False) + _, 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) @@ -895,7 +855,7 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, rho_adv = prognostic(continuity_form(domain, phi, rho), "rho") # Transport term needs special linearisation if self.linearisation_map(rho_adv.terms[0]): - linear_rho_adv = linear_continuity_form(domain, phi, rhobar).label_map( + linear_rho_adv = linear_continuity_form(domain, phi, rho_bar).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u_trial}), t.labels)) @@ -905,7 +865,7 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, theta_adv = prognostic(advection_form(domain, gamma, theta), "theta") # Transport term needs special linearisation if self.linearisation_map(theta_adv.terms[0]): - linear_theta_adv = linear_advection_form(domain, gamma, thetabar).label_map( + linear_theta_adv = linear_advection_form(domain, gamma, theta_bar).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u_trial}), t.labels)) @@ -1203,8 +1163,7 @@ def __init__(self, domain, parameters, Omega=None, w, phi, gamma = self.tests[0:3] u, p, b = split(self.X) u_trial = split(self.trials)[0] - bbar = self.fields("bbar", space=domain.spaces("theta"), dump=False) - bbar = self.fields("pbar", space=domain.spaces("DG"), dump=False) + b_bar = split(self.X_ref)[2] # -------------------------------------------------------------------- # # Time Derivative Terms @@ -1236,7 +1195,7 @@ def __init__(self, domain, parameters, Omega=None, # Buoyancy transport b_adv = prognostic(advection_form(domain, gamma, b), "b") if self.linearisation_map(b_adv.terms[0]): - linear_b_adv = linear_advection_form(domain, gamma, bbar).label_map( + linear_b_adv = linear_advection_form(domain, gamma, b_bar).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u_trial}), t.labels)) diff --git a/gusto/fields.py b/gusto/fields.py index 63e8b9bc6..f3e674f73 100644 --- a/gusto/fields.py +++ b/gusto/fields.py @@ -1,6 +1,6 @@ -from firedrake import Function +from firedrake import Function, MixedElement, functionspaceimpl -__all__ = ["TimeLevelFields", "StateFields"] +__all__ = ["PrescribedFields", "TimeLevelFields", "StateFields"] class Fields(object): @@ -17,8 +17,7 @@ def __init__(self, equation): def add_field(self, name, space, subfield_names=None): """ - Adds a new field to the :class:`FieldCreator`. - + Adds a new field to the :class:`Fields` object. Args: name (str): the name of the prognostic variable. space (:class:`FunctionSpace`): the space to create the field in. @@ -42,15 +41,46 @@ def add_field(self, name, space, subfield_names=None): def __call__(self, name): """ - Returns a specified field from the :class:`FieldCreator`. + Returns a specified field from the :class:`Fields` object. + Args: + name (str): the name of the field. + Returns: + :class:`Function`: the desired field. + """ + return getattr(self, name) + + def __iter__(self): + """Returns an iterable of the contained fields.""" + return iter(self.fields) + + +class PrescribedFields(Fields): + """Object to hold and create a specified set of prescribed fields.""" + def __init__(self): + self.fields = [] + + def __call__(self, name, space=None): + """ + Returns a specified field from the :class:`PrescribedFields`. If a named + field does not yet exist in the :class:`PrescribedFields` object, then + the space argument must be specified so that it can be created and added + to the object. Args: name (str): the name of the field. + space (:class:`FunctionSpace`, optional): the function space to + create the field in. Defaults to None. Returns: :class:`Function`: the desired field. """ - return getattr(self, name) + if hasattr(self, name): + # Field already exists in object, so return it + return getattr(self, name) + else: + # Create field + self.add_field(name, space) + return getattr(self, name) def __iter__(self): """Returns an iterable of the contained fields.""" @@ -58,20 +88,50 @@ def __iter__(self): class StateFields(Fields): - """Creates the prognostic fields for the model's equation.""" + """ + Container for all of the model's fields. + + The `StateFields` are a container for all the fields to be used by a time + stepper. In the case of the prognostic fields, these are pointers to the + time steppers' fields at the (n+1) time level. Prescribed fields are + pointers to the respective equation sets, while diagnostic fields are + created here. + """ - def __init__(self, *fields_to_dump): + def __init__(self, prognostic_fields, prescribed_fields, *fields_to_dump): """ Args: + prognostic_fields (:class:`Fields`): the (n+1) time level fields. + prescribed_fields (iter): an iterable of (name, function_space) + tuples, that are used to create the prescribed fields. *fields_to_dump (str): the names of fields to be dumped. """ self.fields = [] - self.output_specified = len(fields_to_dump) > 0 + output_specified = len(fields_to_dump) > 0 self.to_dump = set((fields_to_dump)) self.to_pickup = set(()) + self._field_types = [] + self._field_names = [] + + # Add pointers to prognostic fields + for field in prognostic_fields.fields: + # Don't add the mixed field + if type(field.ufl_element()) is not MixedElement: + # If fields_to_dump not specified, dump by default + to_dump = field.name() in fields_to_dump or not output_specified + self.__call__(field.name(), field=field, dump=to_dump, + pickup=True, field_type="prognostic") + else: + self.__call__(field.name(), field=field, dump=False, + pickup=False, field_type="prognostic") + + for field in prescribed_fields.fields: + to_dump = field.name() in fields_to_dump + self.__call__(field.name(), field=field, dump=to_dump, + pickup=True, field_type="prescribed") - def __call__(self, name, space=None, subfield_names=None, dump=True, - pickup=False): + def __call__(self, name, field=None, space=None, dump=True, pickup=False, + field_type=None): """ Returns a field from or adds a field to the :class:`StateFields`. @@ -79,13 +139,16 @@ def __call__(self, name, space=None, subfield_names=None, dump=True, the optional arguments must be specified so that it can be created and added to the :class:`StateFields`. + If "field" is specified, then the pointer to the field is added to the + :class:`StateFields` object. If "space" is specified, then the field + itself is created. + Args: name (str): name of the field to be returned/added. + field (:class:`Function`, optional): an existing field to be added + to the :class:`StateFields` object. space (:class:`FunctionSpace`, optional): the function space to create the field in. Defaults to None. - subfield_names (list, optional): a list of names of the constituent - prognostic variables to be created, if the provided space is - actually a :class:`MixedFunctionSpace`. Defaults to None. dump (bool, optional): whether the created field should be outputted. Defaults to True. pickup (bool, optional): whether the created field should be picked @@ -94,22 +157,68 @@ def __call__(self, name, space=None, subfield_names=None, dump=True, Returns: :class:`Function`: the specified field. """ - try: + if hasattr(self, name): + # Field already exists in object, so return it return getattr(self, name) - except AttributeError: - self.add_field(name, space, subfield_names) + else: + # Field does not yet exist in StateFields + if field is None and space is None: + raise ValueError(f'Field {name} does not exist in StateFields. ' + + 'Either field or space argument must be ' + + 'specified to add this field to StateFields') + elif field is not None and space is not None: + raise ValueError('Cannot specify both field and space to StateFields') + + if field is not None: + # Field pointer, so just add existing field to StateFields + assert isinstance(field, Function), \ + f'field argument for creating field {name} must be a Function, not {type(field)}' + setattr(self, name, field) + self.fields.append(field) + else: + # Create field + assert isinstance(space, functionspaceimpl.WithGeometry), \ + f'space argument for creating field {name} must be FunctionSpace, not {type(space)}' + self.add_field(name, space) + if dump: - if subfield_names is not None: - self.to_dump.update(subfield_names) - else: - self.to_dump.add(name) + self.to_dump.add(name) if pickup: - if subfield_names is not None: - self.to_pickup.update(subfield_names) + self.to_pickup.add(name) + + # Work out field type + if field_type is None: + # Prognostics can only be specified through __init__ + if pickup: + field_type = "prescribed" + elif dump: + field_type = "diagnostic" else: - self.to_pickup.add(name) + field_type = "derived" + else: + permitted_types = ["prognostic", "prescribed", "diagnostic", "derived"] + assert field_type in permitted_types, \ + f'field_type {field_type} not in permitted types {permitted_types}' + self._field_types.append(field_type) + self._field_names.append(name) + return getattr(self, name) + def field_type(self, field_name): + """ + Returns the type (e.g. prognostic/diagnostic) of a field held in the + :class:`StateFields`. + + Args: + field_name (str): name of the field to return the type of. + + Returns: + str: a string describing the type (e.g. prognostic) of the field. + """ + assert hasattr(self, field_name), f'StateFields has no field {field_name}' + idx = self._field_names.index(field_name) + return self._field_types[idx] + class TimeLevelFields(object): """Creates the fields required in the :class:`Timestepper` object.""" @@ -149,17 +258,15 @@ def add_fields(self, equation, levels=None): except AttributeError: setattr(self, level, Fields(equation)) - def initialise(self, equation): - # TODO: should this be IO? + def initialise(self, state_fields): """ Initialises the time fields from those currently in the equation. Args: - equation (:class:`PrognosticEquation`): the model's prognostic - equation object. + state_fields (:class:`StateFields`): the model's field container. """ for field in self.n: - field.assign(equation.fields(field.name())) + field.assign(state_fields(field.name())) self.np1(field.name()).assign(field) def update(self): diff --git a/gusto/function_spaces.py b/gusto/function_spaces.py index 5ff3eab83..7adf0cdd6 100644 --- a/gusto/function_spaces.py +++ b/gusto/function_spaces.py @@ -13,6 +13,7 @@ # the creation of the de Rham complex spaces? # TODO: how do we create HCurl spaces if we want them? + class Spaces(object): """Object to create and hold the model's finite element spaces.""" def __init__(self, mesh): @@ -58,13 +59,12 @@ def __call__(self, name, family=None, degree=None, :class:`FunctionSpace`: the desired function space. """ - try: - # First attempt to return the space based on the name, if it exists + if hasattr(self, name) and family is None and V is None: + # We have requested a space that should already have been created return getattr(self, name) - except AttributeError: - - # Space does not exist in creator + else: + # Space does not exist in creator or needs overwriting if V is not None: # The space itself has been provided (to add it to the creator) value = V diff --git a/gusto/initialisation_tools.py b/gusto/initialisation_tools.py index 65adeb661..5807686a1 100644 --- a/gusto/initialisation_tools.py +++ b/gusto/initialisation_tools.py @@ -279,8 +279,9 @@ def remove_initial_w(u): u.assign(uin) -def saturated_hydrostatic_balance(equation, theta_e, mr_t, exner0=None, - top=False, exner_boundary=Constant(1.0), +def saturated_hydrostatic_balance(equation, state_fields, theta_e, mr_t, + exner0=None, top=False, + exner_boundary=Constant(1.0), max_outer_solve_count=40, max_theta_solve_count=5, max_inner_solve_count=3): @@ -300,6 +301,7 @@ def saturated_hydrostatic_balance(equation, theta_e, mr_t, exner0=None, Args: equation (:class:`PrognosticEquation`): the model's equation object. + state_fields (:class:`StateFields`): the model's field container. theta_e (:class:`ufl.Expr`): expression for the desired wet equivalent potential temperature field. mr_t (:class:`ufl.Expr`): expression for the total moisture content. @@ -326,9 +328,9 @@ def saturated_hydrostatic_balance(equation, theta_e, mr_t, exner0=None, number of iterations. """ - theta0 = equation.fields('theta') - rho0 = equation.fields('rho') - mr_v0 = equation.fields('water_vapour') + theta0 = state_fields('theta') + rho0 = state_fields('rho') + mr_v0 = state_fields('water_vapour') # Calculate hydrostatic exner pressure domain = equation.domain @@ -409,8 +411,9 @@ def saturated_hydrostatic_balance(equation, theta_e, mr_t, exner0=None, mr_t=mr_t, solve_for_rho=True) -def unsaturated_hydrostatic_balance(equation, theta_d, H, exner0=None, - top=False, exner_boundary=Constant(1.0), +def unsaturated_hydrostatic_balance(equation, state_fields, theta_d, H, + exner0=None, top=False, + exner_boundary=Constant(1.0), max_outer_solve_count=40, max_inner_solve_count=20): """ @@ -428,6 +431,7 @@ def unsaturated_hydrostatic_balance(equation, theta_d, H, exner0=None, Args: equation (:class:`PrognosticEquation`): the model's equation object. + state_fields (:class:`StateFields`): the model's field container. theta_d (:class:`ufl.Expr`): the specified dry potential temperature field. H (:class:`ufl.Expr`): the specified relative humidity field. @@ -452,9 +456,9 @@ def unsaturated_hydrostatic_balance(equation, theta_d, H, exner0=None, number of iterations. """ - theta0 = equation.fields('theta') - rho0 = equation.fields('rho') - mr_v0 = equation.fields('water_vapour') + theta0 = state_fields('theta') + rho0 = state_fields('rho') + mr_v0 = state_fields('water_vapour') # Calculate hydrostatic exner pressure domain = equation.domain diff --git a/gusto/io.py b/gusto/io.py index d6212c5ec..9d649baa4 100644 --- a/gusto/io.py +++ b/gusto/io.py @@ -5,10 +5,10 @@ from netCDF4 import Dataset import sys import time -from gusto.diagnostics import Diagnostics, Perturbation, SteadyStateError +from gusto.diagnostics import Diagnostics from firedrake import (FiniteElement, TensorProductElement, VectorFunctionSpace, interval, Function, Mesh, functionspaceimpl, File, - Constant, op2, DumbCheckpoint, FILE_CREATE, FILE_READ) + op2, DumbCheckpoint, FILE_CREATE, FILE_READ) import numpy as np from gusto.configuration import logger, set_log_handler @@ -132,17 +132,17 @@ def __init__(self, filename, diagnostics, description, comm, create=True): for diagnostic in diagnostics.available_diagnostics: group.createVariable(diagnostic, np.float64, ("time", )) - def dump(self, equation, t): + def dump(self, state_fields, t): """ Output the global diagnostics. - equation (:class:`PrognosticEquation`): the model's equation object. + state_fields (:class:`StateFields`): the model's field container. t (float): simulation time at which the output occurs. """ diagnostics = [] for fname in self.diagnostics.fields: - field = equation.fields(fname) + field = state_fields(fname) for dname in self.diagnostics.available_diagnostics: diagnostic = getattr(self.diagnostics, dname) diagnostics.append((fname, dname, diagnostic(field))) @@ -160,16 +160,11 @@ def dump(self, equation, t): class IO(object): """Controls the model's input, output and diagnostics.""" - def __init__(self, domain, equation, - output=None, - parameters=None, - diagnostics=None, - diagnostic_fields=None): + def __init__(self, domain, output, diagnostics=None, diagnostic_fields=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the mesh and the compatible function spaces. - equation (:class:`PrognosticEquation`): the prognostic equation. output (:class:`OutputParameters`, optional): holds and describes the options for outputting. Defaults to None. diagnostics (:class:`Diagnostics`, optional): object holding and @@ -181,16 +176,9 @@ def __init__(self, domain, equation, RuntimeError: if no output is provided. TypeError: if `dt` cannot be cast to a :class:`Constant`. """ - - self.equation = equation + self.domain = domain self.mesh = domain.mesh - - if output is None: - # TODO: maybe this shouldn't be an optional argument then? - raise RuntimeError("You must provide a directory name for dumping results") - else: - self.output = output - self.parameters = parameters + self.output = output if diagnostics is not None: self.diagnostics = diagnostics @@ -201,13 +189,6 @@ def __init__(self, domain, equation, else: self.diagnostic_fields = [] - # TODO: quick way of ensuring that diagnostics are registered - if hasattr(equation, "field_names"): - for fname in equation.field_names: - self.diagnostics.register(fname) - else: - self.diagnostics.register(equation.field_name) - if self.output.dumplist is None: self.output.dumplist = [] @@ -218,33 +199,45 @@ def __init__(self, domain, equation, # setup logger logger.setLevel(output.log_level) set_log_handler(self.mesh.comm) - if parameters is not None: - logger.info("Physical parameters that take non-default values:") - logger.info(", ".join("%s: %s" % (k, float(v)) for (k, v) in vars(parameters).items())) - # Constant to hold current time - self.t = Constant(0.0) + def log_parameters(self, equation): + """ + Logs an equation's physical parameters that take non-default values. + + Args: + equation (:class:`PrognosticEquation`): the model's equation which + contains any physical parameters used in the model run. + """ + if hasattr(equation, 'parameters') and equation.parameters is not None: + logger.info("Physical parameters that take non-default values:") + logger.info(", ".join("%s: %s" % (k, float(v)) for (k, v) in vars(equation.parameters).items())) - def setup_diagnostics(self): - """Concatenates the various types of diagnostic field.""" - for name in self.output.perturbation_fields: - f = Perturbation(name) - self.diagnostic_fields.append(f) + def setup_diagnostics(self, state_fields): + """ + Prepares the I/O for computing the model's global diagnostics and + diagnostic fields. - for name in self.output.steady_state_error_fields: - f = SteadyStateError(self.equation, name) - self.diagnostic_fields.append(f) + Args: + state_fields (:class:`StateFields`): the model's field container. + """ - fields = set([f.name() for f in self.equation.fields]) - field_deps = [(d, sorted(set(d.required_fields).difference(fields),)) for d in self.diagnostic_fields] + diagnostic_names = [diagnostic.name for diagnostic in self.diagnostic_fields] + non_diagnostics = [fname for fname in state_fields._field_names if state_fields.field_type(fname) != "diagnostic" or fname not in diagnostic_names] + # Filter out non-diagnostic fields + field_deps = [(d, sorted(set(d.required_fields).difference(non_diagnostics),)) for d in self.diagnostic_fields] schedule = topo_sort(field_deps) self.diagnostic_fields = schedule for diagnostic in self.diagnostic_fields: - # TODO: for diagnostics to see equation and IO, change the setup here - diagnostic.setup(self.equation) + diagnostic.setup(self.domain, state_fields) self.diagnostics.register(diagnostic.name) - def setup_dump(self, t, tmax, pickup=False): + # Register fields for global diagnostics + # TODO: it should be possible to specify which global diagnostics are used + for fname in state_fields._field_names: + if fname in state_fields.to_dump: + self.diagnostics.register(fname) + + def setup_dump(self, state_fields, t, tmax, pickup=False): """ Sets up a series of things used for outputting. @@ -254,6 +247,7 @@ def setup_dump(self, t, tmax, pickup=False): checkpointing file. Args: + state_fields (:class:`StateFields`): the model's field container. t (float): the current model time. tmax (float): the end time of the model's simulation. pickup (bool, optional): whether to pick up the model's initial @@ -287,7 +281,7 @@ def setup_dump(self, t, tmax, pickup=False): comm=self.mesh.comm) # make list of fields to dump - self.to_dump = [f for f in self.equation.fields if f.name() in self.equation.fields.to_dump] + self.to_dump = [f for f in state_fields.fields if f.name() in state_fields.to_dump] # make dump counter self.dumpcount = itertools.count() @@ -304,7 +298,7 @@ def setup_dump(self, t, tmax, pickup=False): # make functions on latlon mesh, as specified by dumplist_latlon self.to_dump_latlon = [] for name in self.output.dumplist_latlon: - f = self.equation.fields(name) + f = state_fields(name) field = Function( functionspaceimpl.WithGeometry.create( f.function_space(), mesh_ll), @@ -324,11 +318,11 @@ def setup_dump(self, t, tmax, pickup=False): if len(self.output.point_data) > 0: # set up point data output pointdata_filename = self.dumpdir+"/point_data.nc" - ndt = int(tmax/float(self.dt)) + ndt = int(tmax/float(self.domain.dt)) self.pointdata_output = PointDataOutput(pointdata_filename, ndt, self.output.point_data, self.output.dirname, - self.equation.fields, + state_fields, self.mesh.comm, self.output.tolerance, create=not pickup) @@ -349,21 +343,21 @@ def setup_dump(self, t, tmax, pickup=False): mode=FILE_CREATE) # make list of fields to pickup (this doesn't include # diagnostic fields) - self.to_pickup = [f for f in self.equation.fields if f.name() in self.equation.fields.to_pickup] + self.to_pickup = [state_fields(f) for f in state_fields.to_pickup] # if we want to checkpoint then make a checkpoint counter if self.output.checkpoint: self.chkptcount = itertools.count() # dump initial fields - self.dump(t) + self.dump(state_fields, t) - def pickup_from_checkpoint(self): + def pickup_from_checkpoint(self, state_fields): """Picks up the model's variables from a checkpoint file.""" # TODO: this duplicates some code from setup_dump. Can this be avoided? # It is because we don't know if we are picking up or setting dump first if self.to_pickup is None: - self.to_pickup = [f for f in self.equation.fields if f.name() in self.equation.fields.to_pickup] + self.to_pickup = [state_fields(f) for f in state_fields.to_pickup] # Set dumpdir if has not been done already if self.dumpdir is None: self.dumpdir = path.join("results", self.output.dirname) @@ -386,7 +380,7 @@ def pickup_from_checkpoint(self): return t - def dump(self, t): + def dump(self, state_fields, t): """ Dumps all of the required model output. @@ -395,6 +389,7 @@ def dump(self, t): a checkpoint file if specified. Args: + state_fields (:class:`StateFields`): the model's field container. t (float): the simulation's current time. """ output = self.output @@ -402,15 +397,15 @@ def dump(self, t): # Diagnostics: # Compute diagnostic fields for field in self.diagnostic_fields: - field(self.equation) + field.compute() if output.dump_diagnostics: # Output diagnostic data - self.diagnostic_output.dump(self.equation, t) + self.diagnostic_output.dump(state_fields, t) if len(output.point_data) > 0 and (next(self.pddumpcount) % output.pddumpfreq) == 0: # Output pointwise data - self.pointdata_output.dump(self.equation.fields, t) + self.pointdata_output.dump(state_fields, t) # Dump all the fields to the checkpointing file (backup version) if output.checkpoint and (next(self.chkptcount) % output.chkptfreq) == 0: @@ -426,20 +421,6 @@ def dump(self, t): if len(output.dumplist_latlon) > 0: self.dumpfile_ll.write(*self.to_dump_latlon) - def initialise(self, initial_conditions): - """ - Initialise the state's prognostic variables. - - Args: - initial_conditions (list): an iterable of pairs: (field_name, expr), - where 'field_name' is the string giving the name of the - prognostic field and expr is the :class:`ufl.Expr` whose value - is used to set the initial field. - """ - for name, ic in initial_conditions: - f_init = getattr(self.equation.fields, name) - f_init.assign(ic) - f_init.rename(name) def get_latlon_mesh(mesh): """ @@ -522,6 +503,7 @@ def topo_sort(field_deps): name2field = dict((f.name, f) for f, _ in field_deps) # map node: (input_deps, output_deps) graph = dict((f.name, (list(deps), [])) for f, deps in field_deps) + roots = [] for f, input_deps in field_deps: if len(input_deps) == 0: diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index e4704d1dc..547fe99d7 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -14,6 +14,7 @@ from firedrake.petsc import flatten_parameters from pyop2.profiling import timed_function, timed_region +from gusto.active_tracers import TracerVariableType from gusto.configuration import logger, DEBUG from gusto.labels import linearisation, time_derivative, hydrostatic from gusto import thermodynamics @@ -123,7 +124,7 @@ class CompressibleSolver(TimesteppingSolver): def __init__(self, equations, alpha=0.5, quadrature_degree=None, solver_parameters=None, - overwrite_solver_parameters=False, moisture=None): + overwrite_solver_parameters=False): """ Args: equations (:class:`PrognosticEquation`): the model's equation. @@ -139,11 +140,8 @@ def __init__(self, equations, alpha=0.5, `solver_parameters` that have been passed in. If False then update the default parameters with the `solver_parameters` passed in. Defaults to False. - moisture (list, optional): list of names of moisture fields. - Defaults to None. """ self.equations = equations - self.moisture = moisture if quadrature_degree is not None: self.quadrature_degree = quadrature_degree @@ -198,8 +196,7 @@ def _setup_solver(self): n = FacetNormal(equations.domain.mesh) # Get background fields - thetabar = equations.fields("thetabar") - rhobar = equations.fields("rhobar") + _, rhobar, thetabar = split(equations.X_ref)[0:3] exnerbar = thermodynamics.exner_pressure(equations.parameters, rhobar, thetabar) exnerbar_rho = thermodynamics.dexner_drho(equations.parameters, rhobar, thetabar) exnerbar_theta = thermodynamics.dexner_dtheta(equations.parameters, rhobar, thetabar) @@ -230,14 +227,21 @@ def V(u): ds_tbp = (ds_t(degree=(self.quadrature_degree)) + ds_b(degree=(self.quadrature_degree))) - # Add effect of density of water upon theta - # TODO: this has to be done using active tracers - if self.moisture is not None: - water_t = Function(Vtheta).assign(0.0) - for water in self.moisture: - water_t += self.equations.fields(water) - theta_w = theta / (1 + water_t) - thetabar_w = thetabar / (1 + water_t) + # Add effect of density of water upon theta, using moisture reference profiles + # TODO: Explore if this is the right thing to do for the linear problem + if equations.active_tracers is not None: + mr_t = Constant(0.0)*thetabar + for tracer in equations.active_tracers: + if tracer.chemical == 'H2O': + if tracer.variable_type == TracerVariableType.mixing_ratio: + idx = equations.field_names.index(tracer.name) + mr_bar = split(equations.X_ref)[idx] + mr_t += mr_bar + else: + raise NotImplementedError('Only mixing ratio tracers are implemented') + + theta_w = theta / (1 + mr_t) + thetabar_w = thetabar / (1 + mr_t) else: theta_w = theta thetabar_w = thetabar @@ -260,18 +264,12 @@ def L_tr(f): rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg) exner_avg_prb = LinearVariationalProblem(a_tr, L_tr(exnerbar), exnerbar_avg) - rho_avg_solver = LinearVariationalSolver(rho_avg_prb, - solver_parameters=cg_ilu_parameters, - options_prefix='rhobar_avg_solver') - exner_avg_solver = LinearVariationalSolver(exner_avg_prb, - solver_parameters=cg_ilu_parameters, - options_prefix='exnerbar_avg_solver') - - with timed_region("Gusto:HybridProjectRhobar"): - rho_avg_solver.solve() - - with timed_region("Gusto:HybridProjectExnerbar"): - exner_avg_solver.solve() + self.rho_avg_solver = LinearVariationalSolver(rho_avg_prb, + solver_parameters=cg_ilu_parameters, + options_prefix='rhobar_avg_solver') + self.exner_avg_solver = LinearVariationalSolver(exner_avg_prb, + solver_parameters=cg_ilu_parameters, + options_prefix='exnerbar_avg_solver') # "broken" u, rho, and trace system # NOTE: no ds_v integrals since equations are defined on @@ -306,6 +304,7 @@ def L_tr(f): + dl*dot(u, n)*(ds_tbp + ds_vp) ) + # TODO: can we get this term using FML? # contribution of the sponge term if hasattr(self.equations, "mu"): eqn += dt*self.equations.mu*inner(w, k)*inner(u, k)*dx @@ -364,6 +363,13 @@ def solve(self, xrhs, dy): """ self.xrhs.assign(xrhs) + # TODO: can we avoid computing these each time the solver is called? + with timed_region("Gusto:HybridProjectRhobar"): + self.rho_avg_solver.solve() + + with timed_region("Gusto:HybridProjectExnerbar"): + self.exner_avg_solver.solve() + # Solve the hybridized system self.hybridized_solver.solve() @@ -443,7 +449,7 @@ def _setup_solver(self): u, p = TrialFunctions(M) # Get background fields - bbar = equation.fields("bbar") + bbar = split(equation.X_ref)[2] # Analytical (approximate) elimination of theta k = equation.domain.k # Upward pointing unit vector diff --git a/gusto/physics.py b/gusto/physics.py index f45c63cba..6eb863850 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -29,8 +29,20 @@ class Physics(object, metaclass=ABCMeta): """Base class for the parametrisation of physical processes for Gusto.""" - def __init__(self): - pass + def __init__(self, equation, parameters=None): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + 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.equation = equation + if parameters is None and hasattr(equation, 'parameters'): + self.parameters = equation.parameters + else: + self.parameters = parameters @abstractmethod def evaluate(self): @@ -76,6 +88,8 @@ def __init__(self, equation, vapour_name='water_vapour', CompressibleEulerEquations. """ + super().__init__(equation, 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 @@ -86,9 +100,8 @@ def __init__(self, equation, vapour_name='water_vapour', assert cloud_name in equation.field_names, 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.equation = equation - parameters = equation.parameters if parameters is None else parameters self.latent_heat = latent_heat # Vapour and cloud variables are needed for every form of this scheme @@ -269,8 +282,8 @@ def __init__(self, equation, rain_name, domain, moments=AdvectedMoments.M3): test = equation.tests[rain_idx] Vu = domain.spaces("HDiv") - # TODO: how do we allow this to be output? - v = equation.fields(name='rainfall_velocity', space=Vu) + # 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 @@ -455,7 +468,7 @@ class EvaporationOfRain(Physics): """ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', - latent_heat=True): + latent_heat=True, parameters=None): """ Args: equation (:class:`PrognosticEquationSet`): the model's equation. @@ -465,11 +478,17 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', Defaults to 'water_vapour'. 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. """ + + super().__init__(equation, 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 @@ -481,8 +500,7 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', # Make prognostic for physics scheme self.X = Function(equation.X.function_space()) - self.equation = equation - parameters = equation.parameters + parameters = self.parameters self.latent_heat = latent_heat # Vapour and cloud variables are needed for every form of this scheme @@ -607,7 +625,7 @@ def evaluate(self, x_in, dt): interpolator.interpolate() -class InstantRain(object): +class InstantRain(Physics): """ The process of converting vapour above the saturation curve to rain. @@ -619,7 +637,8 @@ class InstantRain(object): """ def __init__(self, equation, saturation_curve, vapour_name="water_vapour", - rain_name=None, convective_feedback=False, set_tau_to_dt=False): + rain_name=None, convective_feedback=False, set_tau_to_dt=False, + parameters=None): """ Args: equation (:class: 'PrognosticEquationSet'): the model's equation. @@ -635,9 +654,14 @@ def __init__(self, equation, saturation_curve, vapour_name="water_vapour", conversion is equal to the timestep and False if not. If False then the user must provide a timescale, tau, that gets passed to the parameters list. + parameters (:class:`Configuration`, optional): parameters containing + the values of gas constants. Defaults to None, in which case the + parameters are obtained from the equation. """ - parameters = equation.parameters + super().__init__(equation, parameters=None) + + parameters = self.parameters self.convective_feedback = convective_feedback self.set_tau_to_dt = set_tau_to_dt diff --git a/gusto/state.py b/gusto/state.py deleted file mode 100644 index b4c4ed2cb..000000000 --- a/gusto/state.py +++ /dev/null @@ -1,807 +0,0 @@ -""" -Provides the model's state object, which controls IO and other core functions. - -The model's :class:`State` object is defined in this module. It controls various -input/output (IO) aspects, as well as setting up the compatible finite element -spaces and holding the mesh. In some ways it acts as a bucket, holding core -parts of the model. -""" -from os import path, makedirs -import itertools -from netCDF4 import Dataset -import sys -import time -from gusto.diagnostics import Diagnostics, Perturbation, SteadyStateError -from firedrake import (FiniteElement, TensorProductElement, HDiv, - FunctionSpace, VectorFunctionSpace, - interval, Function, Mesh, functionspaceimpl, - File, SpatialCoordinate, sqrt, Constant, inner, - op2, DumbCheckpoint, FILE_CREATE, FILE_READ, interpolate, - CellNormal, cross, as_vector) -import numpy as np -from gusto.configuration import logger, set_log_handler -from gusto.fields import StateFields - -__all__ = ["State"] - - -class SpaceCreator(object): - """Object to create and hold the model's finite element spaces.""" - def __init__(self, mesh): - """ - Args: - mesh (:class:`Mesh`): the model's mesh. - """ - self.mesh = mesh - self.extruded_mesh = hasattr(mesh, "_base_mesh") - self._initialised_base_spaces = False - - def __call__(self, name, family=None, degree=None, V=None): - """ - Returns a space, and also creates it if it is not created yet. - - If a space needs creating, it may be that more arguments (such as the - family and degree) need to be provided. Alternatively a space can be - passed in to be stored in the creator. - - Args: - name (str): the name of the space. - family (str, optional): name of the finite element family to be - created. Defaults to None. - degree (int, optional): the degree of the finite element space to be - created. Defaults to None. - V (:class:`FunctionSpace`, optional): an existing space, to be - stored in the creator object. If this is provided, it will be - added to the creator and no other action will be taken. This - space will be returned. Defaults to None. - - Returns: - :class:`FunctionSpace`: the desired function space. - """ - - try: - return getattr(self, name) - except AttributeError: - if V is not None: - value = V - elif name == "HDiv" and family in ["BDM", "RT", "CG", "RTCF"]: - value = self.build_hdiv_space(family, degree) - elif name == "theta": - value = self.build_theta_space(degree) - elif name == "DG1_equispaced": - value = self.build_dg_space(1, variant='equispaced') - elif family == "DG": - value = self.build_dg_space(degree) - elif family == "CG": - value = self.build_cg_space(degree) - else: - raise ValueError(f'State has no space corresponding to {name}') - setattr(self, name, value) - return value - - def build_compatible_spaces(self, family, degree): - """ - Builds the sequence of compatible finite element spaces for the mesh. - - If the mesh is not extruded, this builds and returns the spaces: - (HDiv, DG). - If the mesh is extruded, this builds and returns the following spaces: - (HDiv, DG, theta). - The 'theta' space corresponds to the vertical component of the velocity. - - Args: - family (str): the family of the horizontal part of the HDiv space. - degree (int): the polynomial degree of the DG space. - - Returns: - tuple: the created compatible :class:`FunctionSpace` objects. - """ - if self.extruded_mesh and not self._initialised_base_spaces: - self.build_base_spaces(family, degree) - Vu = self.build_hdiv_space(family, degree) - setattr(self, "HDiv", Vu) - Vdg = self.build_dg_space(degree) - setattr(self, "DG", Vdg) - Vth = self.build_theta_space(degree) - setattr(self, "theta", Vth) - return Vu, Vdg, Vth - else: - Vu = self.build_hdiv_space(family, degree) - setattr(self, "HDiv", Vu) - Vdg = self.build_dg_space(degree) - setattr(self, "DG", Vdg) - return Vu, Vdg - - def build_base_spaces(self, family, degree): - """ - Builds the :class:`FiniteElement` objects for the base mesh. - - Args: - family (str): the family of the horizontal part of the HDiv space. - degree (int): the polynomial degree of the DG space. - """ - cell = self.mesh._base_mesh.ufl_cell().cellname() - - # horizontal base spaces - self.S1 = FiniteElement(family, cell, degree+1) - self.S2 = FiniteElement("DG", cell, degree) - - # vertical base spaces - self.T0 = FiniteElement("CG", interval, degree+1) - self.T1 = FiniteElement("DG", interval, degree) - - self._initialised_base_spaces = True - - def build_hdiv_space(self, family, degree): - """ - Builds and returns the HDiv :class:`FunctionSpace`. - - Args: - family (str): the family of the horizontal part of the HDiv space. - degree (int): the polynomial degree of the space. - - Returns: - :class:`FunctionSpace`: the HDiv space. - """ - if self.extruded_mesh: - if not self._initialised_base_spaces: - self.build_base_spaces(family, degree) - Vh_elt = HDiv(TensorProductElement(self.S1, self.T1)) - Vt_elt = TensorProductElement(self.S2, self.T0) - Vv_elt = HDiv(Vt_elt) - V_elt = Vh_elt + Vv_elt - else: - cell = self.mesh.ufl_cell().cellname() - V_elt = FiniteElement(family, cell, degree+1) - return FunctionSpace(self.mesh, V_elt, name='HDiv') - - def build_dg_space(self, degree, variant=None): - """ - Builds and returns the DG :class:`FunctionSpace`. - - Args: - degree (int): the polynomial degree of the space. - variant (str): the variant of the underlying :class:`FiniteElement` - to use. Defaults to None, which will call the default variant. - - Returns: - :class:`FunctionSpace`: the DG space. - """ - if self.extruded_mesh: - if not self._initialised_base_spaces or self.T1.degree() != degree or self.T1.variant() != variant: - cell = self.mesh._base_mesh.ufl_cell().cellname() - S2 = FiniteElement("DG", cell, degree, variant=variant) - T1 = FiniteElement("DG", interval, degree, variant=variant) - else: - S2 = self.S2 - T1 = self.T1 - V_elt = TensorProductElement(S2, T1) - else: - cell = self.mesh.ufl_cell().cellname() - V_elt = FiniteElement("DG", cell, degree, variant=variant) - name = f'DG{degree}_equispaced' if variant == 'equispaced' else f'DG{degree}' - return FunctionSpace(self.mesh, V_elt, name=name) - - def build_theta_space(self, degree): - """ - Builds and returns the 'theta' space. - - This corresponds to the non-Piola mapped space of the vertical component - of the velocity. The space will be discontinuous in the horizontal but - continuous in the vertical. - - Args: - degree (int): degree of the corresponding density space. - - Raises: - AssertionError: the mesh is not extruded. - - Returns: - :class:`FunctionSpace`: the 'theta' space. - """ - assert self.extruded_mesh - if not self._initialised_base_spaces: - cell = self.mesh._base_mesh.ufl_cell().cellname() - self.S2 = FiniteElement("DG", cell, degree) - self.T0 = FiniteElement("CG", interval, degree+1) - V_elt = TensorProductElement(self.S2, self.T0) - return FunctionSpace(self.mesh, V_elt, name='Vtheta') - - def build_cg_space(self, degree): - """ - Builds the continuous scalar space at the top of the de Rham complex. - - Args: - degree (int): degree of the continuous space. - - Returns: - :class:`FunctionSpace`: the continuous space. - """ - return FunctionSpace(self.mesh, "CG", degree, name=f'CG{degree}') - - -class PointDataOutput(object): - """Object for outputting field point data.""" - def __init__(self, filename, ndt, field_points, description, - field_creator, comm, tolerance=None, create=True): - """ - Args: - filename (str): name of file to output to. - ndt (int): number of time points to output at. TODO: remove as this - is unused. - field_points (list): some iterable of pairs, matching fields with - arrays of evaluation points: (field_name, evaluation_points). - description (str): a description of the simulation to be included in - the output. - field_creator (:class:`FieldCreator`): the field creator, used to - determine the datatype and shape of fields. - comm (:class:`MPI.Comm`): MPI communicator. - tolerance (float, optional): tolerance to use for the evaluation of - fields at points. Defaults to None. - create (bool, optional): whether the output file needs creating, or - if it already exists. Defaults to True. - """ - # Overwrite on creation. - self.dump_count = 0 - self.filename = filename - self.field_points = field_points - self.tolerance = tolerance - self.comm = comm - if not create: - return - if self.comm.rank == 0: - with Dataset(filename, "w") as dataset: - dataset.description = "Point data for simulation {desc}".format(desc=description) - dataset.history = "Created {t}".format(t=time.ctime()) - # FIXME add versioning information. - dataset.source = "Output from Gusto model" - # Appendable dimension, timesteps in the model - dataset.createDimension("time", None) - - var = dataset.createVariable("time", np.float64, ("time")) - var.units = "seconds" - # Now create the variable group for each field - for field_name, points in field_points: - group = dataset.createGroup(field_name) - npts, dim = points.shape - group.createDimension("points", npts) - group.createDimension("geometric_dimension", dim) - var = group.createVariable("points", points.dtype, - ("points", "geometric_dimension")) - var[:] = points - - # Get the UFL shape of the field - field_shape = field_creator(field_name).ufl_shape - # Number of geometric dimension occurences should be the same as the length of the UFL shape - field_len = len(field_shape) - field_count = field_shape.count(dim) - assert field_len == field_count, "Geometric dimension occurrences do not match UFL shape" - # Create the variable with the required shape - dimensions = ("time", "points") + field_count*("geometric_dimension",) - group.createVariable(field_name, field_creator(field_name).dat.dtype, dimensions) - - def dump(self, field_creator, t): - """ - Evaluate and output field data at points. - - Args: - field_creator (:class:`FieldCreator`): gives access to the fields. - t (float): simulation time at which the output occurs. - """ - - val_list = [] - for field_name, points in self.field_points: - val_list.append((field_name, np.asarray(field_creator(field_name).at(points, tolerance=self.tolerance)))) - - if self.comm.rank == 0: - with Dataset(self.filename, "a") as dataset: - # Add new time index - dataset.variables["time"][self.dump_count] = t - for field_name, vals in val_list: - group = dataset.groups[field_name] - var = group.variables[field_name] - var[self.dump_count, :] = vals - - self.dump_count += 1 - - -class DiagnosticsOutput(object): - """Object for outputting global diagnostic data.""" - def __init__(self, filename, diagnostics, description, comm, create=True): - """ - Args: - filename (str): name of file to output to. - diagnostics (:class:`Diagnostics`): the object holding and - controlling the diagnostic evaluation. - description (str): a description of the simulation to be included in - the output. - comm (:class:`MPI.Comm`): MPI communicator. - create (bool, optional): whether the output file needs creating, or - if it already exists. Defaults to True. - """ - self.filename = filename - self.diagnostics = diagnostics - self.comm = comm - if not create: - return - if self.comm.rank == 0: - with Dataset(filename, "w") as dataset: - dataset.description = "Diagnostics data for simulation {desc}".format(desc=description) - dataset.history = "Created {t}".format(t=time.ctime()) - dataset.source = "Output from Gusto model" - dataset.createDimension("time", None) - var = dataset.createVariable("time", np.float64, ("time", )) - var.units = "seconds" - for name in diagnostics.fields: - group = dataset.createGroup(name) - for diagnostic in diagnostics.available_diagnostics: - group.createVariable(diagnostic, np.float64, ("time", )) - - def dump(self, state, t): - """ - Output the global diagnostics. - - state (:class:`State`): the model's state object. - t (float): simulation time at which the output occurs. - """ - - diagnostics = [] - for fname in self.diagnostics.fields: - field = state.fields(fname) - for dname in self.diagnostics.available_diagnostics: - diagnostic = getattr(self.diagnostics, dname) - diagnostics.append((fname, dname, diagnostic(field))) - - if self.comm.rank == 0: - with Dataset(self.filename, "a") as dataset: - idx = dataset.dimensions["time"].size - dataset.variables["time"][idx:idx + 1] = t - for fname, dname, value in diagnostics: - group = dataset.groups[fname] - var = group.variables[dname] - var[idx:idx + 1] = value - - -class State(object): - """Keeps the model's mesh and variables, and controls its IO.""" - - def __init__(self, mesh, dt, - output=None, - parameters=None, - diagnostics=None, - diagnostic_fields=None): - """ - Args: - mesh (:class:`Mesh`): the model's mesh. - dt (:class:`Constant`): the time taken to perform a single model - step. If a float or int is passed, it will be cast to a - :class:`Constant`. - output (:class:`OutputParameters`, optional): holds and describes - the options for outputting. Defaults to None. - parameters (:class:`Configuration`, optional): an object containing - the model's physical parameters. Defaults to None. - diagnostics (:class:`Diagnostics`, optional): object holding and - controlling the model's diagnostics. Defaults to None. - diagnostic_fields (list, optional): an iterable of `DiagnosticField` - objects. Defaults to None. - - Raises: - RuntimeError: if no output is provided. - TypeError: if `dt` cannot be cast to a :class:`Constant`. - """ - - if output is None: - # TODO: maybe this shouldn't be an optional argument then? - raise RuntimeError("You must provide a directory name for dumping results") - else: - self.output = output - self.parameters = parameters - - if diagnostics is not None: - self.diagnostics = diagnostics - else: - self.diagnostics = Diagnostics() - if diagnostic_fields is not None: - self.diagnostic_fields = diagnostic_fields - else: - self.diagnostic_fields = [] - - # The mesh - self.mesh = mesh - - self.spaces = SpaceCreator(mesh) - - if self.output.dumplist is None: - - self.output.dumplist = [] - - self.fields = StateFields(*self.output.dumplist) - - self.dumpdir = None - self.dumpfile = None - self.to_pickup = None - - # figure out if we're on a sphere - try: - self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3 and mesh._base_mesh.topological_dimension() == 2) - except AttributeError: - self.on_sphere = (mesh.geometric_dimension() == 3 and mesh.topological_dimension() == 2) - - # build the vertical normal and define perp for 2d geometries - dim = mesh.topological_dimension() - if self.on_sphere: - x = SpatialCoordinate(mesh) - R = sqrt(inner(x, x)) - self.k = interpolate(x/R, mesh.coordinates.function_space()) - if dim == 2: - outward_normals = CellNormal(mesh) - self.perp = lambda u: cross(outward_normals, u) - else: - kvec = [0.0]*dim - kvec[dim-1] = 1.0 - self.k = Constant(kvec) - if dim == 2: - self.perp = lambda u: as_vector([-u[1], u[0]]) - - # setup logger - logger.setLevel(output.log_level) - set_log_handler(mesh.comm) - if parameters is not None: - logger.info("Physical parameters that take non-default values:") - logger.info(", ".join("%s: %s" % (k, float(v)) for (k, v) in vars(parameters).items())) - - # Constant to hold current time - self.t = Constant(0.0) - if type(dt) is Constant: - self.dt = dt - elif type(dt) in (float, int): - self.dt = Constant(dt) - else: - raise TypeError(f'dt must be a Constant, float or int, not {type(dt)}') - - def setup_diagnostics(self): - """Concatenates the various types of diagnostic field.""" - for name in self.output.perturbation_fields: - f = Perturbation(name) - self.diagnostic_fields.append(f) - - for name in self.output.steady_state_error_fields: - f = SteadyStateError(self, name) - self.diagnostic_fields.append(f) - - fields = set([f.name() for f in self.fields]) - field_deps = [(d, sorted(set(d.required_fields).difference(fields),)) for d in self.diagnostic_fields] - schedule = topo_sort(field_deps) - self.diagnostic_fields = schedule - for diagnostic in self.diagnostic_fields: - diagnostic.setup(self) - self.diagnostics.register(diagnostic.name) - - def setup_dump(self, t, tmax, pickup=False): - """ - Sets up a series of things used for outputting. - - This prepares the model for outputting. First it checks for the - existence the specified outputting directory, so prevent it being - overwritten unintentionally. It then sets up the output files and the - checkpointing file. - - Args: - t (float): the current model time. - tmax (float): the end time of the model's simulation. - pickup (bool, optional): whether to pick up the model's initial - state from a checkpointing file. Defaults to False. - - Raises: - IOError: if the results directory already exists, and the model is - not picking up or running in test mode. - """ - - if any([self.output.dump_vtus, self.output.dumplist_latlon, - self.output.dump_diagnostics, self.output.point_data, - self.output.checkpoint and not pickup]): - # setup output directory and check that it does not already exist - self.dumpdir = path.join("results", self.output.dirname) - running_tests = '--running-tests' in sys.argv or "pytest" in self.output.dirname - if self.mesh.comm.rank == 0: - if not running_tests and path.exists(self.dumpdir) and not pickup: - raise IOError("results directory '%s' already exists" - % self.dumpdir) - else: - if not running_tests: - makedirs(self.dumpdir) - - if self.output.dump_vtus: - - # setup pvd output file - outfile = path.join(self.dumpdir, "field_output.pvd") - self.dumpfile = File( - outfile, project_output=self.output.project_fields, - comm=self.mesh.comm) - - # make list of fields to dump - self.to_dump = [f for f in self.fields if f.name() in self.fields.to_dump] - - # make dump counter - self.dumpcount = itertools.count() - - # if there are fields to be dumped in latlon coordinates, - # setup the latlon coordinate mesh and make output file - if len(self.output.dumplist_latlon) > 0: - mesh_ll = get_latlon_mesh(self.mesh) - outfile_ll = path.join(self.dumpdir, "field_output_latlon.pvd") - self.dumpfile_ll = File(outfile_ll, - project_output=self.output.project_fields, - comm=self.mesh.comm) - - # make functions on latlon mesh, as specified by dumplist_latlon - self.to_dump_latlon = [] - for name in self.output.dumplist_latlon: - f = self.fields(name) - field = Function( - functionspaceimpl.WithGeometry.create( - f.function_space(), mesh_ll), - val=f.topological, name=name+'_ll') - self.to_dump_latlon.append(field) - - # we create new netcdf files to write to, unless pickup=True, in - # which case we just need the filenames - if self.output.dump_diagnostics: - diagnostics_filename = self.dumpdir+"/diagnostics.nc" - self.diagnostic_output = DiagnosticsOutput(diagnostics_filename, - self.diagnostics, - self.output.dirname, - self.mesh.comm, - create=not pickup) - - if len(self.output.point_data) > 0: - # set up point data output - pointdata_filename = self.dumpdir+"/point_data.nc" - ndt = int(tmax/float(self.dt)) - self.pointdata_output = PointDataOutput(pointdata_filename, ndt, - self.output.point_data, - self.output.dirname, - self.fields, - self.mesh.comm, - self.output.tolerance, - create=not pickup) - - # make point data dump counter - self.pddumpcount = itertools.count() - - # set frequency of point data output - defaults to - # dumpfreq if not set by user - if self.output.pddumpfreq is None: - self.output.pddumpfreq = self.output.dumpfreq - - # if we want to checkpoint and are not picking up from a previous - # checkpoint file, setup the checkpointing - if self.output.checkpoint: - if not pickup: - self.chkpt = DumbCheckpoint(path.join(self.dumpdir, "chkpt"), - mode=FILE_CREATE) - # make list of fields to pickup (this doesn't include - # diagnostic fields) - self.to_pickup = [f for f in self.fields if f.name() in self.fields.to_pickup] - - # if we want to checkpoint then make a checkpoint counter - if self.output.checkpoint: - self.chkptcount = itertools.count() - - # dump initial fields - self.dump(t) - - def pickup_from_checkpoint(self): - """Picks up the model's variables from a checkpoint file.""" - # TODO: this duplicates some code from setup_dump. Can this be avoided? - # It is because we don't know if we are picking up or setting dump first - if self.to_pickup is None: - self.to_pickup = [f for f in self.fields if f.name() in self.fields.to_pickup] - # Set dumpdir if has not been done already - if self.dumpdir is None: - self.dumpdir = path.join("results", self.output.dirname) - - if self.output.checkpoint: - # Open the checkpointing file for writing - if self.output.checkpoint_pickup_filename is not None: - chkfile = self.output.checkpoint_pickup_filename - else: - chkfile = path.join(self.dumpdir, "chkpt") - with DumbCheckpoint(chkfile, mode=FILE_READ) as chk: - # Recover all the fields from the checkpoint - for field in self.to_pickup: - chk.load(field) - t = chk.read_attribute("/", "time") - # Setup new checkpoint - self.chkpt = DumbCheckpoint(path.join(self.dumpdir, "chkpt"), mode=FILE_CREATE) - else: - raise ValueError("Must set checkpoint True if pickup") - - return t - - def dump(self, t): - """ - Dumps all of the required model output. - - This includes point data, global diagnostics and general field data to - paraview data files. Also writes the model's prognostic variables to - a checkpoint file if specified. - - Args: - t (float): the simulation's current time. - """ - output = self.output - - # Diagnostics: - # Compute diagnostic fields - for field in self.diagnostic_fields: - field(self) - - if output.dump_diagnostics: - # Output diagnostic data - self.diagnostic_output.dump(self, t) - - if len(output.point_data) > 0 and (next(self.pddumpcount) % output.pddumpfreq) == 0: - # Output pointwise data - self.pointdata_output.dump(self.fields, t) - - # Dump all the fields to the checkpointing file (backup version) - if output.checkpoint and (next(self.chkptcount) % output.chkptfreq) == 0: - for field in self.to_pickup: - self.chkpt.store(field) - self.chkpt.write_attribute("/", "time", t) - - if output.dump_vtus and (next(self.dumpcount) % output.dumpfreq) == 0: - # dump fields - self.dumpfile.write(*self.to_dump) - - # dump fields on latlon mesh - if len(output.dumplist_latlon) > 0: - self.dumpfile_ll.write(*self.to_dump_latlon) - - def initialise(self, initial_conditions): - """ - Initialise the state's prognostic variables. - - Args: - initial_conditions (list): an iterable of pairs: (field_name, expr), - where 'field_name' is the string giving the name of the - prognostic field and expr is the :class:`ufl.Expr` whose value - is used to set the initial field. - """ - for name, ic in initial_conditions: - f_init = getattr(self.fields, name) - f_init.assign(ic) - f_init.rename(name) - - def set_reference_profiles(self, reference_profiles): - """ - Initialise the state'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 name, profile in reference_profiles: - if name+'bar' in self.fields: - # For reference profiles already added to state, allow - # interpolation from expressions - ref = self.fields(name+'bar') - elif isinstance(profile, Function): - # Need to add reference profile to state so profile must be - # a Function - ref = self.fields(name+'bar', space=profile.function_space(), dump=False) - else: - raise ValueError(f'When initialising reference profile {name}' - + ' the passed profile must be a Function') - ref.interpolate(profile) - - -def get_latlon_mesh(mesh): - """ - Construct a planar latitude-longitude mesh from a spherical mesh. - - Args: - mesh (:class:`State`): the mesh on which the simulation is performed. - """ - coords_orig = mesh.coordinates - coords_fs = coords_orig.function_space() - - if coords_fs.extruded: - cell = mesh._base_mesh.ufl_cell().cellname() - DG1_hori_elt = FiniteElement("DG", cell, 1, variant="equispaced") - DG1_vert_elt = FiniteElement("DG", interval, 1, variant="equispaced") - DG1_elt = TensorProductElement(DG1_hori_elt, DG1_vert_elt) - else: - cell = mesh.ufl_cell().cellname() - DG1_elt = FiniteElement("DG", cell, 1, variant="equispaced") - vec_DG1 = VectorFunctionSpace(mesh, DG1_elt) - coords_dg = Function(vec_DG1).interpolate(coords_orig) - coords_latlon = Function(vec_DG1) - shapes = {"nDOFs": vec_DG1.finat_element.space_dimension(), 'dim': 3} - - radius = np.min(np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2)) - # lat-lon 'x' = atan2(y, x) - coords_latlon.dat.data[:, 0] = np.arctan2(coords_dg.dat.data[:, 1], coords_dg.dat.data[:, 0]) - # lat-lon 'y' = asin(z/sqrt(x^2 + y^2 + z^2)) - coords_latlon.dat.data[:, 1] = np.arcsin(coords_dg.dat.data[:, 2]/np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2)) - # our vertical coordinate is radius - the minimum radius - coords_latlon.dat.data[:, 2] = np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2) - radius - -# We need to ensure that all points in a cell are on the same side of the branch cut in longitude coords -# This kernel amends the longitude coords so that all longitudes in one cell are close together - kernel = op2.Kernel(""" -#define PI 3.141592653589793 -#define TWO_PI 6.283185307179586 -void splat_coords(double *coords) {{ - double max_diff = 0.0; - double diff = 0.0; - - for (int i=0; i<{nDOFs}; i++) {{ - for (int j=0; j<{nDOFs}; j++) {{ - diff = coords[i*{dim}] - coords[j*{dim}]; - if (fabs(diff) > max_diff) {{ - max_diff = diff; - }} - }} - }} - - if (max_diff > PI) {{ - for (int i=0; i<{nDOFs}; i++) {{ - if (coords[i*{dim}] < 0) {{ - coords[i*{dim}] += TWO_PI; - }} - }} - }} -}} -""".format(**shapes), "splat_coords") - - op2.par_loop(kernel, coords_latlon.cell_set, - coords_latlon.dat(op2.RW, coords_latlon.cell_node_map())) - return Mesh(coords_latlon) - - -def topo_sort(field_deps): - """ - Perform a topological sort to determine the order to evaluate diagnostics. - - Args: - field_deps (list): a list of tuples, pairing diagnostic fields with the - fields that they are to be evaluated from. - - Raises: - RuntimeError: if there is a cyclic dependency in the diagnostic fields. - - Returns: - list: a list specifying the order in which to evaluate the diagnostics. - """ - name2field = dict((f.name, f) for f, _ in field_deps) - # map node: (input_deps, output_deps) - graph = dict((f.name, (list(deps), [])) for f, deps in field_deps) - roots = [] - for f, input_deps in field_deps: - if len(input_deps) == 0: - # No dependencies, candidate for evaluation - roots.append(f.name) - for d in input_deps: - # add f as output dependency - graph[d][1].append(f.name) - - schedule = [] - while roots: - n = roots.pop() - schedule.append(n) - output_deps = list(graph[n][1]) - for m in output_deps: - # Remove edge - graph[m][0].remove(n) - graph[n][1].remove(m) - # If m now as no input deps, candidate for evaluation - if len(graph[m][0]) == 0: - roots.append(m) - if any(len(i) for i, _ in graph.values()): - cycle = "\n".join("%s -> %s" % (f, i) for f, (i, _) in graph.items() - if f not in schedule) - raise RuntimeError("Field dependencies have a cycle:\n\n%s" % cycle) - return list(map(name2field.__getitem__, schedule)) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 326c20efe..e5406df66 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -132,7 +132,7 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): if self.field_name is not None and hasattr(equation, "field_names"): self.idx = equation.field_names.index(self.field_name) - self.fs = equation.fields(self.field_name).function_space() + self.fs = equation.spaces[self.idx] self.residual = self.residual.label_map( lambda t: t.get(prognostic) == self.field_name, lambda t: Term( @@ -183,7 +183,7 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): if self.idx is None: self.x_projected = Function(equation.function_space) else: - self.x_projected = Function(equation.fields(self.field_name).function_space()) + self.x_projected = Function(equation.spaces[self.idx]) new_test = TestFunction(self.fs) parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', @@ -256,7 +256,10 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): if self.discretisation_option == "recovered": # set up the necessary functions - self.x_in = Function(equation.fields(self.field_name).function_space()) + if self.idx is not None: + self.x_in = Function(equation.spaces[self.idx]) + else: + self.x_in = Function(equation.function_space) # Operator to recover to higher discontinuous space self.x_recoverer = ReversibleRecoverer(self.x_in, self.xdg_in, options) @@ -374,7 +377,10 @@ def replace_transport_term(self): # Do the options specify a different ibp to the old transport term? if old_transport_term.labels['ibp'] != self.options.ibp: # Set up a new transport term - field = self.equation.fields(self.field_name) + if self.idx is not None: + field = self.equation.X.split()[self.idx] + else: + field = self.equation.X test = TestFunction(self.fs) # Set up new transport term (depending on the type of transport equation) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index b0409ae49..1b4df1146 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -4,12 +4,13 @@ from firedrake import Function, Projector, Constant from pyop2.profiling import timed_stage from gusto.configuration import logger +from gusto.equations import PrognosticEquationSet from gusto.forcing import Forcing from gusto.fml.form_manipulation_labelling import drop from gusto.labels import (transport, diffusion, time_derivative, linearisation, prognostic, physics) from gusto.linear_solvers import LinearTimesteppingSolver -from gusto.fields import TimeLevelFields +from gusto.fields import TimeLevelFields, StateFields from gusto.time_discretisation import ExplicitTimeDiscretisation __all__ = ["Timestepper", "SplitPhysicsTimestepper", "SemiImplicitQuasiNewton", @@ -30,10 +31,13 @@ def __init__(self, equation, io): self.io = io self.dt = self.equation.domain.dt self.t = Constant(0.0) + self.reference_profiles_initialised = False self.setup_fields() self.setup_scheme() + self.io.log_parameters(equation) + @abstractproperty def transporting_velocity(self): return NotImplementedError @@ -41,6 +45,7 @@ def transporting_velocity(self): @abstractmethod def setup_fields(self): """Set up required fields. Must be implemented in child classes""" + # TODO: should we actually implement this? pass @abstractmethod @@ -63,20 +68,16 @@ def run(self, t, tmax, pickup=False): pickup: (bool): specify whether to pickup from a previous run """ - io = self.io - if pickup: - t = io.pickup_from_checkpoint() + t = self.io.pickup_from_checkpoint(self.fields) - io.setup_diagnostics() + self.io.setup_diagnostics(self.fields) with timed_stage("Dump output"): - io.setup_dump(t, tmax, pickup) + self.io.setup_dump(self.fields, t, tmax, pickup) self.t.assign(t) - self.x.initialise(self.equation) - while float(self.t) < tmax - 0.5*float(self.dt): logger.info(f'at start of timestep, t={float(self.t)}, dt={float(self.dt)}') @@ -84,19 +85,66 @@ def run(self, t, tmax, pickup=False): self.timestep() - for field in self.x.np1: - self.equation.fields(field.name()).assign(field) - self.t.assign(self.t + self.dt) with timed_stage("Dump output"): - io.dump(float(self.t)) + self.io.dump(self.fields, float(self.t)) - if io.output.checkpoint: - io.chkpt.close() + if self.io.output.checkpoint: + 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. + """ + # TODO: come back and consider all aspects of this + 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(), dump=False) + else: + raise ValueError(f'When initialising reference profile {field_name}' + + ' the passed profile must be a Function') + ref.interpolate(profile) + + # Assign profile to X_ref belonging to equation + if isinstance(self.equation, PrognosticEquationSet): + assert field_name in self.equation.field_names, \ + f'Cannot set reference profile as field {field_name} not found' + idx = self.equation.field_names.index(field_name) + X_ref = self.equation.X_ref.split()[idx] + X_ref.assign(ref) + + self.reference_profiles_initialised = True + + # TODO: do we need this interface? If so, should we use it in all examples? + def initialise(self, initial_conditions): + """ + Initialise the state's fields. + + Args: + initial_conditions (list): an iterable of pairs: (field_name, expr), + where 'field_name' is the string giving the name of the + prognostic field and expr is the :class:`ufl.Expr` whose value + is used to set the initial field. + """ + for field_name, ic in initial_conditions: + f_init = getattr(self.fields, field_name) + f_init.assign(ic) + f_init.rename(field_name) + class Timestepper(BaseTimestepper): """ @@ -120,6 +168,8 @@ def transporting_velocity(self): def setup_fields(self): self.x = TimeLevelFields(self.equation, self.scheme.nlevels) + self.fields = StateFields(self.x.np1, self.equation.prescribed_fields, + *self.io.output.dumplist) def setup_scheme(self): self.scheme.setup(self.equation, self.transporting_velocity) @@ -182,6 +232,8 @@ def transporting_velocity(self): def setup_fields(self): self.x = TimeLevelFields(self.equation, self.scheme.nlevels) + self.fields = StateFields(self.x.np1, self.equation.prescribed_fields, + *self.io.output.dumplist) def setup_scheme(self): from gusto.labels import coriolis, pressure_gradient @@ -203,7 +255,7 @@ def timestep(self): class SemiImplicitQuasiNewton(BaseTimestepper): """ Implements a semi-implicit quasi-Newton discretisation, - with Strang splitting and auxilliary semi-Lagrangian transport. + 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 @@ -269,20 +321,21 @@ def __init__(self, equation_set, io, transport_schemes, assert scheme.field_name in equation_set.field_names self.diffusion_schemes.append((scheme.field_name, scheme)) - if not equation_set.reference_profiles_initialised: - raise RuntimeError('Reference profiles for equation set must be initialised to use Semi-Implicit Timestepper') - - super().__init__(equation_set, io) - if auxiliary_equations_and_schemes is not None: for eqn, scheme in auxiliary_equations_and_schemes: - self.x.add_fields(eqn) - scheme.setup(eqn, self.transporting_velocity) + 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: + aux_scheme.setup(aux_eqn, self.transporting_velocity) self.tracers_to_copy = [] for name in equation_set.field_names: @@ -326,6 +379,13 @@ 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")) + 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.np1, self.equation.prescribed_fields, + *self.io.output.dumplist) def setup_scheme(self): """Sets up transport, diffusion and physics schemes""" @@ -399,6 +459,7 @@ def timestep(self): for name, scheme in self.auxiliary_schemes: # transports a field from xn and puts result in xnp1 scheme.apply(xnp1(name), xn(name)) + print('xnp1 tracer', xnp1('tracer').dat.data.min(), xnp1('tracer').dat.data.max()) with timed_stage("Diffusion"): for name, scheme in self.diffusion_schemes: @@ -408,6 +469,21 @@ def timestep(self): for _, scheme in self.physics_schemes: scheme.apply(xnp1(scheme.field_name), xnp1(scheme.field_name)) + def run(self, t, tmax, pickup=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 + pickup: (bool): specify whether to pickup from a previous run + """ + + assert self.reference_profiles_initialised, \ + 'Reference profiles for must be initialised to use Semi-Implicit Timestepper' + + super().run(t, tmax, pickup=pickup) + class PrescribedTransport(Timestepper): """ @@ -450,16 +526,18 @@ def __init__(self, equation, scheme, io, physics_schemes=None, if prescribed_transporting_velocity is not None: self.velocity_projection = Projector( prescribed_transporting_velocity(self.t), - self.equation.fields('u')) + self.fields('u')) else: self.velocity_projection = None @property def transporting_velocity(self): - return self.equation.fields('u') + return self.fields('u') def setup_fields(self): self.x = TimeLevelFields(self.equation, self.scheme.nlevels) + self.fields = StateFields(self.x.np1, self.equation.prescribed_fields, + *self.io.output.dumplist) def setup_scheme(self): self.scheme.setup(self.equation, self.transporting_velocity) diff --git a/integration-tests/balance/test_compressible_balance.py b/integration-tests/balance/test_compressible_balance.py index 41db83aa3..f8efda306 100644 --- a/integration-tests/balance/test_compressible_balance.py +++ b/integration-tests/balance/test_compressible_balance.py @@ -12,38 +12,31 @@ def setup_balance(dirname): - # set up grid and time stepping parameters + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Parameters dt = 1. tmax = 5. deltax = 400 L = 2000. H = 10000. - nlayers = int(H/deltax) ncolumns = int(L/deltax) + # Domain m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) domain = Domain(mesh, dt, "CG", 1) - output = OutputParameters(dirname=dirname+'/dry_balance', dumpfreq=10, dumplist=['u']) + # Equation parameters = CompressibleParameters() eqns = CompressibleEulerEquations(domain, parameters) - io = IO(domain, eqns, output=output) - - # Initial conditions - rho0 = eqns.fields("rho") - theta0 = eqns.fields("theta") - - # Isentropic background state - Tsurf = Constant(300.) - theta0.interpolate(Tsurf) - - # Calculate hydrostatic exner - compressible_hydrostatic_balance(eqns, theta0, rho0, solve_for_rho=True) - eqns.set_reference_profiles([('rho', rho0), - ('theta', theta0)]) + # I/O + output = OutputParameters(dirname=dirname+'/dry_balance', dumpfreq=10, dumplist=['u']) + io = IO(domain, output) # Set up transport schemes transported_fields = [ImplicitMidpoint(domain, "u"), @@ -57,6 +50,23 @@ def setup_balance(dirname): stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, linear_solver=linear_solver) + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + + # Isentropic background state + Tsurf = Constant(300.) + theta0.interpolate(Tsurf) + + # Calculate hydrostatic exner + compressible_hydrostatic_balance(eqns, theta0, rho0, solve_for_rho=True) + + stepper.set_reference_profiles([('rho', rho0), + ('theta', theta0)]) + return stepper, tmax diff --git a/integration-tests/balance/test_saturated_balance.py b/integration-tests/balance/test_saturated_balance.py index 6c9b1f7a0..41dce0f18 100644 --- a/integration-tests/balance/test_saturated_balance.py +++ b/integration-tests/balance/test_saturated_balance.py @@ -14,24 +14,27 @@ def setup_saturated(dirname, recovered): - # set up grid and time stepping parameters + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Parameters dt = 1. tmax = 3. deltax = 400. L = 2000. H = 10000. - nlayers = int(H/deltax) ncolumns = int(L/deltax) - degree = 0 if recovered else 1 + # Domain m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) domain = Domain(mesh, dt, "CG", degree) + # Equation tracers = [WaterVapour(), CloudWater()] - if recovered: u_transport_option = "vector_advection_form" else: @@ -40,31 +43,10 @@ def setup_saturated(dirname, recovered): eqns = CompressibleEulerEquations( domain, parameters, u_transport_option=u_transport_option, active_tracers=tracers) + # I/O output = OutputParameters(dirname=dirname+'/saturated_balance', dumpfreq=1, dumplist=['u']) - diagnostic_fields = [Theta_e()] - io = IO(domain, eqns, output=output, diagnostic_fields=diagnostic_fields) - - # Initial conditions - rho0 = eqns.fields("rho") - theta0 = eqns.fields("theta") - water_v0 = eqns.fields("water_vapour") - water_c0 = eqns.fields("cloud_water") - moisture = ['water_vapour', 'cloud_water'] - - # spaces - Vt = theta0.function_space() - - # Isentropic background state - Tsurf = Constant(300.) - total_water = Constant(0.02) - theta_e = Function(Vt).interpolate(Tsurf) - water_t = Function(Vt).interpolate(total_water) - - # Calculate hydrostatic exner - saturated_hydrostatic_balance(eqns, theta_e, water_t) - water_c0.assign(water_t - water_v0) - - eqns.set_reference_profiles([('rho', rho0), ('theta', theta0)]) + diagnostic_fields = [Theta_e(eqns)] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) # Set up transport schemes if recovered: @@ -102,16 +84,41 @@ def setup_saturated(dirname, recovered): else: transported_fields.append(ImplicitMidpoint(domain, 'u')) - linear_solver = CompressibleSolver(eqns, moisture=moisture) + # Linear solver + linear_solver = CompressibleSolver(eqns) - # add physics + # Physics schemes physics_schemes = [(SaturationAdjustment(eqns), ForwardEuler(domain))] - # build time stepper + # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, linear_solver=linear_solver, physics_schemes=physics_schemes) + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + water_v0 = stepper.fields("water_vapour") + water_c0 = stepper.fields("cloud_water") + + # spaces + Vt = theta0.function_space() + + # Isentropic background state + Tsurf = Constant(300.) + total_water = Constant(0.02) + theta_e = Function(Vt).interpolate(Tsurf) + water_t = Function(Vt).interpolate(total_water) + + # Calculate hydrostatic exner + saturated_hydrostatic_balance(eqns, stepper.fields, theta_e, water_t) + water_c0.assign(water_t - water_v0) + + stepper.set_reference_profiles([('rho', rho0), ('theta', theta0)]) + return stepper, tmax diff --git a/integration-tests/balance/test_unsaturated_balance.py b/integration-tests/balance/test_unsaturated_balance.py index b9c59c209..e2377ceed 100644 --- a/integration-tests/balance/test_unsaturated_balance.py +++ b/integration-tests/balance/test_unsaturated_balance.py @@ -14,7 +14,11 @@ def setup_unsaturated(dirname, recovered): - # set up grid and time stepping parameters + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Parameters dt = 1. tmax = 3. deltax = 400 @@ -26,10 +30,12 @@ def setup_unsaturated(dirname, recovered): degree = 0 if recovered else 1 + # Domain m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) domain = Domain(mesh, dt, "CG", degree) + # Equation tracers = [WaterVapour(), CloudWater()] if recovered: @@ -40,30 +46,12 @@ def setup_unsaturated(dirname, recovered): eqns = CompressibleEulerEquations( domain, parameters, u_transport_option=u_transport_option, active_tracers=tracers) + # I/O output = OutputParameters(dirname=dirname+'/unsaturated_balance', dumpfreq=1) - diagnostic_fields = [Theta_d(), RelativeHumidity()] - io = IO(domain, eqns, output=output, diagnostic_fields=diagnostic_fields) - - # Initial conditions - rho0 = eqns.fields("rho") - theta0 = eqns.fields("theta") - moisture = ['water_vapour', 'cloud_water'] - - # spaces - Vt = theta0.function_space() - - # Isentropic background state - Tsurf = Constant(300.) - humidity = Constant(0.5) - theta_d = Function(Vt).interpolate(Tsurf) - RH = Function(Vt).interpolate(humidity) - - # Calculate hydrostatic exner - unsaturated_hydrostatic_balance(eqns, theta_d, RH) - - eqns.set_reference_profiles([('rho', rho0), ('theta', theta0)]) + diagnostic_fields = [Theta_d(eqns), RelativeHumidity(eqns)] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) - # Set up transport schemes + # Transport schemes if recovered: VDG1 = domain.spaces("DG1_equispaced") VCG1 = FunctionSpace(mesh, "CG", 1) @@ -91,16 +79,38 @@ def setup_unsaturated(dirname, recovered): else: transported_fields.append(ImplicitMidpoint(domain, "u")) - linear_solver = CompressibleSolver(eqns, moisture=moisture) + # Linear solver + linear_solver = CompressibleSolver(eqns) # Set up physics physics_schemes = [(SaturationAdjustment(eqns), ForwardEuler(domain))] - # build time stepper + # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, linear_solver=linear_solver, physics_schemes=physics_schemes) + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + + # spaces + Vt = theta0.function_space() + + # Isentropic background state + Tsurf = Constant(300.) + humidity = Constant(0.5) + theta_d = Function(Vt).interpolate(Tsurf) + RH = Function(Vt).interpolate(humidity) + + # Calculate hydrostatic exner + unsaturated_hydrostatic_balance(eqns, stepper.fields, theta_d, RH) + + stepper.set_reference_profiles([('rho', rho0), ('theta', theta0)]) + return stepper, tmax diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py index d31f16939..92ae1deca 100644 --- a/integration-tests/conftest.py +++ b/integration-tests/conftest.py @@ -9,7 +9,7 @@ from collections import namedtuple import pytest -opts = ('domain', 'tmax', 'output', 'f_init', 'f_end', 'degree', +opts = ('domain', 'tmax', 'io', 'f_init', 'f_end', 'degree', 'uexpr', 'umax', 'radius', 'tol') TracerSetup = namedtuple('TracerSetup', opts) TracerSetup.__new__.__defaults__ = (None,)*len(opts) @@ -30,6 +30,7 @@ def tracer_sphere(tmpdir, degree): dt = pi/3. * 0.02 output = OutputParameters(dirname=str(tmpdir), dumpfreq=15) domain = Domain(mesh, dt, family="BDM", degree=degree) + io = IO(domain, output) umax = 1.0 uexpr = as_vector([- umax * x[1] / radius, umax * x[0] / radius, 0.0]) @@ -40,7 +41,7 @@ def tracer_sphere(tmpdir, degree): tol = 0.05 - return TracerSetup(domain, tmax, output, f_init, f_end, degree, + return TracerSetup(domain, tmax, io, f_init, f_end, degree, uexpr, umax, radius, tol) @@ -57,6 +58,7 @@ def tracer_slice(tmpdir, degree): tmax = 0.75 output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) domain = Domain(mesh, dt, family="CG", degree=degree) + io = IO(domain, output) uexpr = as_vector([2.0, 0.0]) @@ -73,7 +75,7 @@ def tracer_slice(tmpdir, degree): tol = 0.12 - return TracerSetup(domain, tmax, output, f_init, f_end, degree, uexpr, tol=tol) + return TracerSetup(domain, tmax, io, f_init, f_end, degree, uexpr, tol=tol) def tracer_blob_slice(tmpdir, degree): @@ -84,12 +86,13 @@ def tracer_blob_slice(tmpdir, degree): output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) domain = Domain(mesh, dt, family="CG", degree=degree) + io = IO(domain, output) tmax = 1. x = SpatialCoordinate(mesh) f_init = exp(-((x[0]-0.5*L)**2 + (x[1]-0.5*L)**2)) - return TracerSetup(domain, tmax, output, f_init, degree=degree) + return TracerSetup(domain, tmax, io, f_init, degree=degree) @pytest.fixture() diff --git a/integration-tests/diffusion/test_diffusion.py b/integration-tests/diffusion/test_diffusion.py index 680330657..9cc43b38f 100644 --- a/integration-tests/diffusion/test_diffusion.py +++ b/integration-tests/diffusion/test_diffusion.py @@ -8,11 +8,9 @@ import pytest -def run(equation, diffusion_scheme, io, tmax): - - timestepper = Timestepper(equation, diffusion_scheme, io) +def run(timestepper, tmax): timestepper.run(0., tmax) - return timestepper.equation.fields("f") + return timestepper.fields("f") @pytest.mark.parametrize("DG", [True, False]) @@ -36,12 +34,12 @@ def test_scalar_diffusion(tmpdir, DG, tracer_setup): diffusion_params = DiffusionParameters(kappa=kappa, mu=mu) eqn = DiffusionEquation(domain, V, "f", diffusion_parameters=diffusion_params) - io = IO(domain, eqn, output=setup.output) - diffusion_scheme = BackwardEuler(domain) + timestepper = Timestepper(eqn, diffusion_scheme, setup.io) - eqn.fields("f").interpolate(f_init) - f_end = run(eqn, diffusion_scheme, io, tmax) + # Initial conditions + timestepper.fields("f").interpolate(f_init) + f_end = run(timestepper, tmax) assert errornorm(f_end_expr, f_end) < tol @@ -69,14 +67,14 @@ def test_vector_diffusion(tmpdir, DG, tracer_setup): diffusion_params = DiffusionParameters(kappa=kappa, mu=mu) eqn = DiffusionEquation(domain, V, "f", diffusion_parameters=diffusion_params) - io = IO(domain, eqn, output=setup.output) + diffusion_scheme = BackwardEuler(domain) + timestepper = Timestepper(eqn, diffusion_scheme, setup.io) + # Initial conditions if DG: - eqn.fields("f").interpolate(f_init) + timestepper.fields("f").interpolate(f_init) else: - eqn.fields("f").project(f_init) - - diffusion_scheme = BackwardEuler(domain) + timestepper.fields("f").project(f_init) - f_end = run(eqn, diffusion_scheme, io, tmax) + f_end = run(timestepper, tmax) assert errornorm(f_end_expr, f_end) < tol diff --git a/integration-tests/equations/test_advection_diffusion.py b/integration-tests/equations/test_advection_diffusion.py index b1fb096ba..6c91c1e35 100644 --- a/integration-tests/equations/test_advection_diffusion.py +++ b/integration-tests/equations/test_advection_diffusion.py @@ -10,12 +10,18 @@ def run_advection_diffusion(tmpdir): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain dt = 0.02 tmax = 1.0 L = 10 mesh = PeriodicIntervalMesh(20, L) domain = Domain(mesh, dt, "CG", 1) + # Equation diffusion_params = DiffusionParameters(kappa=0.75, mu=5) V = domain.spaces("DG", "DG", 1) Vu = VectorFunctionSpace(mesh, "CG", 1) @@ -23,11 +29,17 @@ def run_advection_diffusion(tmpdir): equation = AdvectionDiffusionEquation(domain, V, "f", Vu=Vu, diffusion_parameters=diffusion_params) + # I/O output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) - io = IO(domain, equation, output=output) + io = IO(domain, output) + # Time stepper + stepper = PrescribedTransport(equation, SSPRK3(domain), io) + # ------------------------------------------------------------------------ # # Initial conditions + # ------------------------------------------------------------------------ # + x = SpatialCoordinate(mesh) xc_init = 0.25*L xc_end = 0.75*L @@ -47,15 +59,18 @@ def run_advection_diffusion(tmpdir): f_init_expr = f_init*exp(-(x_init / f_width_init)**2) f_end_expr = f_end*exp(-(x_end / f_width_end)**2) - equation.fields('f').interpolate(f_init_expr) - equation.fields('u').interpolate(as_vector([Constant(umax)])) - f_end = equation.fields('f_end', V).interpolate(f_end_expr) + stepper.fields('f').interpolate(f_init_expr) + stepper.fields('u').interpolate(as_vector([Constant(umax)])) + f_end = stepper.fields('f_end', space=V) + f_end.interpolate(f_end_expr) - # Time stepper - timestepper = PrescribedTransport(equation, SSPRK3(domain), io) - timestepper.run(0, tmax=tmax) + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(0, tmax=tmax) - error = norm(equation.fields('f') - f_end) / norm(f_end) + error = norm(stepper.fields('f') - f_end) / norm(f_end) return error diff --git a/integration-tests/equations/test_dry_compressible.py b/integration-tests/equations/test_dry_compressible.py index cdeb221ba..5a016daae 100644 --- a/integration-tests/equations/test_dry_compressible.py +++ b/integration-tests/equations/test_dry_compressible.py @@ -12,6 +12,11 @@ def run_dry_compressible(tmpdir): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain dt = 6.0 tmax = 2*dt nlayers = 10 # horizontal layers @@ -22,18 +27,36 @@ def run_dry_compressible(tmpdir): mesh = ExtrudedMesh(m, layers=nlayers, layer_height=Lz/nlayers) domain = Domain(mesh, dt, "CG", 1) + # Equation parameters = CompressibleParameters() - R_d = parameters.R_d - g = parameters.g eqn = CompressibleEulerEquations(domain, parameters) + # I/O output = OutputParameters(dirname=tmpdir+"/dry_compressible", dumpfreq=2, chkptfreq=2) - io = IO(domain, eqn, output=output) + io = IO(domain, output) + + # Transport schemes + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta")] + # Linear solver + linear_solver = CompressibleSolver(eqn) + + # Time stepper + stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, + linear_solver=linear_solver) + + # ------------------------------------------------------------------------ # # Initial conditions - rho0 = eqn.fields("rho") - theta0 = eqn.fields("theta") + # ------------------------------------------------------------------------ # + + R_d = parameters.R_d + g = parameters.g + + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") # Approximate hydrostatic balance x, z = SpatialCoordinate(mesh) @@ -43,26 +66,17 @@ def run_dry_compressible(tmpdir): theta0.interpolate(tde.theta(parameters, T, p)) rho0.interpolate(p / (R_d * T)) - eqn.set_reference_profiles([('rho', rho0), ('theta', theta0)]) + stepper.set_reference_profiles([('rho', rho0), ('theta', theta0)]) # Add perturbation r = sqrt((x-Lx/2)**2 + (z-Lz/2)**2) theta_pert = 1.0*exp(-(r/(Lx/5))**2) theta0.interpolate(theta0 + theta_pert) - # Set up transport schemes - transported_fields = [ImplicitMidpoint(domain, "u"), - SSPRK3(domain, "rho"), - SSPRK3(domain, "theta")] - - # Set up linear solver for the timestepping scheme - linear_solver = CompressibleSolver(eqn) - - # build time stepper - stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, - linear_solver=linear_solver) - + # ------------------------------------------------------------------------ # # Run + # ------------------------------------------------------------------------ # + stepper.run(t=0, tmax=tmax) # IO for checking checkpoints @@ -71,22 +85,22 @@ def run_dry_compressible(tmpdir): check_eqn = CompressibleEulerEquations(domain, parameters) check_output = OutputParameters(dirname=tmpdir+"/dry_compressible", checkpoint_pickup_filename=new_path) - check_io = IO(domain, check_eqn, output=check_output) - check_eqn.set_reference_profiles([]) + check_io = IO(domain, check_output) check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, []) + check_stepper.set_reference_profiles([]) check_stepper.run(t=0, tmax=0, pickup=True) - return eqn, check_eqn + return stepper, check_stepper def test_dry_compressible(tmpdir): dirname = str(tmpdir) - eqn, check_eqn = run_dry_compressible(dirname) + stepper, check_stepper = run_dry_compressible(dirname) for variable in ['u', 'rho', 'theta']: - new_variable = eqn.fields(variable) - check_variable = check_eqn.fields(variable) + new_variable = stepper.fields(variable) + check_variable = check_stepper.fields(variable) error = norm(new_variable - check_variable) / norm(check_variable) # Slack values chosen to be robust to different platforms diff --git a/integration-tests/equations/test_forced_advection.py b/integration-tests/equations/test_forced_advection.py index 420a08726..e9e4574b9 100644 --- a/integration-tests/equations/test_forced_advection.py +++ b/integration-tests/equations/test_forced_advection.py @@ -8,13 +8,18 @@ """ from gusto import * -from firedrake import (PeriodicIntervalMesh, SpatialCoordinate, FunctionSpace, - VectorFunctionSpace, conditional, acos, cos, pi, +from firedrake import (PeriodicIntervalMesh, SpatialCoordinate, + VectorFunctionSpace, conditional, acos, cos, pi, sin, as_vector, errornorm) def run_forced_advection(tmpdir): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain Lx = 100 delta_x = 2.0 nx = int(Lx/delta_x) @@ -24,10 +29,10 @@ def run_forced_advection(tmpdir): dt = 0.2 domain = Domain(mesh, dt, "CG", 1) - VD = FunctionSpace(mesh, "DG", 1) + VD = domain.spaces("DG") Vu = VectorFunctionSpace(mesh, "CG", 1) - # set up parameters and initial conditions + # Equation u_max = 1 C0 = 0.6 K0 = 0.3 @@ -40,27 +45,34 @@ def run_forced_advection(tmpdir): msat = Function(VD) msat.interpolate(msat_expr) - # initial moisture profile - mexpr = C0 + K0*cos((2*pi*x)/Lx) - rain = Rain(space='tracer', transport_eqn=TransportEquationType.no_transport) meqn = ForcedAdvectionEquation(domain, VD, field_name="water_vapour", Vu=Vu, active_tracers=[rain]) physics_schemes = [(InstantRain(meqn, msat, rain_name="rain", - set_tau_to_dt=True), ForwardEuler(domain))] + set_tau_to_dt=True, parameters=None), ForwardEuler(domain))] + # I/O output = OutputParameters(dirname=str(tmpdir), dumpfreq=1) diagnostic_fields = [CourantNumber()] - io = IO(domain, meqn, output=output, diagnostic_fields=diagnostic_fields) + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Time Stepper + stepper = PrescribedTransport(meqn, RK4(domain), io, + physics_schemes=physics_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + # initial moisture profile + mexpr = C0 + K0*cos((2*pi*x)/Lx) - meqn.fields("u").project(as_vector([u_max])) - qv = meqn.fields("water_vapour") - qv.project(mexpr) + stepper.fields("u").project(as_vector([u_max])) + stepper.fields("water_vapour").project(mexpr) # exact rainfall profile (analytically) - r_exact = meqn.fields("r_exact", VD) + r_exact = stepper.fields("r_exact", space=VD) lim1 = Lx/(2*pi) * acos((C0 + K0 - Csat)/Ksat) lim2 = Lx/2 coord = (Ksat*cos(2*pi*x/Lx) + Csat - C0)/K0 @@ -68,13 +80,13 @@ def run_forced_advection(tmpdir): r_expr = conditional(x < lim2, conditional(x > lim1, exact_expr, 0), 0) r_exact.interpolate(r_expr) - # build time stepper - stepper = PrescribedTransport(meqn, RK4(domain), io, - physics_schemes=physics_schemes) + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # stepper.run(0, tmax=tmax) - error = errornorm(r_exact, meqn.fields("rain")) + error = errornorm(r_exact, stepper.fields("rain")) return error diff --git a/integration-tests/equations/test_incompressible.py b/integration-tests/equations/test_incompressible.py index c1be5e5c5..51c6ac9f4 100644 --- a/integration-tests/equations/test_incompressible.py +++ b/integration-tests/equations/test_incompressible.py @@ -11,6 +11,11 @@ def run_incompressible(tmpdir): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain dt = 6.0 tmax = 2*dt nlayers = 10 # horizontal layers @@ -21,16 +26,33 @@ def run_incompressible(tmpdir): mesh = ExtrudedMesh(m, layers=nlayers, layer_height=Lz/nlayers) domain = Domain(mesh, dt, "CG", 1) + # Equation parameters = CompressibleParameters() eqn = IncompressibleBoussinesqEquations(domain, parameters) + # I/O output = OutputParameters(dirname=tmpdir+"/incompressible", dumpfreq=2, chkptfreq=2) - io = IO(domain, eqn, output=output) + io = IO(domain, output) + + # Transport Schemes + b_opts = SUPGOptions() + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "b", options=b_opts)] + # Linear solver + linear_solver = IncompressibleSolver(eqn) + + # Time stepper + stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, + linear_solver=linear_solver) + + # ------------------------------------------------------------------------ # # Initial conditions - p0 = eqn.fields("p") - b0 = eqn.fields("b") + # ------------------------------------------------------------------------ # + + p0 = stepper.fields("p") + b0 = stepper.fields("b") # z.grad(bref) = N**2 x, z = SpatialCoordinate(mesh) @@ -39,50 +61,41 @@ def run_incompressible(tmpdir): b_b = Function(b0.function_space()).interpolate(bref) incompressible_hydrostatic_balance(eqn, b_b, p0) - eqn.set_reference_profiles([('p', p0), ('b', b_b)]) + stepper.set_reference_profiles([('p', p0), ('b', b_b)]) # Add perturbation r = sqrt((x-Lx/2)**2 + (z-Lz/2)**2) b_pert = 0.1*exp(-(r/(Lx/5)**2)) b0.interpolate(b_b + b_pert) - # Set up transport schemes - b_opts = SUPGOptions() - transported_fields = [ImplicitMidpoint(domain, "u"), - SSPRK3(domain, "b", options=b_opts)] - - # Set up linear solver for the timestepping scheme - linear_solver = IncompressibleSolver(eqn) - - # build time stepper - stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, - linear_solver=linear_solver) - + # ------------------------------------------------------------------------ # # Run + # ------------------------------------------------------------------------ # + stepper.run(t=0, tmax=tmax) # State for checking checkpoints checkpoint_name = 'incompressible_chkpt' new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') check_eqn = IncompressibleBoussinesqEquations(domain, parameters) - check_eqn.set_reference_profiles([]) check_output = OutputParameters(dirname=tmpdir+"/incompressible", checkpoint_pickup_filename=new_path) - check_io = IO(domain, check_eqn, output=check_output) + check_io = IO(domain, check_output) check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, []) + check_stepper.set_reference_profiles([]) check_stepper.run(t=0, tmax=0, pickup=True) - return eqn, check_eqn + return stepper, check_stepper def test_incompressible(tmpdir): dirname = str(tmpdir) - eqn, check_eqn = run_incompressible(dirname) + stepper, check_stepper = run_incompressible(dirname) for variable in ['u', 'b', 'p']: - new_variable = eqn.fields(variable) - check_variable = check_eqn.fields(variable) + new_variable = stepper.fields(variable) + check_variable = check_stepper.fields(variable) error = norm(new_variable - check_variable) / norm(check_variable) # Slack values chosen to be robust to different platforms diff --git a/integration-tests/equations/test_moist_compressible.py b/integration-tests/equations/test_moist_compressible.py index 2b9e71ddc..c51852e5e 100644 --- a/integration-tests/equations/test_moist_compressible.py +++ b/integration-tests/equations/test_moist_compressible.py @@ -12,6 +12,11 @@ def run_moist_compressible(tmpdir): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain dt = 6.0 tmax = 2*dt nlayers = 10 # horizontal layers @@ -22,22 +27,39 @@ def run_moist_compressible(tmpdir): mesh = ExtrudedMesh(m, layers=nlayers, layer_height=Lz/nlayers) domain = Domain(mesh, dt, "CG", 1) + # Equation parameters = CompressibleParameters() - R_d = parameters.R_d - R_v = parameters.R_v - g = parameters.g - tracers = [WaterVapour(name='vapour_mixing_ratio'), CloudWater(name='cloud_liquid_mixing_ratio')] eqn = CompressibleEulerEquations(domain, parameters, active_tracers=tracers) + # I/O output = OutputParameters(dirname=tmpdir+"/moist_compressible", dumpfreq=2, chkptfreq=2) - io = IO(domain, eqn, output=output) + io = IO(domain, output) + + # Transport schemes + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta")] + # Linear solver + linear_solver = CompressibleSolver(eqn) + + # Time stepper + stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, + linear_solver=linear_solver) + + # ------------------------------------------------------------------------ # # Initial conditions - rho0 = eqn.fields("rho") - theta0 = eqn.fields("theta") - m_v0 = eqn.fields("vapour_mixing_ratio") + # ------------------------------------------------------------------------ # + + R_d = parameters.R_d + R_v = parameters.R_v + g = parameters.g + + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + m_v0 = stepper.fields("vapour_mixing_ratio") # Approximate hydrostatic balance x, z = SpatialCoordinate(mesh) @@ -49,50 +71,42 @@ def run_moist_compressible(tmpdir): theta0.interpolate(tde.theta(parameters, T_vd, p)) rho0.interpolate(p / (R_d * T)) - eqn.set_reference_profiles([('rho', rho0), ('theta', theta0)]) + stepper.set_reference_profiles([('rho', rho0), ('theta', theta0), + ('vapour_mixing_ratio', m_v0)]) # Add perturbation r = sqrt((x-Lx/2)**2 + (z-Lz/2)**2) theta_pert = 1.0*exp(-(r/(Lx/5))**2) theta0.interpolate(theta0 + theta_pert) - # Set up transport schemes - transported_fields = [ImplicitMidpoint(domain, "u"), - SSPRK3(domain, "rho"), - SSPRK3(domain, "theta")] - - # Set up linear solver for the timestepping scheme - linear_solver = CompressibleSolver(eqn, moisture=['vapour_mixing_ratio']) - - # build time stepper - stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, - linear_solver=linear_solver) - + # ------------------------------------------------------------------------ # # Run + # ------------------------------------------------------------------------ # + stepper.run(t=0, tmax=tmax) # State for checking checkpoints checkpoint_name = 'moist_compressible_chkpt' new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') check_eqn = CompressibleEulerEquations(domain, parameters, active_tracers=tracers) - check_eqn.set_reference_profiles([]) check_output = OutputParameters(dirname=tmpdir+"/moist_compressible", checkpoint_pickup_filename=new_path) - check_io = IO(domain, check_eqn, output=check_output) + check_io = IO(domain, output=check_output) check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, []) + check_stepper.set_reference_profiles([]) check_stepper.run(t=0, tmax=0, pickup=True) - return eqn, check_eqn + return stepper, check_stepper def test_moist_compressible(tmpdir): dirname = str(tmpdir) - eqn, check_eqn = run_moist_compressible(dirname) + stepper, check_stepper = run_moist_compressible(dirname) for variable in ['u', 'rho', 'theta', 'vapour_mixing_ratio']: - new_variable = eqn.fields(variable) - check_variable = check_eqn.fields(variable) + new_variable = stepper.fields(variable) + check_variable = check_stepper.fields(variable) error = norm(new_variable - check_variable) / norm(check_variable) # Slack values chosen to be robust to different platforms diff --git a/integration-tests/equations/test_sw_linear_triangle.py b/integration-tests/equations/test_sw_linear_triangle.py index 6ad9f18b4..908e182b7 100644 --- a/integration-tests/equations/test_sw_linear_triangle.py +++ b/integration-tests/equations/test_sw_linear_triangle.py @@ -12,6 +12,11 @@ def setup_sw(dirname): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain dt = 3600. refinements = 3 # number of horizontal cells = 20*(4^refinements) @@ -26,19 +31,29 @@ def setup_sw(dirname): domain = Domain(mesh, dt, "BDM", degree=1) - # Coriolis + # Equation parameters = ShallowWaterParameters(H=H) Omega = parameters.Omega fexpr = 2*Omega*x[2]/R eqns = LinearShallowWaterEquations(domain, parameters, fexpr=fexpr) - output = OutputParameters(dirname=dirname+"/sw_linear_w2", steady_state_error_fields=['u', 'D'], dumpfreq=12) - io = IO(domain, eqns, output=output) + # I/O + diagnostic_fields = [SteadyStateError('u'), SteadyStateError('D')] + output = OutputParameters(dirname=dirname+"/sw_linear_w2", dumpfreq=12) + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes + transport_schemes = [ForwardEuler(domain, "D")] + + # Time stepper + stepper = SemiImplicitQuasiNewton(eqns, io, transport_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # - # interpolate initial conditions - # Initial/current conditions - u0 = eqns.fields("u") - D0 = eqns.fields("D") + u0 = stepper.fields("u") + D0 = stepper.fields("D") u_max = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s) uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) g = parameters.g @@ -47,12 +62,7 @@ def setup_sw(dirname): D0.interpolate(Dexpr) Dbar = Function(D0.function_space()).assign(H) - eqns.set_reference_profiles([('D', Dbar)]) - - transport_schemes = [ForwardEuler(domain, "D")] - - # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, io, transport_schemes) + stepper.set_reference_profiles([('D', Dbar)]) return stepper, 2*day diff --git a/integration-tests/equations/test_sw_triangle.py b/integration-tests/equations/test_sw_triangle.py index 406427ac0..d2ce2303a 100644 --- a/integration-tests/equations/test_sw_triangle.py +++ b/integration-tests/equations/test_sw_triangle.py @@ -22,6 +22,11 @@ def setup_sw(dirname, dt, u_transport_option): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain refinements = 3 # number of horizontal cells = 20*(4^refinements) mesh = IcosahedralSphereMesh(radius=R, @@ -30,12 +35,14 @@ def setup_sw(dirname, dt, u_transport_option): x = SpatialCoordinate(mesh) mesh.init_cell_orientations(x) + # Equation parameters = ShallowWaterParameters(H=H) Omega = parameters.Omega fexpr = 2*Omega*x[2]/R eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, u_transport_option=u_transport_option) + # I/O diagnostic_fields = [RelativeVorticity(), AbsoluteVorticity(), PotentialVorticity(), ShallowWaterPotentialEnstrophy('RelativeVorticity'), @@ -53,36 +60,44 @@ def setup_sw(dirname, dt, u_transport_option): 'SWPotentialEnstrophy_from_AbsoluteVorticity'), MeridionalComponent('u'), ZonalComponent('u'), - RadialComponent('u')] - output = OutputParameters(dirname=dirname+"/sw", dumplist_latlon=['D', 'D_error'], steady_state_error_fields=['D', 'u']) - io = IO(domain, eqns, output=output, diagnostic_fields=diagnostic_fields) + RadialComponent('u'), + SteadyStateError('D'), + SteadyStateError('u')] + output = OutputParameters(dirname=dirname+"/sw", dumplist_latlon=['D', 'D_error']) + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + return domain, eqns, io + + +def set_up_initial_conditions(domain, equation, stepper): + + x = SpatialCoordinate(domain.mesh) # interpolate initial conditions - u0 = eqns.fields("u") - D0 = eqns.fields("D") + u0 = stepper.fields("u") + D0 = stepper.fields("D") uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) - g = parameters.g + g = equation.parameters.g + Omega = equation.parameters.Omega Dexpr = H - ((R * Omega * u_max + u_max*u_max/2.0)*(x[2]*x[2]/(R*R)))/g u0.project(uexpr) D0.interpolate(Dexpr) Dbar = Function(D0.function_space()).assign(H) - eqns.set_reference_profiles([('D', Dbar)]) + stepper.set_reference_profiles([('D', Dbar)]) vspace = FunctionSpace(domain.mesh, "CG", 3) vexpr = (2*u_max/R)*x[2]/R - # TODO: these fields should not be in eqns - f = eqns.fields("coriolis") - vrel_analytical = eqns.fields("AnalyticalRelativeVorticity", vspace) + + f = stepper.fields("coriolis") + vrel_analytical = stepper.fields("AnalyticalRelativeVorticity", space=vspace) vrel_analytical.interpolate(vexpr) - vabs_analytical = eqns.fields("AnalyticalAbsoluteVorticity", vspace) + vabs_analytical = stepper.fields("AnalyticalAbsoluteVorticity", space=vspace) vabs_analytical.interpolate(vexpr + f) - pv_analytical = eqns.fields("AnalyticalPotentialVorticity", vspace) + pv_analytical = stepper.fields("AnalyticalPotentialVorticity", space=vspace) pv_analytical.interpolate((vexpr+f)/D0) - return domain, eqns, io - def check_results(dirname): filename = path.join(dirname, "sw/diagnostics.nc") @@ -140,10 +155,18 @@ def test_sw_setup(tmpdir, u_transport_option): dt = 1500 domain, eqns, io = setup_sw(dirname, dt, u_transport_option) + # Transport schemes transported_fields = [] transported_fields.append((ImplicitMidpoint(domain, "u"))) transported_fields.append((SSPRK3(domain, "D"))) + + # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields) + + # Initial conditions + set_up_initial_conditions(domain, eqns, stepper) + + # Run stepper.run(t=0, tmax=0.25*day) check_results(dirname) @@ -159,6 +182,11 @@ def test_sw_ssprk3(tmpdir, u_transport_option): domain, eqns, io = setup_sw(dirname, dt, u_transport_option) stepper = Timestepper(eqns, SSPRK3(domain), io) + + # Initial conditions + set_up_initial_conditions(domain, eqns, stepper) + + # Run stepper.run(t=0, tmax=0.01*day) check_results(dirname) diff --git a/integration-tests/model/test_checkpointing.py b/integration-tests/model/test_checkpointing.py index 694aebe32..e478aeb4e 100644 --- a/integration-tests/model/test_checkpointing.py +++ b/integration-tests/model/test_checkpointing.py @@ -27,9 +27,7 @@ def setup_checkpointing(dirname): output = OutputParameters(dirname=dirname, dumpfreq=1, chkptfreq=2, log_level='INFO') - io = IO(domain, eqns, output=output) - - initialise_fields(eqns) + io = IO(domain, output) # Set up transport schemes transported_fields = [] @@ -44,18 +42,20 @@ def setup_checkpointing(dirname): stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, linear_solver=linear_solver) - return eqns, stepper, dt + initialise_fields(eqns, stepper) + + return stepper, dt -def initialise_fields(eqns): +def initialise_fields(eqns, stepper): L = 1.e5 H = 1.0e4 # Height position of the model top # Initial conditions - u0 = eqns.fields("u") - rho0 = eqns.fields("rho") - theta0 = eqns.fields("theta") + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") # spaces Vt = theta0.function_space() @@ -84,15 +84,15 @@ def initialise_fields(eqns): rho0.assign(rho_b) u0.project(as_vector([20.0, 0.0])) - eqns.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) + stepper.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) def test_checkpointing(tmpdir): dirname_1 = str(tmpdir)+'/checkpointing_1' dirname_2 = str(tmpdir)+'/checkpointing_2' - eqns_1, stepper_1, dt = setup_checkpointing(dirname_1) - eqns_2, stepper_2, dt = setup_checkpointing(dirname_2) + stepper_1, dt = setup_checkpointing(dirname_1) + stepper_2, dt = setup_checkpointing(dirname_2) # ------------------------------------------------------------------------ # # Run for 4 time steps and store values @@ -107,9 +107,9 @@ def test_checkpointing(tmpdir): stepper_2.run(t=0.0, tmax=2*dt) # Wipe fields, then pickup - eqns_2.fields('u').project(as_vector([-10.0, 0.0])) - eqns_2.fields('rho').interpolate(Constant(0.0)) - eqns_2.fields('theta').interpolate(Constant(0.0)) + stepper_2.fields('u').project(as_vector([-10.0, 0.0])) + stepper_2.fields('rho').interpolate(Constant(0.0)) + stepper_2.fields('theta').interpolate(Constant(0.0)) stepper_2.run(t=2*dt, tmax=4*dt, pickup=True) @@ -121,14 +121,14 @@ def test_checkpointing(tmpdir): # This is the best way to compare fields from different meshes for field_name in ['u', 'rho', 'theta']: with DumbCheckpoint(dirname_1+'/chkpt', mode=FILE_READ) as chkpt: - field_1 = Function(eqns_1.fields(field_name).function_space(), + field_1 = Function(stepper_1.fields(field_name).function_space(), name=field_name) chkpt.load(field_1) # These are preserved in the comments for when we can use CheckpointFile # mesh = chkpt.load_mesh(name='firedrake_default_extruded') # field_1 = chkpt.load_function(mesh, name=field_name) with DumbCheckpoint(dirname_2+'/chkpt', mode=FILE_READ) as chkpt: - field_2 = Function(eqns_1.fields(field_name).function_space(), + field_2 = Function(stepper_1.fields(field_name).function_space(), name=field_name) chkpt.load(field_2) # These are preserved in the comments for when we can use CheckpointFile diff --git a/integration-tests/model/test_passive_tracer.py b/integration-tests/model/test_passive_tracer.py index 08eb82e27..63f1f5d5f 100644 --- a/integration-tests/model/test_passive_tracer.py +++ b/integration-tests/model/test_passive_tracer.py @@ -15,10 +15,14 @@ def run_tracer(setup): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + # Get initial conditions from shared config domain = setup.domain mesh = domain.mesh - output = setup.output + io = setup.io x = SpatialCoordinate(mesh) H = 0.1 @@ -29,17 +33,29 @@ def run_tracer(setup): R = setup.radius fexpr = 2*Omega*x[2]/R - # Need to create a new state containing parameters - # Equations eqns = LinearShallowWaterEquations(domain, parameters, fexpr=fexpr) tracer_eqn = AdvectionEquation(domain, domain.spaces("DG"), "tracer") - io = IO(domain, eqns, output=output) + + # set up transport schemes + transport_schemes = [ForwardEuler(domain, "D")] + + # Set up tracer transport + tracer_transport = [(tracer_eqn, SSPRK3(domain))] + + # build time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transport_schemes, + auxiliary_equations_and_schemes=tracer_transport) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # # Specify initial prognostic fields - u0 = eqns.fields("u") - D0 = eqns.fields("D") - tracer0 = tracer_eqn.fields("tracer", D0.function_space()) + u0 = stepper.fields("u") + D0 = stepper.fields("D") + tracer0 = stepper.fields("tracer") tracer_end = Function(D0.function_space()) # Expressions for initial fields corresponding to Williamson 2 test case @@ -50,22 +66,11 @@ def run_tracer(setup): tracer0.interpolate(setup.f_init) tracer_end.interpolate(setup.f_end) - eqns.set_reference_profiles([('D', Dbar)]) - - # set up transport schemes - transport_schemes = [ForwardEuler(domain, "D")] - - # Set up tracer transport - tracer_transport = [(tracer_eqn, SSPRK3(domain))] - - # build time stepper - stepper = SemiImplicitQuasiNewton( - eqns, io, transport_schemes, - auxiliary_equations_and_schemes=tracer_transport) + stepper.set_reference_profiles([('D', Dbar)]) stepper.run(t=0, tmax=setup.tmax) - error = norm(tracer_eqn.fields("tracer") - tracer_end) / norm(tracer_end) + error = norm(stepper.fields("tracer") - tracer_end) / norm(tracer_end) return error diff --git a/integration-tests/model/test_prescribed_transport.py b/integration-tests/model/test_prescribed_transport.py index 754f6880b..758a6349d 100644 --- a/integration-tests/model/test_prescribed_transport.py +++ b/integration-tests/model/test_prescribed_transport.py @@ -8,11 +8,9 @@ from firedrake import sin, cos, norm, pi, as_vector -def run(eqn, transport_scheme, io, tmax, f_end, prescribed_u): - timestepper = PrescribedTransport(eqn, transport_scheme, io, - prescribed_transporting_velocity=prescribed_u) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(eqn.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) def test_prescribed_transport_setup(tmpdir, tracer_setup): @@ -26,20 +24,22 @@ def test_prescribed_transport_setup(tmpdir, tracer_setup): V = domain.spaces("DG") # Make equation eqn = AdvectionEquation(domain, V, "f") - io = IO(domain, eqn, output=setup.output) # Initialise fields def u_evaluation(t): return as_vector([2.0*cos(2*pi*t/setup.tmax), sin(2*pi*t/setup.tmax)*sin(pi*z)]) - eqn.fields("f").interpolate(setup.f_init) - eqn.fields("u").project(u_evaluation(Constant(0.0))) - transport_scheme = SSPRK3(domain) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, + prescribed_transporting_velocity=u_evaluation) + + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(u_evaluation(Constant(0.0))) + # Run and check error - error = run(eqn, transport_scheme, io, setup.tmax, - setup.f_init, u_evaluation) + error = run(timestepper, setup.tmax, setup.f_init) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py index 0ba151ef7..0fe7c7f12 100644 --- a/integration-tests/model/test_time_discretisation.py +++ b/integration-tests/model/test_time_discretisation.py @@ -3,10 +3,9 @@ import pytest -def run(eqn, transport_scheme, io, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, io) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(eqn.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("scheme", ["ssprk", "implicit_midpoint", @@ -18,10 +17,6 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup): V = domain.spaces("DG") eqn = AdvectionEquation(domain, V, "f") - io = IO(domain, eqn, output=setup.output) - - eqn.fields("f").interpolate(setup.f_init) - eqn.fields("u").project(setup.uexpr) if scheme == "ssprk": transport_scheme = SSPRK3(domain) @@ -33,4 +28,11 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup): transport_scheme = Heun(domain) elif scheme == "BDF2": transport_scheme = BDF2(domain) - assert run(eqn, transport_scheme, io, setup.tmax, setup.f_end) < setup.tol + + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) + + assert run(timestepper, setup.tmax, setup.f_end) < setup.tol diff --git a/integration-tests/physics/test_condensation.py b/integration-tests/physics/test_condensation.py index d48b514df..0b92d6f11 100644 --- a/integration-tests/physics/test_condensation.py +++ b/integration-tests/physics/test_condensation.py @@ -16,6 +16,10 @@ def run_cond_evap(dirname, process): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + dt = 2.0 # declare grid shape, with length L and height H @@ -24,33 +28,46 @@ def run_cond_evap(dirname, process): nlayers = int(H / 10.) ncolumns = int(L / 10.) - # make mesh + # make mesh and domain m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) - domain = Domain(mesh, dt, "CG", 1) x, z = SpatialCoordinate(mesh) # spaces - Vt = domain.spaces("theta", degree=1) - Vr = domain.spaces("DG", "DG", degree=1) - # Set up equation -- use compressible to set up these spaces + # Set up equation tracers = [WaterVapour(), CloudWater()] parameters = CompressibleParameters() eqn = CompressibleEulerEquations(domain, parameters, active_tracers=tracers) + # I/O output = OutputParameters(dirname=dirname+"/cond_evap", dumpfreq=1, dumplist=['u']) - io = IO(domain, eqn, output=output, diagnostic_fields=[Sum('water_vapour', 'cloud_water')]) + io = IO(domain, output, diagnostic_fields=[Sum('water_vapour', 'cloud_water')]) + + # Physics scheme + physics_schemes = [(SaturationAdjustment(eqn, parameters=parameters), ForwardEuler(domain))] + + # Time stepper + scheme = ForwardEuler(domain) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, + physics_schemes=physics_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + Vt = domain.spaces("theta", degree=1) + Vr = domain.spaces("DG", "DG", degree=1) # Declare prognostic fields - rho0 = eqn.fields("rho") - theta0 = eqn.fields("theta") - water_v0 = eqn.fields("water_vapour", Vt) - water_c0 = eqn.fields("cloud_water", Vt) + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + water_v0 = stepper.fields("water_vapour") + water_c0 = stepper.fields("cloud_water") # Set a background state with constant pressure and temperature pressure = Function(Vr).interpolate(Constant(100000.)) @@ -90,27 +107,24 @@ def run_cond_evap(dirname, process): eqn.residual = eqn.residual.label_map(lambda t: t.has_label(time_derivative), map_if_true=identity, map_if_false=drop) - physics_schemes = [(SaturationAdjustment(eqn, parameters=parameters), ForwardEuler(domain))] - - # build time stepper - scheme = ForwardEuler(domain) - stepper = SplitPhysicsTimestepper(eqn, scheme, io, - physics_schemes=physics_schemes) + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # stepper.run(t=0, tmax=dt) - return eqn, mv_true, mc_true, theta_d_true, mc_init + return eqn, stepper, mv_true, mc_true, theta_d_true, mc_init @pytest.mark.parametrize("process", ["evaporation", "condensation"]) def test_cond_evap(tmpdir, process): dirname = str(tmpdir) - eqn, mv_true, mc_true, theta_d_true, mc_init = run_cond_evap(dirname, process) + eqn, stepper, mv_true, mc_true, theta_d_true, mc_init = run_cond_evap(dirname, process) - water_v = eqn.fields('water_vapour') - water_c = eqn.fields('cloud_water') - theta_vd = eqn.fields('theta') + water_v = stepper.fields('water_vapour') + water_c = stepper.fields('cloud_water') + theta_vd = stepper.fields('theta') theta_d = Function(theta_vd.function_space()) theta_d.interpolate(theta_vd/(1 + water_v * eqn.parameters.R_v / eqn.parameters.R_d)) diff --git a/integration-tests/physics/test_instant_rain.py b/integration-tests/physics/test_instant_rain.py index c35cca902..bfb65cd70 100644 --- a/integration-tests/physics/test_instant_rain.py +++ b/integration-tests/physics/test_instant_rain.py @@ -14,7 +14,11 @@ def run_instant_rain(dirname): - # set up mesh + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # set up mesh and domain L = 10 nx = 10 mesh = PeriodicSquareMesh(nx, nx, L) @@ -27,6 +31,7 @@ def run_instant_rain(dirname): g = 10 fexpr = Constant(0) + # Equation vapour = WaterVapour(name="water_vapour", space='DG') rain = Rain(name="rain", space="DG", transport_eqn=TransportEquationType.no_transport) @@ -35,13 +40,28 @@ def run_instant_rain(dirname): eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, active_tracers=[vapour, rain]) + # I/O output = OutputParameters(dirname=dirname+"/instant_rain", dumpfreq=1, dumplist=['vapour', "rain"]) diagnostic_fields = [CourantNumber()] - io = IO(domain, eqns, output=output, diagnostic_fields=diagnostic_fields) + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Physics schemes + # define saturation function + saturation = Constant(0.5) + physics_schemes = [(InstantRain(eqns, saturation, rain_name="rain", + set_tau_to_dt=True), ForwardEuler(domain))] + + # Time stepper + stepper = PrescribedTransport(eqns, RK4(domain), io, + physics_schemes=physics_schemes) - vapour0 = eqns.fields("water_vapour") + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + vapour0 = stepper.fields("water_vapour") # set up vapour xc = L/2 @@ -55,29 +75,24 @@ def run_instant_rain(dirname): VD = FunctionSpace(mesh, "DG", 1) initial_vapour = Function(VD).interpolate(vapour_expr) - # define saturation function - saturation = Constant(0.5) - # define expected solutions; vapour should be equal to saturation and rain # should be (initial vapour - saturation) vapour_true = Function(VD).interpolate(saturation) rain_true = Function(VD).interpolate(vapour0 - saturation) - physics_schemes = [(InstantRain(eqns, saturation, rain_name="rain", - set_tau_to_dt=True), ForwardEuler(domain))] - - stepper = PrescribedTransport(eqns, RK4(domain), io, - physics_schemes=physics_schemes) + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # stepper.run(t=0, tmax=5*dt) - return eqns, saturation, initial_vapour, vapour_true, rain_true + return stepper, saturation, initial_vapour, vapour_true, rain_true def test_instant_rain_setup(tmpdir): dirname = str(tmpdir) - eqns, saturation, initial_vapour, vapour_true, rain_true = run_instant_rain(dirname) - v = eqns.fields("water_vapour") - r = eqns.fields("rain") + stepper, saturation, initial_vapour, vapour_true, rain_true = run_instant_rain(dirname) + v = stepper.fields("water_vapour") + r = stepper.fields("rain") # check that the maximum of the vapour field is equal to the saturation assert v.dat.data.max() - saturation.values() < 0.001, "The maximum of the final vapour field should be equal to saturation" diff --git a/integration-tests/physics/test_precipitation.py b/integration-tests/physics/test_precipitation.py index 995908069..aded26c70 100644 --- a/integration-tests/physics/test_precipitation.py +++ b/integration-tests/physics/test_precipitation.py @@ -13,6 +13,10 @@ def setup_fallout(dirname): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + # declare grid shape, with length L and height H dt = 0.1 L = 10. @@ -20,25 +24,36 @@ def setup_fallout(dirname): nlayers = 10 ncolumns = 10 - # make mesh + # Domain m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) domain = Domain(mesh, dt, "CG", 1) x = SpatialCoordinate(mesh) + # Equation Vrho = domain.spaces("DG1_equispaced") active_tracers = [Rain(space='DG1_equispaced')] eqn = ForcedAdvectionEquation(domain, Vrho, "rho", active_tracers=active_tracers) + # I/O output = OutputParameters(dirname=dirname+"/fallout", dumpfreq=10, dumplist=['rain']) diagnostic_fields = [Precipitation()] - io = IO(domain, eqn, output=output, diagnostic_fields=diagnostic_fields) + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Physics schemes + physics_schemes = [(Fallout(eqn, 'rain', domain), SSPRK3(domain, 'rain'))] + # build time stepper scheme = ForwardEuler(domain) - eqn.fields("rho").assign(1.) + stepper = PrescribedTransport(eqn, scheme, io, + physics_schemes=physics_schemes) - physics_schemes = [(Fallout(eqn, 'rain', domain), SSPRK3(domain, 'rain'))] - rain0 = eqn.fields("rain") + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + stepper.fields("rho").assign(1.) + rain0 = stepper.fields("rain") # set up rain xc = L / 2 @@ -49,10 +64,6 @@ def setup_fallout(dirname): rain0.interpolate(rain_expr) - # build time stepper - stepper = PrescribedTransport(eqn, scheme, io, - physics_schemes=physics_schemes) - return stepper, 10.0 diff --git a/integration-tests/transport/test_dg_transport.py b/integration-tests/transport/test_dg_transport.py index 4d4003102..9b4abe75a 100644 --- a/integration-tests/transport/test_dg_transport.py +++ b/integration-tests/transport/test_dg_transport.py @@ -8,10 +8,9 @@ import pytest -def run(eqn, transport_scheme, io, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, io) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(eqn.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("geometry", ["slice", "sphere"]) @@ -26,12 +25,15 @@ def test_dg_transport_scalar(tmpdir, geometry, equation_form, tracer_setup): else: eqn = ContinuityEquation(domain, V, "f") - io = IO(domain, eqn, output=setup.output) - eqn.fields("f").interpolate(setup.f_init) - eqn.fields("u").project(setup.uexpr) - transport_scheme = SSPRK3(domain) - error = run(eqn, transport_scheme, io, setup.tmax, setup.f_end) + + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) + + error = run(timestepper, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' @@ -49,11 +51,14 @@ def test_dg_transport_vector(tmpdir, geometry, equation_form, tracer_setup): else: eqn = ContinuityEquation(domain, V, "f") - io = IO(domain, eqn, output=setup.output) - eqn.fields("f").interpolate(f_init) - eqn.fields("u").project(setup.uexpr) - transport_schemes = SSPRK3(domain) + transport_scheme = SSPRK3(domain) + + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + + # Initial conditions + timestepper.fields("f").interpolate(f_init) + timestepper.fields("u").project(setup.uexpr) f_end = as_vector([setup.f_end]*gdim) - error = run(eqn, transport_schemes, io, setup.tmax, f_end) + error = run(timestepper, setup.tmax, f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_embedded_dg_advection.py b/integration-tests/transport/test_embedded_dg_advection.py index c6238a811..d760a82f6 100644 --- a/integration-tests/transport/test_embedded_dg_advection.py +++ b/integration-tests/transport/test_embedded_dg_advection.py @@ -8,10 +8,9 @@ import pytest -def run(eqn, transport_schemes, io, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_schemes, io) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(eqn.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("ibp", [IntegrateByParts.ONCE, IntegrateByParts.TWICE]) @@ -33,12 +32,13 @@ def test_embedded_dg_advection_scalar(tmpdir, ibp, equation_form, space, else: eqn = ContinuityEquation(domain, V, "f", ibp=ibp) - io = IO(domain, eqn, output=setup.output) - eqn.fields("f").interpolate(setup.f_init) - eqn.fields("u").project(setup.uexpr) - transport_schemes = SSPRK3(domain, options=opts) + timestepper = PrescribedTransport(eqn, transport_schemes, setup.io) + + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) - error = run(eqn, transport_schemes, io, setup.tmax, setup.f_end) + error = run(timestepper, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_limiters.py b/integration-tests/transport/test_limiters.py index 2cc02d645..f5454624b 100644 --- a/integration-tests/transport/test_limiters.py +++ b/integration-tests/transport/test_limiters.py @@ -26,14 +26,12 @@ def setup_limiters(dirname, space): rotations = 0.25 # ------------------------------------------------------------------------ # - # Mesh and spaces + # Build model objects # ------------------------------------------------------------------------ # + # Domain m = PeriodicIntervalMesh(20, Ld) mesh = ExtrudedMesh(m, layers=20, layer_height=(Ld/20)) - output = OutputParameters(dirname=dirname+'/limiters', - dumpfreq=1, dumplist=['u', 'tracer', 'true_tracer']) - degree = 0 if space in ['DG0', 'Vtheta_degree_0'] else 1 domain = Domain(mesh, dt, family="CG", degree=degree) @@ -57,17 +55,47 @@ def setup_limiters(dirname, space): Vpsi = domain.spaces('CG', 'CG', degree+1) - # set up the equation + # Equation eqn = AdvectionEquation(domain, V, 'tracer') - io = IO(domain, eqn, output=output) + # I/O + output = OutputParameters(dirname=dirname+'/limiters', + dumpfreq=1, dumplist=['u', 'tracer', 'true_tracer']) + io = IO(domain, output) + + # ------------------------------------------------------------------------ # + # Set up transport scheme + # ------------------------------------------------------------------------ # + + if space in ['DG0', 'Vtheta_degree_0']: + opts = RecoveryOptions(embedding_space=VDG1, + recovered_space=VCG1, + project_low_method='recover', + boundary_method=BoundaryMethod.taylor) + transport_schemes = SSPRK3(domain, options=opts, + limiter=VertexBasedLimiter(VDG1)) + + elif space == 'DG1': + transport_schemes = SSPRK3(domain, limiter=DG1Limiter(V)) + + elif space == 'DG1_equispaced': + transport_schemes = SSPRK3(domain, limiter=VertexBasedLimiter(V)) + + elif space == 'Vtheta_degree_1': + opts = EmbeddedDGOptions() + transport_schemes = SSPRK3(domain, options=opts, limiter=ThetaLimiter(V)) + else: + raise NotImplementedError + + # Build time stepper + stepper = PrescribedTransport(eqn, transport_schemes, io) # ------------------------------------------------------------------------ # # Initial condition # ------------------------------------------------------------------------ # - tracer0 = eqn.fields('tracer', V) - true_field = eqn.fields('true_tracer', V) + tracer0 = stepper.fields('tracer', V) + true_field = stepper.fields('true_tracer', space=V) x, z = SpatialCoordinate(mesh) @@ -127,7 +155,7 @@ def setup_limiters(dirname, space): # ------------------------------------------------------------------------ # psi = Function(Vpsi) - u = eqn.fields('u') + u = stepper.fields('u') # set up solid body rotation for transport # we do this slightly complicated stream function to make the velocity 0 at edges @@ -151,34 +179,7 @@ def setup_limiters(dirname, space): gradperp = lambda v: as_vector([-v.dx(1), v.dx(0)]) u.project(gradperp(psi)) - # ------------------------------------------------------------------------ # - # Set up transport scheme - # ------------------------------------------------------------------------ # - - if space in ['DG0', 'Vtheta_degree_0']: - opts = RecoveryOptions(embedding_space=VDG1, - recovered_space=VCG1, - project_low_method='recover', - boundary_method=BoundaryMethod.taylor) - transport_schemes = SSPRK3(domain, options=opts, - limiter=VertexBasedLimiter(VDG1)) - - elif space == 'DG1': - transport_schemes = SSPRK3(domain, limiter=DG1Limiter(V)) - - elif space == 'DG1_equispaced': - transport_schemes = SSPRK3(domain, limiter=VertexBasedLimiter(V)) - - elif space == 'Vtheta_degree_1': - opts = EmbeddedDGOptions() - transport_schemes = SSPRK3(domain, options=opts, limiter=ThetaLimiter(V)) - else: - raise NotImplementedError - - # build time stepper - stepper = PrescribedTransport(eqn, transport_schemes, io) - - return stepper, tmax, eqn, true_field + return stepper, tmax, true_field @pytest.mark.parametrize('space', ['Vtheta_degree_0', 'Vtheta_degree_1', @@ -187,9 +188,9 @@ def test_limiters(tmpdir, space): # Setup and run dirname = str(tmpdir) - stepper, tmax, eqn, true_field = setup_limiters(dirname, space) + stepper, tmax, true_field = setup_limiters(dirname, space) stepper.run(t=0, tmax=tmax) - final_field = eqn.fields('tracer') + final_field = stepper.fields('tracer') # Check tracer is roughly in the correct place assert norm(true_field - final_field) / norm(true_field) < 0.05, \ diff --git a/integration-tests/transport/test_recovered_transport.py b/integration-tests/transport/test_recovered_transport.py index 6584e210c..87160db95 100644 --- a/integration-tests/transport/test_recovered_transport.py +++ b/integration-tests/transport/test_recovered_transport.py @@ -9,10 +9,9 @@ import pytest -def run(eqn, transport_scheme, io, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, io) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(eqn.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("geometry", ["slice", "sphere"]) @@ -31,12 +30,6 @@ def test_recovered_space_setup(tmpdir, geometry, tracer_setup): # Make equation eqn = ContinuityEquation(domain, VDG0, "f") - io = IO(domain, eqn, output=setup.output) - - # Initialise fields - eqn.fields("f").interpolate(setup.f_init) - eqn.fields("u").project(setup.uexpr) - # Declare transport scheme recovery_opts = RecoveryOptions(embedding_space=VDG1, recovered_space=VCG1, @@ -44,7 +37,13 @@ def test_recovered_space_setup(tmpdir, geometry, tracer_setup): transport_scheme = SSPRK3(domain, options=recovery_opts) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + + # Initialise fields + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) + # Run and check error - error = run(eqn, transport_scheme, io, setup.tmax, setup.f_end) + error = run(timestepper, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_subcycling.py b/integration-tests/transport/test_subcycling.py index 4b496314c..09915ba18 100644 --- a/integration-tests/transport/test_subcycling.py +++ b/integration-tests/transport/test_subcycling.py @@ -8,10 +8,9 @@ import pytest -def run(eqn, transport_scheme, io, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, io) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(eqn.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("equation_form", ["advective", "continuity"]) @@ -25,12 +24,14 @@ def test_subcyling(tmpdir, equation_form, tracer_setup): else: eqn = ContinuityEquation(domain, V, "f") - io = IO(domain, eqn, output=setup.output) + transport_scheme = SSPRK3(domain, subcycles=2) - eqn.fields("f").interpolate(setup.f_init) - eqn.fields("u").project(setup.uexpr) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) - transport_scheme = SSPRK3(domain, subcycles=2) - error = run(eqn, transport_scheme, io, setup.tmax, setup.f_end) + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) + + error = run(timestepper, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_supg_transport.py b/integration-tests/transport/test_supg_transport.py index f88190ec0..9475234f0 100644 --- a/integration-tests/transport/test_supg_transport.py +++ b/integration-tests/transport/test_supg_transport.py @@ -8,10 +8,9 @@ import pytest -def run(eqn, transport_scheme, io, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, io) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(eqn.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("equation_form", ["advective", "continuity"]) @@ -37,17 +36,18 @@ def test_supg_transport_scalar(tmpdir, equation_form, scheme, space, else: eqn = ContinuityEquation(domain, V, "f") - io = IO(domain, eqn, output=setup.output) - - eqn.fields("f").interpolate(setup.f_init) - eqn.fields("u").project(setup.uexpr) - if scheme == "ssprk": transport_scheme = SSPRK3(domain, options=opts) elif scheme == "implicit_midpoint": transport_scheme = ImplicitMidpoint(domain, options=opts) - error = run(eqn, transport_scheme, io, setup.tmax, setup.f_end) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) + + error = run(timestepper, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' @@ -77,20 +77,22 @@ def test_supg_transport_vector(tmpdir, equation_form, scheme, space, else: eqn = ContinuityEquation(domain, V, "f") - io = IO(domain, eqn, output=setup.output) + if scheme == "ssprk": + transport_scheme = SSPRK3(domain, options=opts) + elif scheme == "implicit_midpoint": + transport_scheme = ImplicitMidpoint(domain, options=opts) + + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) - f = eqn.fields("f") + # Initial conditions + f = timestepper.fields("f") if space == "CG": f.interpolate(f_init) else: f.project(f_init) - eqn.fields("u").project(setup.uexpr) - if scheme == "ssprk": - transport_scheme = SSPRK3(domain, options=opts) - elif scheme == "implicit_midpoint": - transport_scheme = ImplicitMidpoint(domain, options=opts) + timestepper.fields("u").project(setup.uexpr) f_end = as_vector([setup.f_end]*gdim) - error = run(eqn, transport_scheme, io, setup.tmax, f_end) + error = run(timestepper, setup.tmax, f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_vector_recovered_space.py b/integration-tests/transport/test_vector_recovered_space.py index 8c0f2b88b..4b4c15221 100644 --- a/integration-tests/transport/test_vector_recovered_space.py +++ b/integration-tests/transport/test_vector_recovered_space.py @@ -9,11 +9,10 @@ import pytest -def run(eqn, transport_scheme, io, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, io) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(eqn.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("geometry", ["slice"]) @@ -41,18 +40,17 @@ def test_vector_recovered_space_setup(tmpdir, geometry, tracer_setup): # Make equation eqn = AdvectionEquation(domain, Vu, "f") - io = IO(domain, eqn, output=setup.output) + transport_scheme = SSPRK3(domain, options=rec_opts) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) # Initialise fields f_init = as_vector([setup.f_init]*gdim) - eqn.fields("f").project(f_init) - eqn.fields("u").project(setup.uexpr) - - transport_scheme = SSPRK3(domain, options=rec_opts) + timestepper.fields("f").project(f_init) + timestepper.fields("u").project(setup.uexpr) f_end = as_vector([setup.f_end]*gdim) # Run and check error - error = run(eqn, transport_scheme, io, setup.tmax, f_end) + error = run(timestepper, setup.tmax, f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' From e4f59523c5c3e05c7e8e4c919ebb5a18bb36f1df Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Wed, 28 Dec 2022 13:52:01 +0000 Subject: [PATCH 09/12] PR #320: remove messed up merge and fix a couple of outstanding bugs --- gusto/__init__.py | 1 - gusto/diagnostics.py | 12 ++++++------ integration-tests/transport/test_limiters.py | 2 +- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/gusto/__init__.py b/gusto/__init__.py index cc74c9d6b..dabfcdf2a 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -14,7 +14,6 @@ from gusto.physics import * # noqa from gusto.preconditioners import * # noqa from gusto.recovery import * # noqa -from gusto.state import * # noqa from gusto.time_discretisation import * # noqa from gusto.timeloop import * # noqa from gusto.transport_forms import * # noqa diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index 7e12a8655..c8692461d 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -432,8 +432,8 @@ def _spherical_polar_unit_vectors(self, domain): R = sqrt(x**2 + y**2) # distance from z axis r = sqrt(x**2 + y**2 + z**2) # distance from origin - lambda_hat = (-x*z/R * x_hat - y*z/R * y_hat + R * z_hat) / r - phi_hat = (x * y_hat - y * x_hat) / R + lambda_hat = (x * y_hat - y * x_hat) / R + phi_hat = (-x*z/R * x_hat - y*z/R * y_hat + R * z_hat) / r r_hat = (x * x_hat + y * y_hat + z * z_hat) / r return lambda_hat, phi_hat, r_hat @@ -473,8 +473,8 @@ def setup(self, domain, state_fields): """ f = state_fields(self.fname) self._check_args(domain, f) - lambda_hat, _, _ = self._spherical_polar_unit_vectors(domain) - self.expr = dot(f, lambda_hat) + _, phi_hat, _ = self._spherical_polar_unit_vectors(domain) + self.expr = dot(f, phi_hat) super().setup(domain, state_fields) @@ -495,8 +495,8 @@ def setup(self, domain, state_fields): """ f = state_fields(self.fname) self._check_args(domain, f) - _, phi_hat, _ = self._spherical_polar_unit_vectors(domain) - self.expr = dot(f, phi_hat) + lambda_hat, _, _ = self._spherical_polar_unit_vectors(domain) + self.expr = dot(f, lambda_hat) super().setup(domain, state_fields) diff --git a/integration-tests/transport/test_limiters.py b/integration-tests/transport/test_limiters.py index f5454624b..e98a0842e 100644 --- a/integration-tests/transport/test_limiters.py +++ b/integration-tests/transport/test_limiters.py @@ -60,7 +60,7 @@ def setup_limiters(dirname, space): # I/O output = OutputParameters(dirname=dirname+'/limiters', - dumpfreq=1, dumplist=['u', 'tracer', 'true_tracer']) + dumpfreq=1, dumplist=['u', 'tracer', 'true_tracer']) io = IO(domain, output) # ------------------------------------------------------------------------ # From 34d74cae007e7876c89a3d2febffa72f560c18d8 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sat, 7 Jan 2023 13:34:02 +0000 Subject: [PATCH 10/12] various bug fixes --- gusto/diagnostics.py | 9 ++- gusto/equations.py | 33 ++++++----- gusto/function_spaces.py | 52 +++++++++++------- gusto/io.py | 7 +-- gusto/physics.py | 8 +-- gusto/timeloop.py | 2 - .../data/dry_compressible_chkpt.h5 | Bin 22648 -> 22648 bytes .../data/incompressible_chkpt.h5 | Bin 22648 -> 22648 bytes .../data/moist_compressible_chkpt.h5 | Bin 34104 -> 34128 bytes .../equations/test_advection_diffusion.py | 2 +- .../equations/test_dry_compressible.py | 6 +- .../equations/test_incompressible.py | 6 +- .../equations/test_moist_compressible.py | 6 +- 13 files changed, 75 insertions(+), 56 deletions(-) diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index c8692461d..34ac90fd4 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -190,6 +190,11 @@ def setup(self, domain, state_fields, space=None): 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' + domain.spaces(space.name, V=space) + self.field = state_fields(self.name, space=space, dump=True, pickup=False) if self.method != 'solve': @@ -339,7 +344,7 @@ def setup(self, domain, state_fields): except IndexError: field_dim = 1 shape = (mesh_dim, ) * field_dim - space = TensorFunctionSpace(domain.mesh, "CG", 1, shape=shape) + space = TensorFunctionSpace(domain.mesh, "CG", 1, shape=shape, name=f'Tensor{field_dim}_CG1') if self.method != 'solve': self.expr = grad(f) @@ -1355,7 +1360,7 @@ def setup(self, domain, state_fields, vorticity_type=None): dgspace = domain.spaces("DG") # TODO: should this be degree + 1? cg_degree = dgspace.ufl_element().degree() + 2 - space = FunctionSpace(domain.mesh, "CG", cg_degree) + space = FunctionSpace(domain.mesh, "CG", cg_degree, name=f"CG{cg_degree}") u = state_fields("u") if vorticity_type in ["absolute", "potential"]: diff --git a/gusto/equations.py b/gusto/equations.py index 9f10e20a7..1dc1281a0 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -73,7 +73,7 @@ def __init__(self, domain, function_space, field_name, Vu=None, **kwargs): super().__init__(domain, function_space, field_name) if Vu is not None: - V = domain.spaces("HDiv", V=Vu) + V = domain.spaces("HDiv", V=Vu, overwrite_space=True) else: V = domain.spaces("HDiv") self.prescribed_fields("u", V) @@ -106,7 +106,7 @@ def __init__(self, domain, function_space, field_name, Vu=None, **kwargs): super().__init__(domain, function_space, field_name) if Vu is not None: - V = domain.spaces("HDiv", V=Vu) + V = domain.spaces("HDiv", V=Vu, overwrite_space=True) else: V = domain.spaces("HDiv") self.prescribed_fields("u", V) @@ -170,7 +170,7 @@ def __init__(self, domain, function_space, field_name, Vu=None, super().__init__(domain, function_space, field_name) if Vu is not None: - V = domain.spaces("HDiv", V=Vu) + V = domain.spaces("HDiv", V=Vu, overwrite_space=True) else: V = domain.spaces("HDiv") self.prescribed_fields("u", V) @@ -500,7 +500,7 @@ def __init__(self, domain, function_space, field_name, Vu=None, PrognosticEquation.__init__(self, domain, W, full_field_name) if Vu is not None: - V = domain.spaces("HDiv", V=Vu) + V = domain.spaces("HDiv", V=Vu, overwrite_space=True) else: V = domain.spaces("HDiv") self.prescribed_fields("u", V) @@ -570,11 +570,11 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, if linearisation_map == 'default': # Default linearisation is time derivatives, pressure gradient and - # transport term from depth equation + # transport term from depth equation. Don't include active tracers linearisation_map = lambda t: \ - (any(t.has_label(time_derivative, pressure_gradient)) - or (t.get(prognostic) == "D" and t.has_label(transport))) - + t.get(prognostic) in ["u", "D"] \ + and (any(t.has_label(time_derivative, pressure_gradient)) + or (t.get(prognostic) == "D" and t.has_label(transport))) super().__init__(field_names, domain, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, @@ -649,7 +649,6 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, if fexpr is not None: V = FunctionSpace(domain.mesh, "CG", 1) - # TODO: link this to state fields f = self.prescribed_fields("coriolis", V).interpolate(fexpr) coriolis_form = coriolis( subject(prognostic(f*inner(domain.perp(u), w)*dx, "u"), self.X)) @@ -803,10 +802,11 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, if linearisation_map == 'default': # Default linearisation is time derivatives and scalar transport terms + # Don't include active tracers linearisation_map = lambda t: \ - (t.has_label(time_derivative) - or (t.get(prognostic) != "u" and t.has_label(transport))) - + 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, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, @@ -978,7 +978,6 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, # -------------------------------------------------------------------- # # Linearise equations # -------------------------------------------------------------------- # - # TODO: add linearisation states for variables # Add linearisations to equations self.residual = self.generate_linear_terms(residual, self.linearisation_map) @@ -1092,7 +1091,6 @@ def hydrostatic_projection(self, t): class IncompressibleBoussinesqEquations(PrognosticEquationSet): - # TODO: check that these are correct """ Class for the incompressible Boussinesq equations, which evolve the velocity 'u', the pressure 'p' and the buoyancy 'b'. @@ -1149,9 +1147,11 @@ def __init__(self, domain, parameters, Omega=None, if linearisation_map == 'default': # Default linearisation is time derivatives and scalar transport terms + # Don't include active tracers linearisation_map = lambda t: \ - (t.has_label(time_derivative) - or (t.get(prognostic) not in ["u", "p"] and t.has_label(transport))) + 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, linearisation_map=linearisation_map, @@ -1241,6 +1241,5 @@ def __init__(self, domain, parameters, Omega=None, # -------------------------------------------------------------------- # # Linearise equations # -------------------------------------------------------------------- # - # TODO: add linearisation states for variables # Add linearisations to equations self.residual = self.generate_linear_terms(residual, self.linearisation_map) diff --git a/gusto/function_spaces.py b/gusto/function_spaces.py index 7adf0cdd6..e7ee2687d 100644 --- a/gusto/function_spaces.py +++ b/gusto/function_spaces.py @@ -26,13 +26,14 @@ def __init__(self, mesh): self._initialised_base_spaces = False def __call__(self, name, family=None, degree=None, - horizontal_degree=None, vertical_degree=None, V=None): + horizontal_degree=None, vertical_degree=None, + V=None, overwrite_space=False): """ Returns a space, and also creates it if it is not created yet. If a space needs creating, it may be that more arguments (such as the family and degree) need to be provided. Alternatively a space can be - passed in to be stored in the creator. + passed in to be stored in the space container. For extruded meshes, it is possible to seperately specify the horizontal and vertical degrees of the elements. Alternatively, if these degrees @@ -54,24 +55,31 @@ def __call__(self, name, family=None, degree=None, stored in the creator object. If this is provided, it will be added to the creator and no other action will be taken. This space will be returned. Defaults to None. + overwrite_space (bool, optional): Logical to allow space existing in + container to be overwritten by an incoming space. Defaults to + False. Returns: :class:`FunctionSpace`: the desired function space. """ - if hasattr(self, name) and family is None and V is None: + if hasattr(self, name) and (V is None or not overwrite_space): # We have requested a space that should already have been created + if V is not None: + assert getattr(self, name) == V, \ + f'There is a conflict between the space {name} already ' + \ + 'existing in the space container, and the space being passed to it' return getattr(self, name) else: - # Space does not exist in creator or needs overwriting + # Space does not exist in creator if V is not None: # The space itself has been provided (to add it to the creator) value = V elif name == "DG1_equispaced": # Special case as no degree arguments need providing - value = self.build_dg_space(1, 1, variant='equispaced') + value = self.build_dg_space(1, 1, variant='equispaced', name='DG1_equispaced') else: check_degree_args('Spaces', self.mesh, degree, horizontal_degree, vertical_degree) @@ -86,11 +94,11 @@ def __call__(self, name, family=None, degree=None, elif name == "theta": value = self.build_theta_space(horizontal_degree, vertical_degree) elif family == "DG": - value = self.build_dg_space(horizontal_degree, vertical_degree) + value = self.build_dg_space(horizontal_degree, vertical_degree, name=name) elif family == "CG": - value = self.build_cg_space(horizontal_degree, vertical_degree) + value = self.build_cg_space(horizontal_degree, vertical_degree, name=name) else: - raise ValueError(f'State has no space corresponding to {name}') + raise ValueError(f'There is no space corresponding to {name}') setattr(self, name, value) return value @@ -120,7 +128,7 @@ def build_compatible_spaces(self, family, horizontal_degree, self.build_base_spaces(family, horizontal_degree, vertical_degree) Vu = self.build_hdiv_space(family, horizontal_degree, vertical_degree) setattr(self, "HDiv", Vu) - Vdg = self.build_dg_space(horizontal_degree, vertical_degree) + Vdg = self.build_dg_space(horizontal_degree, vertical_degree, name='DG') setattr(self, "DG", Vdg) Vth = self.build_theta_space(horizontal_degree, vertical_degree) setattr(self, "theta", Vth) @@ -128,7 +136,7 @@ def build_compatible_spaces(self, family, horizontal_degree, else: Vu = self.build_hdiv_space(family, horizontal_degree+1) setattr(self, "HDiv", Vu) - Vdg = self.build_dg_space(horizontal_degree, vertical_degree) + Vdg = self.build_dg_space(horizontal_degree, vertical_degree, name='DG') setattr(self, "DG", Vdg) return Vu, Vdg @@ -184,7 +192,7 @@ def build_hdiv_space(self, family, horizontal_degree, vertical_degree=None): V_elt = FiniteElement(family, cell, horizontal_degree) return FunctionSpace(self.mesh, V_elt, name='HDiv') - def build_dg_space(self, horizontal_degree, vertical_degree=None, variant=None): + def build_dg_space(self, horizontal_degree, vertical_degree=None, variant=None, name='DG'): """ Builds and returns the DG :class:`FunctionSpace`. @@ -194,12 +202,17 @@ def build_dg_space(self, horizontal_degree, vertical_degree=None, variant=None): vertical_degree (int, optional): the polynomial degree of the vertical part of the DG space. Defaults to None. Must be specified if the mesh is extruded. - variant (str): the variant of the underlying :class:`FiniteElement` - to use. Defaults to None, which will call the default variant. + variant (str, optional): the variant of the underlying + :class:`FiniteElement` to use. Defaults to None, which will call + the default variant. + name (str, optional): name to assign to the function space. Default + is "DG". Returns: :class:`FunctionSpace`: the DG space. """ + assert not hasattr(self, name), f'There already exists a function space with name {name}' + if self.extruded_mesh: if vertical_degree is None: raise ValueError('vertical_degree must be specified to create DG space on an extruded mesh') @@ -214,8 +227,7 @@ def build_dg_space(self, horizontal_degree, vertical_degree=None, variant=None): else: cell = self.mesh.ufl_cell().cellname() V_elt = FiniteElement("DG", cell, horizontal_degree, variant=variant) - # TODO: how should we name this if vertical degree is different? - name = f'DG{horizontal_degree}_equispaced' if variant == 'equispaced' else f'DG{horizontal_degree}' + return FunctionSpace(self.mesh, V_elt, name=name) def build_theta_space(self, horizontal_degree, vertical_degree): @@ -244,9 +256,9 @@ def build_theta_space(self, horizontal_degree, vertical_degree): self.S2 = FiniteElement("DG", cell, horizontal_degree) self.T0 = FiniteElement("CG", interval, vertical_degree+1) V_elt = TensorProductElement(self.S2, self.T0) - return FunctionSpace(self.mesh, V_elt, name='Vtheta') + return FunctionSpace(self.mesh, V_elt, name='theta') - def build_cg_space(self, horizontal_degree, vertical_degree): + def build_cg_space(self, horizontal_degree, vertical_degree=None, name='CG'): """ Builds the continuous scalar space at the top of the de Rham complex. @@ -256,10 +268,13 @@ def build_cg_space(self, horizontal_degree, vertical_degree): vertical_degree (int, optional): the polynomial degree of the vertical part of the the CG space. Defaults to None. Must be specified if the mesh is extruded. + name (str, optional): name to assign to the function space. Default + is "CG". Returns: :class:`FunctionSpace`: the continuous space. """ + assert not hasattr(self, name), f'There already exists a function space with name {name}' if self.extruded_mesh: if vertical_degree is None: @@ -272,9 +287,6 @@ def build_cg_space(self, horizontal_degree, vertical_degree): cell = self.mesh.ufl_cell().cellname() V_elt = FiniteElement("CG", cell, horizontal_degree) - # How should we name this if the horizontal and vertical degrees are different? - name = f'CG{horizontal_degree}' - return FunctionSpace(self.mesh, V_elt, name=name) diff --git a/gusto/io.py b/gusto/io.py index 9d649baa4..974c16db5 100644 --- a/gusto/io.py +++ b/gusto/io.py @@ -17,13 +17,11 @@ class PointDataOutput(object): """Object for outputting field point data.""" - def __init__(self, filename, ndt, field_points, description, + def __init__(self, filename, field_points, description, field_creator, comm, tolerance=None, create=True): """ Args: filename (str): name of file to output to. - ndt (int): number of time points to output at. TODO: remove as this - is unused. field_points (list): some iterable of pairs, matching fields with arrays of evaluation points: (field_name, evaluation_points). description (str): a description of the simulation to be included in @@ -318,8 +316,7 @@ def setup_dump(self, state_fields, t, tmax, pickup=False): if len(self.output.point_data) > 0: # set up point data output pointdata_filename = self.dumpdir+"/point_data.nc" - ndt = int(tmax/float(self.domain.dt)) - self.pointdata_output = PointDataOutput(pointdata_filename, ndt, + self.pointdata_output = PointDataOutput(pointdata_filename, self.output.point_data, self.output.dirname, state_fields, diff --git a/gusto/physics.py b/gusto/physics.py index 6eb863850..11b221a05 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -124,9 +124,7 @@ def __init__(self, equation, vapour_name='water_vapour', V_idxs.append(theta_idx) # need to evaluate rho at theta-points, and do this via recovery - # TODO: make this bit of code neater if possible using domain object - v_deg = V.ufl_element().degree()[1] - boundary_method = BoundaryMethod.extruded if v_deg == 1 else None + 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) @@ -523,9 +521,7 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', V_idxs.append(theta_idx) # need to evaluate rho at theta-points, and do this via recovery - # TODO: make this bit of code neater if possible using domain object - v_deg = V.ufl_element().degree()[1] - boundary_method = BoundaryMethod.extruded if v_deg == 1 else None + boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 1 else None rho_averaged = Function(V) self.rho_recoverer = Recoverer(rho, rho_averaged, boundary_method=boundary_method) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 1b4df1146..fab64ae67 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -45,7 +45,6 @@ def transporting_velocity(self): @abstractmethod def setup_fields(self): """Set up required fields. Must be implemented in child classes""" - # TODO: should we actually implement this? pass @abstractmethod @@ -459,7 +458,6 @@ def timestep(self): for name, scheme in self.auxiliary_schemes: # transports a field from xn and puts result in xnp1 scheme.apply(xnp1(name), xn(name)) - print('xnp1 tracer', xnp1('tracer').dat.data.min(), xnp1('tracer').dat.data.max()) with timed_stage("Diffusion"): for name, scheme in self.diffusion_schemes: diff --git a/integration-tests/data/dry_compressible_chkpt.h5 b/integration-tests/data/dry_compressible_chkpt.h5 index a3fbe43e1b3fc71f5a1d2e67d55e5f7dc2ecfdf8..d03b5ce5dfadc987816c3a0d88babec1c7f9df19 100644 GIT binary patch literal 22648 zcmeIac{o*X-}i5(%#|^j$A~gd*{6ALGZB@@kOs=IktT|$l%hmZsR+#o71ch;Oi|`J z^H|18GCZGsRoC%c-|zLi@B8>YzvK8F_x-g0@ZPVrmbInvi$22;`DRNZ~nc1*RLnEzw*ql%%uDIgqDVxhPsdc*LJR7=NF~_ zxxw$(c6&QxV;VT}_xXS4uiZ2@Oh2Eo{NC^<|F1GYG}gDFUhuCn{AZO`{Hn9xGQ9Tf z_VFZBE9vKh_7N()@4xNO);|XR82Dr0kAXi1{uuagY+&I%FCz^OKnqFYywt33 z@yw_%4?v?QH0!8ozBwoVer*71S30Hp0MaF^dP`(;09yQ@*%k$P%-!#aM|S`Ue(G+$ z2U42(_Mq`l0Ae(#C0Ro1zj(hWKM{cD?nH2FLHfT7e#Y}H0Oit|jL1QL99s0JTMj@T z+ZbopLEe~r~#40B^H*YlFmilEW@8AwxM z<=ir1GJ2=I$k79NV_tQxP@0SeQ?`G92f1N`lu5248EH#j)y0r8eS3AX)Pd=}>+&J5 z>77)_(jy};WjoWSkU^$n61gU16u@Y@Ef4aDJ==Vt6&Ym(?ceqivU~sM$nvdloa@JH z)GX=EN@#K+qtfK!3~JiYjw?|;CRX-epyqqwiaDxB`h?!G2FUaV)lRBMmiMk*k0DKYf@E9nRHrWno*xGJl2p$h_bjkPz=}_KepamX_MAWa$LAF+X zdoIOCMwctEUTuMVLho6!QiO~I^WL-Egq$}oT=WnGVw4b&RG7JZ=d*Z`TUB zPl32A2RyQyaQME2yk?wOcMUwo?JzINgWM{GDiXkB%$_WcM#!6zn#9N8QL#I*RDL)B zDLav;2f(9C)s~bulL3hN`F$tNxd3$hlPbW z$fz>vR4yAU8C^K^sB?gejP!!`W$UjdqwtLVub%Uh(QT<+-#!bH(M@H)cTS>Y#Ion) zX9aQaeCdo*4|vQ8W@T2FBcr5ZH&Z|G=qkvcXtt4z$jg280vcrWVS7w!Jmj+Txq?3M zNbEZ+CrThA=MRmgmEckHMSWr)q-#{(g*V{w^@BTHZy|%?1V2=Q$2nilt{TX^4;6`2 zk1e@@Wi5~^6Vx2(!K2Js2NOof>!Kv@4)FNp{EHtwkR35}eWl=$SJ3Hn6Xf33mkaN} z$pm zTc6m~JP+p`ZCzqm3_vOo#O(D9WVF6>2ahz=x<^*q{--{q9`eHJsIu$)Mxli0mijU zQ?mFbGJ4di`s1838J*%dJItX;M&+iCE88J$t$IH=gGZ?xAI4(HCvk!8i{SB^?6bRa zrZC^!zTe#h9%rwZY#V~Kc)o6J9e6BJI}y_edDz@)ycImQwI6D4hul}5eS--+4%}_d zV}=w@?;IZhj{?Pg6T^_LU&YmX!DHlySc@^p(+UH50xbck;D$-sCde1NTDl9sqYYOe z?KEUhpPS7*czk7Ea&q4&e1G_(*x8u?^Rl!Sr??z4~=1}N6!w%k5sP;>b|8T;4wjl_T(TWsc`P%C-7Ko z*0Q=E(&}7(xCVH<=n!~M7&2UEZQBxf+~u}CxeT&E{$YMNc)X#PKwJf>yxTm!A3Uz< z3b`x-*=sBHbOAi>oFsC+1+S{Cyzy!@*1;zl)2WjD(bz_!wjk9xpBFw(>$oZSv>b4IW?er$$;qcF@LE zErQ3jclC&0AjRlUexrI+{pO+426=w|{xPb@Czo{tsb23LP0yi@^8vbJuc_nS^oF-_ zA9y@|#wfWHa#OF2RU3GuOL*r~3E3&OzmEYtT59cdXo2xQW5ghF4m>&;dwrLLj85ET zcMa;3t@q(<4UGHfBhpIdP@koedJP|8+^_Imz3T*wb6208jFvwF(1C!u=y%L8u7ZXx zR>8QxI^Q00g&XF}CcnZ7824*?n4@%Ho?MZfq3y^^Mw2Q1lVM_H^vTX}UlWY`eo-6! zTQX$SFgD$_|zVgirWkJ{rj zVZ3iS-4)dg<9vrup@lzWh{u8V%HVNB{?qGBkoE3w-i(6BjR9=y+8}wW)}=RrM+G%0 znfH*6cj<^7;IaFA?3Ol2Z-M=XsO!+2?;wG?F1=2bpQf%uD>6q4sq50mg3^Pp!QStv8NTdRE(DMKZK*N0AY~?cCwgEVS}}a7F%^>TY@DzOcoaO`;aD*gfG&mFw@$-4 zH2i8hVl|}NmrnV1Sck3)=PHEG1t9O_UPdx8s)$`njGFWEmquXz zbZ|RMdIxE*|I)h_JW8onE0jV;3yj`<3m!Gx$re?R@}ohl)cKQ-%cY*WF1g*`oJ#GJ zLfsbhshJ~~xJ>Pn9I`W&skx<08a)Az{v~>+A3_R0T4{3$JT|9~UX6wf{CN9QIC#|A zmqGG_%v)a19R(hXUy&ZWLPiJsWZj&A`SM;>)i%h>^1Y6g-vUsjw`#pJH*2N{l$_HtnV5RJ+!@`ZpXjb=ZnF*R3_^FU?(RTWp;?|qk;AP-Bxx#efjQTi&^V7DUlpW7I3PN~DXO!oDu zfLynWcD5Egs^?#NRS4<6W4lukcuaY*Q2rb;OWfl)29K9_G#{n*O%6N{Tc~{!QxwB9>WPwktEK2GyfJ?>(PVxi{wm*jG4|9R+uQ9d=FO2w*D`%m?l{aw_) zey3etE47bjj;?+EkG_6+Z5XwWFKTJJ>Is>WE!9cwyQdX45zy92AB+&R=_>8j?-NRH$B>jA%F9b)-Pv6zJxC10HW) zuQtwy3=NPBd#(=sB7Lq=Hso!a8)CV7WMmv*kV5V2Mc>*LQ2Tfu+gGO4zFx}fNC>r$ zH~qj$N6p7^jd!UY`*(PpqGp}zrTbKm+e5{!QS;}0dg@){pAS)k|M1{-RN~50Y_uKK`=05!`fq(ywpOF36?Y}+! z^8){M_aFLW;E#bn2L2fMW8ja0|A7X~K6B_@=po@^51wO6vn2dXOGP~8S{sFxJCi{hF+oX?OV~%@FvFt9zi*h;H^a?U(uWK4hbe=8QNCa3W++)| z(MOfjyD2Ht54W#zn55*B(?g4V&G3kp)cf8nGrZxl(gTyM3PN}oM<%Ci8^tCyWqZlU zETPBw(%ld8g9IDhhPpi+W>|GJDan%29GA}CxR~$EggDk|+JB>=Ln)hBPu&+@h2$P; zOc`C^LI)Bo4y@ob$3Y&Ot-2S@Fg6Xz9c$KA!y+^)?`u^Yuv!`Yk40cX~s(*~@Jdy9pUK%@HVkYglVs^ok{ZYK%* zO18U>3i6?F6@gt9-}#W!RMdv>Y)*u}vd5)+aG{Z zNby1lUvnA_`fG`zyd$eLCW-hiA9{~NmJ7Wv zsE16D;c9JrV}`FTdgk07HNz_=l^eRn%&|`Gj|+}5=6ED}p8fDAbDUNmu=(C03w&|V z?vP`H1&(umxP~lmi4XHNkXs0rxH!Pvev;o3$4R-LRkwhXJxt#fnrO?LRoMo|RWe}^%<;(hu7PxDaxZ?B|b9}FUYpu>!Y2?jq zqtlcoi!O|@ZT;dag_<7;>giX&@s9+H==N;uWGQ%k{aqpzZNZ3MXyvE*106C>j?W}pe7TtDtX*M(A zM$H1fp`K-{&_QW8y@nDJKCB?J(c%FKSKT>#+a`h*X~ceB3Z0`ve|>)~UbTO_$7_KD zR5u>a@Upms(7+sbWv|EV{Pb{nrj9ZO&CAax{w`F}7GDMR#T(g&Cjj zP4rM|PU$!+XpUcb9Yc1^GDx+%{YoCM99qw9yEtYkfxf%1?m~i6 zXhQRY?EG~zye9t&-EmJ=NH7sa!;&fzvBGQ6Ef&MDyF)xE`|+ET zfuBj3ovFb-t%ii(NLIW!V99`-vLa+GRx$mp9={JRTeiDcV3u2ixv)wLTpYFfTY!rt zwr4OOJe+EUx7-h2s(ERR+gGhVlvHkmnPdqE&D?GA*jSgY%%Uw`)WxrTH`!sS3-O6P zW42h)wNIjXw=M2|(p!!L#!qD5KFF!@c<#HMZ%78Yk5`uytO)MJbY+QJ#U2Bg8Ci5?E(kNrW z?zB@*r9qsS$&5o=F_j<7v$GD(=7?Y!_ofcMZx(oBgLAv%W-IJ;Qk&mRPaIRGIua9> zWU%z(^3@rhGWcQ>!?v6cN?0eDlkmjS8b@pt-OO;+29KtM&35X@<2Q-)`wNLGcooCW zk6qphIN-flhP9I#HtiByE?aMl&2{ZgtgW`ia`kO4>-Me33~Sy-T5MIreX2gCE_>2g^-q7Yz|vkjuCe9Fp>JnhZpBNWaqm(NjIQO8*RR*#}JQ;B^>=2%c(U)&|h9CLJXwr*v$z#boEwmAq`Vic8>#WrP$ zD^q2z-HW!y0v0!R>1f;F6`QIE3MDq!PDo18>5~m+d8~dW3FgOp59%JhGqAx2`lRnC zc3R<*L*K)TPg>%x?3V()OXfJ)S4Cq`+#H|iF1a+0%U5{*rY4_$foXp|IZ7lmlEeWgfrRRFy7iPt< zIbU0?;B^W7bxbz1I@b)RMio{%Z867wTR$JC*(Z+~YCXS>i74Xirnl=k$;vo{a`TM9 zaaDZ&(N61=Z_TmUjk?dA#TGcSf7!hzOC4|K;`yBJr-jQ%TSt6yweXKS?)x}No3YK_ zzQS;QOI&v7REq47C63n|?TOi=jm_6|zG#os!&hdQH;|+?aXX!`9Zjk(?sNSZdz;$| zYo|$XqU2a&v(Nop>cz^~lSe50%?1su8K}3^BqNQZ8hPmzK5oD^$-O%kjo>)3PBj^V z1r}e5pp57-;wR@6(nWOnu!Hg98nX?nk<-k{-HPUn$V$xc>E6TU_}SJeCPI=K9z)ao zKX@gPPVmCm_VZ%Mls>#&(@_qw>`YhVl$S-SZ&LR-ikM+mu8VCKZ$lsF>$x|gS`Kxe zcz<2OTMixVW+m#(NTWle+d@?6Wl+@A=i$JkBz#7)^!9BM2@f58V#y&Oj;^j4N_cN5 zjvRuROCCPrM{{3%!dRFEkcESK#ts4rk6jAcvO<-FbyG;yYo=MykItF_4#yRLtH)XU zSEg5ZNZ91t_buiEB%D`#|5nX965jW=L?T9xgx^aKPPlJ`dWpFCG)OA53xBNB0f-=*ZYxyCa(^XaFXpQ5E=ZccJd9!*Wp@Q@UjLol5vHL0F4rQU-UJ4G-H)8ewj;M|N%Rjq#0K=f;r}0{D&6hmKZVHaw%@ z#ZX00hekxTF3yN^Asnd}@+s5=du)4rE6dUp+cvtX6Bk4giPx$@rAZuF$k?;psFFe2 zH`kWcc*&xfH^E*##in>H*Y9uufry_>ng)MRkV6lSIa-n8F5ua@8AGmvqi0?2&B*e>#qnXfvD7|(u#9Xd*>CON@+QC+wqI{4arTdak zQ63Yq^8O?#mY49pb59Yblvf~n&%Examw%rhT_~cp4s9m*QzCOl?^zSfVk`W$U&sW% zUhgNbBWR4H#ZLD7nHk~k(CP8TLIYf6UBy@0u8-r`N7!Ub5botJZi)0qcm%5$ydOh& zt;wGL^sD-K|0Ct9wHpj^;Y#nMq9E`d#m??HX^gEO&vkWGn&53nAVKDVDL(Sa>(Kmn zQ>+_&;AO5O5i|K^^sx95@d9^n{dm;h>+xz}wTIXm6a20+<$>oTW1LF7{GBMkg`*4o zZOg8$zzv_L)l-lt-tn~Qo?MgwKJ0Q#FI&VIJDM$}D0LX(1d5At9Z00c1MmpJuZv)c;+gw)*0ZG_aP!m69^k^%rBH)Qo*#}torYiD&Vye{)UF8 z>NtI2>Da9hMI7+?wCYqM0ekyB=<;5*88;0~QtpJS<4hBV(j9#Yc*Xg(mdos_xbgJI zxO=)X__$R&?Qt_*?7LvM%OPwt&SDaFy;8RU&vCPIHa}a7Z*98k+Hzk6>z}wo*7l#J z*m&HSVDKj3#HVTL>DTo!V_Q%IW6TH2(a3$nl}xM9!bz9+8N32${A8e>kGD8FEKqvx zimoAkm+sQv%3+KPEOM{UoRdOgw;KxvL}k&6kfMmBLOHaZCxnGdb3IzU;keP|wI*1D zn9(|`VT$z~t3N+HCx^-x0?XHalSPiw=|*h(rP0Akg`JwgQs{DF(yPwfrug~68?T)N ziP&D;Kj?hC7|I*YEtxJ5MIDFkPxN*2p@Z0};_Y2tA_I&> z6=V19MtIzErA)ks9=5sc)o^A?2TvUEu!&I6!IirtCfaRv@gC(vM=}oT;d3|IGO8u@ zvFmxuh-)z zj*Ad+T-+NKZ}j(i^ynbw?oBYp4=Qz(XjF{xoSVBx|78ZusofvkczuDQFmg_hmbeC2 zTZOwgWwGMI70G&<=|)&tq_iXPxFHU@clFVljCEMNe!IoybE|Q3ym)ecnmAr@FC$QO z3m;ZKWq+0_!~iQAbBDcXN4Wd~!@E=kDf}j`Vw^FMA72+0@v3r=!kKn#kHw$z;AgIx zP1=_U_=)XvS&>tl@ot@9rq!9^cy^$#^PL|%W;JLipiLIQgSM@QB6vPibi4GMr)G5V za1uic?aXFuA7ioSyB`gj5{X!k3>414>tcrO2fYskHydLquE_m|rcAK5pM*%Z5hu@chkr1k_0C6Vgvll!Ze9-*bw&*EXbdC zGREBNKi&*hG{M6y5%IZ#rnn;b^fQgKrr1MZrC{(^Q+)d4mjXL+BIfJ~Sh@e-)FXQc z`Qbhf6U?cXx>a7v1YfUFhlz7w}u(vn8j^% zyORvD*~_C$PvqJ0nblr4ry1rb_eqDnyY@5UL&*|NOWtji7Z&2LZDniR08WyZ5Nz`>%XR86$?Rrsef6D+HbX+AS zzgUHy_|J+ct>HubC*t<(x-E?E6)jZJv5TVfA`tAs41SG>JAC^N^AclRc>T=4D{~VZ*Gc~FyI_JdFEei2v(FT7 z7Hjt}>NUl3Bhs=8P=DI_pPur%5;5)4*1a=_iMX>_qd)rG-{(i|9rp*SmQ8V@Z3Z*# zm??hiv4f(1qnD6%QzuhkzKxKiwc%dnwMoMDFpD3HN<)OS6H1kdZ%y$GyG`3%yeXdB z@zTDxX@;OVz=uIwC5#=qHl$r5?K}x~41e<`Y8A^5o&*J&dCV2M8<36J@6I|ibpd9ByhtCyS zeR!9}h+{jX{CPUKu)2&Wd&4_!O#dO~C>J)tr)yS>v9_Dwf{m}rgogOB#GR-V5q1H5 zKw{v#lawgFP_f`f|4;-^-Da*3Ujd$;dW+2LFvYty`QAoApKzC^soTlM2-NpYN1dCXAxpcr<1b9KX$d|_N1t<$e`b{)PU z>hbjJMItWbaXt7R`ocUunOCi|e0Xw)T2~u4ACB)4drb1+!Vc9@-J~Q=?8|&%+7bFj z!ME-qTT6);i@cvtN~6J(*DDYAbJ5_x&X51u@3+MM^D3SChY(KSdZ~2^y@uIz`x?$Eg_xT=?vy4Fxz--Ai{v#1!H7-FpgkyQc`7R?#gA z0e2OhaV`QrXG1s|Xfs8~{B9B&q%uW#V=YC z5iLuZBvg3V*1rdK?fyJl1iZ&qEYc2qNtc#}D}XQYL6;}Sz;~58o(TiDa&<)=179%) zq~r|nIa}BVTzo`0}BBp=Atgw4U-l9QaDbCt*(TCF*d+ z`YHIzzv98!6Fx=oY?jnH3oKHnziS)#nj5dp*$KWr91gj!0vvl*D^m}+Nw=2C6nx2_ zz9Do6d@&t7U#|cxNAuxU3~=D}YU^_F#kg$r*%y4>EMKXX56ls@O2Qqu#U5P4p9Ge&B2K_&EP`9@Ix~Y??eUYY+$1 zJMcB}G5cj8_V8ks_ZDSx*s$PcHT?Do+t| zcc`8d0AG{6s~3-fuk@EggR|?v_ppkf8?f7joli2rSL_i6`mNyW#B_XJEb#2M@K1)o zWyGT!bHUdsE-s^V@KsPzt&s*?xa_q(5%{A^dCofUb>Cd%_U!&?LSbjAFDLMW^Wpmk zfZe}r@zDcc*$&%WoxvCVmSLU^z#jQe*P8&VC)J=1@KycD{t^j%iLNykD)T&vS}k{&L8w7JQXl>3ui|zOs%k60ZT@9}w1l z2`rV)bG93Ny-@p=B>38M#70U3e2EHl4EX?mzq2u1 zhU)KbcggW{0=^1eb}F)iuT#1uH5tGHwr_nI zfESw&N>jiWX5Bs?0=|v~YiZ;GUvbM$^Z^dL=cFD2zDBsqk7s}{YyQmsy};hiJPDV9 zx4hV6$^gFjF67ug2VcpYYNek}P7w}CY$K;bZX&Rl->Q#-z)>P2jBem7Oz8_n7<_4MkQHzO?znMc z13$2B@Yv-D@O8)aKwAj-lGK}=ISwrMsrDHeSkK7NmFi1Se=@HTd@<-89$^A5dEP@u z^%+=jltU1Fk<{#r*Ml#P%*X~Y;PZj0*VKT?HVpT!fv<0;_6d(IPZ6RCk_`8$?Dh5v zE%+;Df4V&ue64j#6VU`;mo0db9|DVXep22694EhBO&)v=k!J#I!Plke5L;W|6ZvoK zLx9W2A02%RzTCw(+KU+q3m_4##Ir5gC!RfP)3;H#tX?lT+U6uymG zr-AFU4m~{$zDmk6=U0QT>`nc<6M#>8UhP>2e0^~7wmtZgT=wAS17Bm?dCi@Hw|&d$ z5CS&nRv=S-?KmbCM)kGlGIvBcFsY`>lIn9Z_JlnL_|ln^G3owFTJy?qoi-{#-=b!sx-5ZKO+2HSmL*MxJuO zb+Vhn?t?F*W6kHbfUms-{fjNY9*?EV&jWjrU7wqPukF6XR|mkCq5qHZDBu#BoKsJM zQ(td8Q4PLmLKOPWf-l;7&U4&Qk6x!z{HuU3&7Uf_17B|IHNM7xulEvnD$W5nFot|6 z0B$LX*_Z;pJP-e{kO5z2!Vje102f*|9 z+Ckt=3+Y?vVE&W5COe41SKArS2zlU(qG$A`fxj`TyA*-1JL!3sroq?6>NM>+;PbJ} ziHa~En(#=Q*nltph1WUsF#i>8&eWWF#BF={ju@nIfAdAXQ-bU){~nB%)WuZnR3;K&VjGWjQfxE!53Y-;LIK1 zy9Zuy*#pZmA85P-zB0~k3QhoDSFD-mbAWYCb$spuCu>}2GXh__4vwSiz}M~2kt$nY z_eU2aMwQGu_5djy@DKn=4$39OG$^~ERx2&ai17F5NX{Sqp z1(FtdsQpawIs@I;;48S*x$+74aw;GcRRNbBm-eCdGi|~8l}_NRwuf7m+Aj<5%gm+r z(^F?BEUEqUVTP*42=FC;?ur7nUoO~};!W+R-(7fgDEA-!TzVV;H$Oz9Pt783SyNh*bcnt=tRi@j&eS~joL52OOVM424A0rW{VoB z{swN(9|eB4n=gpkFV_dJNTK%2gqu9d)P7of*9;4_pJt=c{PllI8EMwQZPmbi2X<3G bR{cNypTJ+YaesaCl0`H$Hou<#e;fY`j9Uc8 literal 22648 zcmeI4c{G;oyZ3MNkj!J5hs+^EBCfNUib^DDpfWU2xl5u%X-+gzQc+Se7Y)LFs+341 zR0t6g$yCO&&+~HL&w8Jpz1QCB-S1lO+P__YJcr{Z<>NT6@AvpFPseQL&YvhaNsxhm z@$)md8KU%??Ej3BUjZw}Ut#y`_q&ldI7mH8>MhRkH#iu44E#P>^13Mb{K{P8FKj2T zJ37pp$3P#+{`~Ap-pyFbJN_0w`-MOCe-#0nd9#+{FPMzsUsGB{=8O%)bCbu`jZ8dA zG3xBG3|fc{)jG@#=%3g|zP{r6y3_WUoi!2e=hZogzMo?!OeMUuH?f4<-6_`mpR|JQvO z2ZML~dqBW(5C!Gn7=OWpaAj96^nUyvF=gtn=a0z*hRGQDI3DyWem^G%C;bWd$;WeX za^51}EBvY0-!_Igc@uwd?-#BO3`W5J*Y5#--r>1o`vP{2ll%AoTJQdDV;F2se|>@f z{%`y9>K_7s2>c=Nhrk~Ke+c{w5tuJ*xZGFi7fayS8A-E63>a4-ujuuthQ;gc*=1nc z#CqL7GAXaCl@(%kE+yB#o3-umVP#<!4Zq^TE&T9VlzHK@n?jIxd8Eal$}?&`ca(PB+B%l;b~W4n zef`uO>5;Qs_AOM4V|UA!?IQ5@^-27vM?_)#8xU@-?)k?jzJH+z1Z#qtCSCc;`q-Iq zuKi65YoV3@+&&2za9h~Fl=VvnWFK}LICWhF)C;@j-SiX%a8RC^=G4T}+tMaf!q>=Z zbI%kTlxBeTIWyX`ZZg0>Msw-L`kMFO^BC_PzkmI~X$>s354?FKmCqO}^*LB1Y#ky} zI^#(A^{1>t&g(yKP;@(z`Nh^DWEDGMyEBT=;OXNrL9Smceve>T z&4YrlyKvRU8q{}OvN&zrx5`xQq zcmVV`awcFFZDPAI4+R(VOIHfE~27Kc`{W3L69JF@!YBlI8gGIuL zU{9c|6nro#OguGmdBi z(~|Y^unm2G_%)pM!(^bv^>1ra?5BV+1s#i%CEJH)qvYfs}P=z8bBgL#JqTpGPpGN%E`5#Rlu^uGSP-wE?AnD zsP2(9%1V+obzDDJ4%BBx6(lE329iqqlzYv@K+s9E-S&^ff!o`Gc*S1>tgTP3y1uIE zXSv6!Rvh)?0Q-D%sl2b8|Jav%S3d5w@0W!J3F~$b@1G1?Cr+K7E3X8T3N{Kctkj@6 zI9b&aHU-|x(VDd&NEbf*FqG6%Jr(u_=e>8~HG=1!O)S**HG<;T?p1|k8$#U^M~$|w zn+iSVO@HPvq6KYci;Z^stHa71t?fH@D8M^n4-QurNW*oW*%nRN!tkL`?02U}d~j-G zPR!3*2CU5%a$Tq~%Ia@j!YRCR@E`Nomtu2drp08aW?*T2#9k3LH*V@I?-^m`3@&I{ zZpIB`EzB&=jtat-QHxsEX9?&!!}586lL|~Zev$F0Lj&fmm6&QYDh@-X-K<^+E5PEh z(U_};q~V^Mp<&9Fs!*%mJwACz8$N&dF3U>K5biss8@|b121=*L$348J4&}UJ#nX}wyCL+<&-$h zdXP&N7SE~oxzug~-?u!b&i5Nb<>NKC3~uuQjR=7`8uk7#0q9en+Q!Oo6r&qi{kK&=F`wJ+=^E8bWkk<)dE<=xr! zLfn8GjP++uF8adrk9m~7-|~IO$4Ss$0NO-bDnj1OJ;#G~slc&wJT=eGO@USwQC8c> zbfCk7wi`kEM$miZ`&ApFjNxR@icd4R%-|>8rC!Y?X7IV{;LP_)rVyxHO`aZZ40lyU zf2`PS0C@wmJJcJrpr#3{AwgROzO8h#iB6D%>a&+r-_jF>Sq4uVSx@=lhS8frl6M&} z(=zVby}l7vla`zFiq`}Gn8()e-DSU$Wnk-mZr}XYNw9E0FN!ro0B)IC68G@A2-IDn zGW(XI92~7!@LYjc8K#}yw%lcv5-gS63&fy0T=2l9CihxP znLDk|Lr9pDFGy2##uI6$KZ_zB?`tWshXtlt`Th z`!{huNyyNLJHB|&iAj_MaEnE*>7zH4L!N+yQkoiw(zOV_aAgvB=ESe3ET;+2&D?SL z`)2f9kxkaAHq-|1^cvGIjA(!wn|hrCyt+W^`^&ur?6<&~aqZeQm(wW4-X~D-DoV+88$b#1ybip(5@_C=ZnBOYb7 zO&OFF94}5tQvoT$v5$vprGd5Qi~M6;vcPWj;+PAm1FSZIxBOgz zYl=1$QCX?fY^4uNZ+0yTE;592x*_jfJ&j;h;81ITgCW#vRpEW_iI#5ot zhn1$;Y-pX@`Hy-066Lbs@>fnc%`LcWpqvK^*GGR$eJ2L{hf_4x+DkyG{VCGj@~Tj$ z=xj=3ni{+(>!B%}CIHq51VA2PACroy5mztr@kX^_8p+6D7V zCNMOqux*pB6ttCFT6R5H9@@BDSl=HrgA4XomcF_?15#aq9QrC2u(tcE_T7_K@HKbH zaZ5vGSatrhaAbihR9@IR^L)Dn^srdn>;J(T=5~zns^^-+l2NdzSI`PpEvb8KyXQ2x3EiD+#b~J{}2&-nLOVi=<^de7tT}^mgm)A5baVoTO$arz+o-%yCD`Hn` zksRb>Mhza>%@50XSd|hHvarPMd|#}VIAFS6UlSn81@s)b`-8toK#Rd|&u+~Zgia6b ztM9YaK>75>9l?q!ATj?<`a%V5&^=hNIH*<=Y$$#GZL>2 zE)?|c)-lrtdG+UePE}0-ANk!DFKC_u(o|L(OmFRGWzKsWli2Zz<*eB??FNr3SRG=Y zlsQ!uEPFHtJ?M|H-YC$1HW~ML&tqH?bbZT+!)Y1O7V3Et0v@Y zZ%c>u3L~Got^G(?p%p_#rMuUZJ3nN>gU6hQhO9!o%3n6=f3$)ZKR=K3=F1Q9JHYj0 z@<3jOX`=bO`;)C18~1YUQkFq}ID2r#iEmbnknZ$S{-0KiA6;$(O2YY#W^xV5Q^fK^ z-uOSP9O$-U{Pt_bSvx$Bpk%)to=0%XC61oMeG5<0bE%aPMbDvAL^3^>=hP44c?5DB zR^WLAwOrrnIlRpCm7YseRa1Hnhd1w}=koK@2Y4Pq*h_jI0iWq>dJZKs&FQ&hdc@Fk zSaUg;p3B0b{dgY1o_Do)9)Wo28hQ@B3|i>9^m$`R&!NO94;dunwE#-I~!SJUVdJe-b+tG7rDe{q?!`mC4&~sVF(TwL2tP($j=aE{RF+PX! zj@Rh9{9dx2oKW`6uK2kDsM2*n%k^24byK}t{ z)@$I|{L2UFHCWX4f!1lA>-w~AFX&oK>vW!x4z1hjeL+~Sfwh(~v7Od!=bLx2UIU@0-B_=|KE^LvrzIJiY2D_COQChz%$$eT?Y>YKtk+<1 zY%|tt@S#SZ*6Fpb&uQHjYf7he+EhD%)@{AdQ?OoxIL-vD*Wh68GFqp-pXJcHJ*M!A z*6DAvj2>g%j^?8E8k`g7!g>w1zCS_hG_UQrZZ~~0qIJ4>mLRR$stdTWUIUYm^;oX~ zpLYYT)84Oc(YoEi_(AJ*7H1``+XG8zy#_A=tFc~#;dSFW9sI$C-T4z1g#$|qvIrW#z=V!fu`RxYA- zTE9S@*6rWFzt*>Ev3~>|iBGV91X+B$v3~?-%B8S>1oJ(^v3~^5^*&+$2(-V~VE+gl zyI*1d2x>1b#{LngMgGM85%|YNWB&-`xuURt1OlO%*gt}mSMsob1T$>q(x*lg?{ z!K>;L>>t6Qz)jdcf_YoMV*g0Zc;SNm<3Hw6Wa(Ud&w$RwCitEKK2^);J0>)H9evll z7@J1lF)bzj^j%|_4e&h!_EAyzo&kK_-t--_^o=xq*W_6<={v^%h3F-G*X-QA0pBw~ z^7&hQ&j7hCSLr)OQeKn3YnBbIrSBNgpalA^(U0WC_YBzP6^ZW|5ZSYZzGHe8`_Xre zbGkKs#|VF5(svE_)>?eefT;R>e9wT9`q%UwW4ot9U6VMEzGpyA zKYh=D)ctAn9h39gj=pQI3Wd{m%x>Qw^j&j6=o7wY!1)tm_?`i;qPXZgW>lq$zH5BT zhv_>;Jh_{`Ych*o;d=)3YM0}C26${bOW!eNnS=CQ!+ATLzGLRO9ii_UlOMD2Jp+t% zci?*lgq4opG5itV=)2~0;c5DgX`DO!0pB&D$vf~pLp|89h3^?kb5j<5$FMX7i}79a z`#fq*eTVOBP~Wy2-`C*0r98f`K@P(k-`C)})&YE9gOC>n_`U|K>-zA04U~9T_`U`^ zCr`llHP|^w-`C(;cM!g>L1Tn5zOR8nu`s@`L6)%qzOO;P%^iGSgNTZ!_`U|;YeVpT z4U9gO;QJc9T(xHvy00n0kC*X%{f~K+t5U{!4~Us3it`@u!Fd6l^JLGCp>v;Bsr_`$ zv!Y$95a&JzCenEih^$?V^ByomHi*u7O0Fl5=RSLl>6}NGON`Ea?)jzRya%N2r1Kuo zYI~Q?dDMaz(z(yWz9>59@fOhV$GMMM8x!X};6pT>_kgZj<#f(7y`X{4eHeKqbk1{o zS2dmc@DFC=ya(LQx{dQ55WIan=h?B~C7t_RlTM;@p1kNGI`;`Z^%UnlU~JPeoc92; zVGTOxIXb-0_Q#8ZDJwLd#Enm zNIK_X=~&XakKK!Dbj}k~7)s|p-HCBH@1aIEXW_hu%9s#G=RA*Z$k4gZ@AIgr6p8ak zkSYHK=Z|2WS2E5Y!GdNA=Z`=>=?=~xL3H+z56T}wY*`%6AAy-w2F@Qr)$szHKY}%V zV{#~e1d*SF5>WmKmPSb8{1H&bcX9p*&YorB{1I3f&cgX4x&uRT{s;!)#~_qHQiVT! zaQ;YDEfm7}BlZ0>h4aUM%;RTCIh@y1XME4%ydId8@54DgmG>}R0p<3T{-j`>(^I@| z&XZAYPchDW?z9WDYc%NIJc*41IBPpPZ``^=!tTBs$E0`=k?%B zwH(gt!674goYPaPXPvB3ZqGVoA#o1n^eoHKVjGm(vnnnR;k+LBH45Xr9!#t7#W_8U z+&B<|a(nn;9p_Dy)5AosV>q{m;+N>W9>kcb;k+IMM0~?JJ?zQl!MQzLFMAZ{^zhUX zNu1lm6Ztc6UJuqs3|>NcJrLNt0q6Aaq@^a#?ctf6pHomy4?UbqaBdGr3O3`s9_XmL z;JhA0$U5Vk9*HeZ|&*HqEy1kYM=kzeOybkB~ zP=*f%=0uw}%1z z-~CWtPc7Fh!FfGZ&S!yhdU(M963*@6@AEkRb3$By{Tv_E{rw~U-pi6UWY*tb|Gi`X ze8&a8-{14E+5cxhClpElv=4*9uK$gn6DrpmM~FWa8}t8|pA!lZ`VYSUrhk8(&FQZ% z@bB;b)BPdvhrk~Ke+c{`@Q1*^0D%cM&Ck{x@1?f5)owQ5(o3~8ogP_<`mGzzE-^;E zbnE(m)L$@qDfP(PF;=Lz z+I~=vh5Dp;lY|RjdZ^1Uoa%SK?4i8%8#TjGZ_qU>cr)tX1eaBuyw^hwp6K%`z0gBN zEwXkEM!kI7x79uEn2Yx@fzeFd+dsO z{zpO{hNyR3e`?Gd@wz>5An7LJ_2r_j4u$%}siNLNsPEnX%$W!AD&3=}lZ<$UM;6+= zKhQ&sbgb_WMSbJTvyA14SN(*@t1gJw=}FG}^ie+?wSJZb>Wj2^*Uv?~PUU1gi$lCl z-qy?2M*W$eajhp%&pEK*@H50qczDL~ZHQM|fvI9P>SJ&6M7X2=#8mJ4smY`=kcrFf}JIU!zcnI+3kq23|8 ztZp{yy)Nd)+(o=}toGL2K)hOJ&JVwidSe0ojwsaM^;Z&BK)h5qF7NK!+)F9jCU*#< zp5Iez&nMKkmnD17M7(N;+O{|&UViE}+jLQXan1(b*{GM!nSA&S;x*stm7h7{HEa|e z{t)$jb{DIqQ4jU!SbS^mp?-P5t7(YW_f&bu=cr#g@v(X^>W{8}a4ZAyGA-%5*%aMF zReWuZ3q<|8+H2gmP%o%^+5#Y6tpb~`_ipT=Hq9yG zUWZPl%&kMb?pN=M;Xu8}g-xE-sAq;ZsHh@dYSI&ynju~xlSlj~qrPj+(RO{*-zYS6 zNk+W%y9-je5wC%=-SY9MUsAUH@W5CPHM#-(%tXA7X~rgOM7-p0B(J%K`uGs5gk`7~ zE_~s!7V&!6c6Y)t#4DDsN6HrUS@#=VcA&mLX|G*NXb)A^*Pj!McyZP`^u0v=oy!|S z4xwIb<2-6D;`Qc!{gp_>D=1;Dy)Ej8o1CBSN4;;Z>YG@^i?clrs32ZvVoVhxP``f5 zCrSqOdlDnTTEr{6|9q+f;w8GzctB>&VCp@PTo@`#GgcqA9VDn5E=e3RSlH54X^D^P- zO?agfUJ`_77U4NuGtTQO;nhKSz9u}k5MFhJmki-qM|kccyw(z4`v^}#!c&m&@*=!K z2+t>krx@XNhVYUgJXwS%dtU{}zGjns)*}0SknHPKvajrYZYBHti|i|p{rJABlYL%H z_L(DXd|#c&zLt`G7V;e5=RC5nE6BbcA^U7j_IVxIS7oxV?0s$``@Dqgt2f!#I<>wPNEk}iJmZto*X86af0YY5783=q9+|hFBTEK&?S13PW0p> z(F=W|7v4lq)QFzM5x){m{E7wfGcm-^?D@^FG!j1(O8m?%;#am2zYufKT|~f%o*ZW28ds&Bz~ri_?a`r zuN)Y}1;j5eBYyhl+;Kn6N&Ipa@yol3pVlFM zdOh*WI>awOCVrZi`04*FzfBr|5N@yPvwdK{;>%c{{3|} zr@y|yzrXuW_lLkA0)GhnA@GO59|HdZ1VmaQcp1I3Do2M@jV1A8{#?l_FJO#)tJE#JUnB5AU}WfyfUh%T5fOb z=VgN4f_XVod`!?&9M%3wm0)`{0~B z#IgK>uM$MH@2y)M=IB|e^soq3iofs$=D7EzzwTR9@5uHjb0Ch{hn6kpL^btJ>B?%v zvG7%N%zIQ<@Xp?F7jfk1&{>s-s;gt%(aVTqlfT=72dD~;xQ@gkj+tY2qP3`AI0w^m z5Xa>k&MTOI_5*d&YDpuAW5>Cq>i6A#z`n~h#bwYBoXgECygP)xZ=I5@`3xqw+-rT~ zA}13#2ui&d6kq~An;+t~6PTbh^<$Kx7!!PPuKavJf(cfZ@NfPq%LHm$eOmM;GeL1( z+lPb7OrSBvsnAW831-*tF=uF@dET;AceOSXJW;P}IEXkZzA*IIV2I{YWJbFBbSBs@ zaH#kis#R5|hCU;X*6x-TniS%!^QpT8alDywETt9I6+8bjOB>vt?c zZjT?(y?)4Dxz7(caIMyr9PtB^yN1tM|MUZ~o7cSB%gqEC{Tr2cp}7pY`=H_UL?*D4 zU%pibahzwa;M6CDUQd`TQY^;=UJg1b0SahN2Ru$osW3r+k94JtI{N*0i+~N9$L~G| zqoEF(M;@z*J$g*g9F`eoW5fglr=N;tPDArN>z=N6+{K z64f=@+4h_Zm_SfymDp5N4{RO|dxtn$HqB~A* z1M7wMB$}YQ*XwI@0pi$R^7Y~%s)?~*4)A~R18v^B{ryp0bo%PciT!>cO!aV78LCID zj}4^%@&k_tW@ddr`#pQ1?XqoX9*1t|b_t<*JkoV9Hv-M!Mf)0&Lx|(8oncZg@=TyN z^L@^CMJCubU$QA!4ejBih|=8}Ot5U$qe-@C4%ax2W_2QtHVu<1yrwe2^sip%{|i8$ zujh>GHDe~IKbid8%Y+FgYr?ewGnn9nWl~!kszCLlk|;pW+o!KiE~l2Ex&paP1%*nS`srrg4s`gzHnnQI2q(FgVUJkZ@g>IL=X&a7`l|*<95L z$H#gFv_u6F7Ey>=il0Ai z_@;E?lLCluvL`-Cf%v9H#3uz1-}I39Brbz--=si%(k$Ye*gok8@l655C*={}S_@o}G`lR6~5zYvPmIiEmm%d{PhbO?!w>0>n45eY_~~ z^=u!XM0~wA@$oN+ulFZD{u%N0Y#*;geEn?V6;^U_gU;mi+_#?#E zcM~5UL418Y@$p^6*V_^w&-V3v#K+eWUvEu(ygTvram2?55MMt`e0&}8^+LqQn-O2H zMtnSr`1%gwN;<{HovU`F+3Vyv}(&=Xt%J*Ki~YNHQ9;Y!{r|SVPE+k;`1y|M_kln5e-r`2E}ea=1OAnU|IAX#uRQw= z!})}ho8zg~O#1mE{66u2Xt{oU(qHqMl}Zuj%IfDZQn6AoQF(azcpih#<%9(NnWH=Z z@rl3v_-%jZ|A6&Rhi?Bg|NKi0J^t_W<@e_o{F*Nv^go;aNAu;Z?9X;Tefekif8P8L zA@JAvBKm9G3%~N`H-^){-uzyN9e(NG`MT|Q!~di;H!;;;&EVg87xyb~fA{y-JpNC1 z`#-K3H5J1@)`6dYR@a%IJ18~vPm-#?xc;`ZaDMd+mniz@`lHH1rTXJnJ9;>}dY*=w z<{up}{c2B3L-X`ktX8}F9X6`9zuvAk{&1l17!}po|EDonUGG-&>UUl6+otgU+irFI zYZ*iJo6|oB_;3HSKS%!%_(R|ifj{_QdV38$g3EE8`x~b)YN@z5=QE zcPR0T=Qn$cPV^<9%MPbD#Dk_ei~Ok!@xXO+o!*>-s6*%JY_nt++8AcELrrrKJ!h1( z$v@eJ)L9oDBYqI@hCK&L(RTvw)|2mj9@_IJ3LSb6?cc2|W@$kGGvztqLC}9};zXbQ zHv;~RYlocv3IXr28zb9;KSo5I!;Roi*Swrl2Kdi)zxnxg@INz9Nho-kfJ^e%?XX)S z;O?9qmQ-y$sDfW1a2JfXG{BZT0epzNHVNARpOLlOI2{%VIMrJI`Ga2xczwOd(&pS& zw3BEp+yOj;xFa5%0{(g02ds_(|AFVrnD#uhFIGqmo+RM7y~V&Ah*#)^&!+tlufL{d z*?HA)3Gl!E*iE{#Fy8f(ns181Vr}xTl67L)*LsvFN3vkd*;*w@r}ezS|1lO#GN!wlA^f{@fD{W*=wAJxU_iW z@C>UVK3b*vl7qz%-}qtf#Q@};%~2!wql|`lX7_FmHOMP#-`=tFkS{xI#nL~+=j5TY zw^+c>@2%0s7a(4~W-H<~5C>g}_^kU21iZLI`=|(T{#Gu$Z~8Ll4wJ=#a58n`Y#yQ-4-?l7#Tk z5&c=1hwe+8>2l4}r7Y;OFhnAy3U$HYcyu zYC;Wq*FD5{5fO=bsizzKkKa(!mj(V8`w)sVfPdh}FLrOH33yzEsZ$V)cQS8DpRT14 zxfA(|*XC6rYgUb;j=(2uH}%9J$oq&4T&-`6fSb)2m-GYA9U*lE%D|IHx;4`Z_?KsH zxRm@h4GnGCc#iOafGg3o98m5d;1ZW7v$b+FP~@`xnE}!3f15A&-B6b0CPQ4gUwAGH z@^WM9!QDdqhB%4(f|WPq=MryGSUbc~_y`r-Yj()T(@`~k5T{Dhgf)z7Apg&IX^X!ZSHEG%iT^l27d9SVqya@&*NTE?O6Omz{>rK`d%rD0sB zrg$7XH7txf$z zG8;2zcF94(Q1guJ#y_G*DYf!ucM-O`c~tUzAo;Qyr`yPDQP9=?@iX4x+ z6I9}n`c?nl_b}f0D+^ovfR9c>_X`K$^G2LHC*vale;}&5-x}hz=YA(~-FY`8T&&=} z0z3~fJmKumalk?bj5O+izucT9vn<3@Wq-qukS+o)FDrG87UFeb#G}7PVn3#_*@ZPSrWupcC-yRLTyeo=Ac@4CQ;|HYjNCg2mr)Tg5e z<5%>Ub|`@Hb8M)lWI{g9DY7{-j1utA61fW@Lj?R)SlA##-{135oqVi|M$i!du|YbN z3jCx^C=i(#d5m4Tcg9d+eHKQNNojd}=oMC&+MnJQSb_;SuQ2JuI0B7n>;hn3$!^yx z3_f0rQ4H7Ll6yhMx?7_=8vSr=Y@}-XvN5m?sFaUwW9MkX7^nhHFM|Kx2``*Jp*m=S9+(ML4V@QFwS-F~lv>Y5vG9;5ntt zTH_2nJBhE(=Kz1v0c+MxK_%GYGXcU6i0iEHM}cqe33xLxO|Ko|m25$OAOzy|*L)N_ z&RDV=;$?d>;0PJ^X+PHI#lsK>{B=(+GjL|-PmC0R_0WFAI8zk%?Peb5ccH-ji@Wu# zBCL;I6$v>6;Co=2Bgk|$A77K4UIIUUmX#0rz|ScoX09IA$7kujp8|oC?2byte%Pl; zYXYP{z47VYwLXW23ibX*hP?KKO;PI0<_nZkxRW>Wo)=?evYn&+m5a*OIqX zVd=F73o56oF>aTSU+&G+Vss(tDz0>m7{iT{@-|rSFLE9Dd}?*weZ?1jFfOh- zn#ndS{im|DtY@BO!4}zozvid)adiBXlQrvX)dD- zSYZoe=293h{g-*cS1?|^Q*Y$kfKN89&y^(L)8|roGY{fQ-)SjDg1XE6x=Cq1@ca-I z7VZx`$GkNk?FarMWBSM3fPdN3jr+?X{v%S)yeA><+XtMZ&O*F)YEPbI2+#dnKIR=% zZr=&}M)0GC$FyKgMdTgFp$uUxuut9O$A91ic2aVRO9hRt-3`gmc=y1TK=P!x`JZ}-0BQ)Ui0|W7gdPoQ|Xzpq_Jv@FZ%A6*^DY|(?r}# z?Y>4V(f8HD_vAXvhpFE1Dzux7DQ|lV_RFX{pSw>BmcmePC>7d-t$n^ng4wqX^Yih~ zbDM9$?&aFne1vx23VOy@>$`CeK^mQtotWPq(I*8`ZP>xW08dZb9?Y!r%I!;Y?U?_H zPM67Q9VnmX)dP8!>SJguCEAZY%u>A>__YH|lpCf*8xLUk($OuAyq#ET)Tjtcp&cw=Ct@; zfIpG~QcFwV&*{S5hwQ{T;4Xy>|5ug3?hIgueb9mYEqDNHy7<2Anc zNZ>Q@SxS7EWRy~eS-FMFzJYdqm&Ge0VgB9vRCfdio<0W8Q)Ql2VfQSGXbXV<%lkpW zEx^D0YFVBMtOLVO?z&dPy0HC&wwWlzE8+H#CflBzzvZK3(zEIj*f(i2?RL$>dfm;D za5V+iD+^s=c3sHll8@pwYhhjE!+QM(!CGmv%-KP{dv2KjaTxNVdGccTFyw`x$GP!q zz*nWyNZuXz$ywFR@qr(R?zu=R_~EQxqLBvc8C{MxSqRow>6_Jo6JS|5lcmdGy<1}W z_`)96LcA3g`TkPQjL4D1<#aYnZY(&gWi>5`j3+h zPtH|i!MzQadIU(=bC$q>)w(|XsXAr_96}jxs!i^#$1IFHEIIC$WA*}TF3Qk0VS=-* zI%d9=*vn7%ro3UD_aCV9e!R*jB}JM&s2NkcLcH6OP=$?*?csCxZoxz&)%v5qS7AO` zVgZK0>#aoz(+l9W!*D`A7Hh#gqM!G(O;%y|uVLbr1ud9{k-kT1R~7bHb1(aT;OTGH z{W%YKas_#4U4!<1NlRQGCaSP==kZ}5=wDfz6l4zlkESIA+5p$S2p#zt*w?1`3L@pe zpSJH^mc`)D^O|wBec->xX1~3w{-@dYg{A}dPf4V=tJiV$5wu=vCiR$o9`9*M81Kyf zjREPvCq}=_cLDgAIl9Hi!9F^&_6H`q&Jdqmt7)GGJbB;A6chtb`hB*Z#=!reR@a^F z!2kPsfi+K{Zf})Zd}s;RpR{ss&bUFmvTMRVm_fY$8sEW$S7WTOp36iNPttP%jK!r6agPoC5lycZflp?q126Qvv=Ssk8g`0^f_+2|*(8 zmG99J*9Lw{)td4-z;CIpzw-_7X|(3A>Vt77#=RVhhjHu4iamb^^;}yx?}j%pt{Vq> z`Lv-POffx|bQ}B**v?h-LiF$X7)q%=&f5*wcOI>`!=O%fREX`XbWO#Yzp~kNJqp97 zQW(Q4H)LYxKU*nOTVBS_Y!h~mhPqpDgS<;AudG0CRn}bwvGZ3L33?JHrS(AIe*7 zR{);WS=w?Npgkkwxg&O0l2Pok!Ly>NsSfGSrdZZmkc91$gfv)7sZBQ{xf|*^1}LsSY$*Wy<>zV z%JTEE+{B)bvE*60masrG4XVvtU$Q#>@6S49hEf1b+vVU7*3{QJ|uOkp2uN-%g5r8vX)1kP@l5C zY+{7?##KDEz5?;n3{F&%h5F6kkeH4`d}_4}iUuGbN91lRiNkuZQnGmJ0j!h9H!Q6r z!+uDTos!gm{gwCL9qUTi_nQLl+_Q)E{9WcJ-9xZ0%4Cmm!l2p*OpnunNu)6X0gr-3*f)Ju9VHeO@GhF9)4MNEy&03S{B~MaNQHv`pn;D zR|HDf9Lim&^aN$-$-VgIe+O-dWU*JNc#Y0n4DTC<`4w=No)iJ&;B(wYr#trujoZCG zROyE!-=q!f+BKO-YFc(=&2Tw#b;;6VfN_q_RbTr7dA70e@!7-KxoEtLHAdJ z%6zNxQJFzrvJ!VK65})!S`T^1zP_8R40+Z!bo8@ocrkL=-V`3wUXN5?HmLU8D?^@) zkv^CF8`1shejW_=ck>_Q!YjZ_zUgGTl`W3UdE3sP-PDBIWK?Jwi*Xda(K;^oXcKy} z64&bkJe7ijZtjGA5?9}K(+k@7exCBb2<>;0niUl9m!Z~^kNQ%e|60!2`C7<)CPEL)^Oaf~SaZeNh~Kt*9RIFT{*xMG5jRV$G+(HrS^V4n5wc30w(B z{MJdpea2lz3Bl9AclH3iDhKe}e0Sh^58RjBoqa#H0yxfGTRtkYx^G`JU8sZnbYds# zD?xtF2@I%f!}Ze<=ip9B7^f`LtLKN|exaOL^I=_)zvtr=i`v9EtaI}HepX44FMc0p z-XWS?ELd!0ro(}RRf(u(`0pcONo_uc9^>`c&EZkW)$98wUsgyZ;J3(oq4w!B)!2>6 z2o+k@x7g*Sn#;`|4cM_AKD(Bd+cCyt=lSL#uk3qqetQ_-oW0QiQN0=CskJB{;p@i2 zW-}BcP2OVG9J>zCrT1cL6IUiIpdMU&e}*X^I1GqU9GAYgV*c%X8X9^17^&x@YN1X$ z<|_J@ohoSn8@8nHtA};*lD>(RGt|p3mBck)Gk0L^AJ>i_SoaRQb#vP$1=9}f>uy*3 zH(Kwoo^P4{iBOlj$vj(Bgmtn|VVhYwwC^z_PB}n(TMng)X6QeA`-O(J&_C($mqr3` zzac%dkO6gMbB}7&h99liL$|6SPw+>(z{~3e_+LCtqq7_QA3@S>DG;9XiW03F1Z+V~QhIMtbfC%FfbuMi_6Dp}2oM?L$TN+x{1x z6Ka(IiEy>6-- zq}OJs(^njAX^cDI`H4WnydY?^OQ8oeKnFQfEY%(m@FcPIy3Q&DTwpMH??%wIK2P>p zfL8Oq8KYKZfY%=iG&;|2fY*g^(`^K;9c+H@1Zc;zX|YOfyYa`FC0nna*oCiWWEB(z z9VIiiDF8Ib!S&73*YxqdO}g=S_4V-0Y7Pc`pvSTYzC?rWV7nqKsHTfopJAMPl)Mvv z^;z;L8|WN4u2Z)`i&+IS^S;)`$*Htc2j(^LqQXkH?|ItzvP}Y=9_Z^qKK+wF)bTw! zJKi6Drh@yEY?Qh|FBfUrnSz#88sH@wD&l>488`DuoA9$%oP4i9C*)^8e*t>5YWAwU zi7;xXE5jh1X_})$sr?&zV6>1>kkiAWwn=w8IOUP6^PL zr>0E>fY-G%msJVC>*fm+j2*NH=P5T$&@I!;BYeP%<9(m%A>hR|6n>8x^oFGf78B65 z4_nLE0WXiJAv!PMRsV#_iw^Yr`=y0PKpX9n3SWR!;-Gl&)ZvF~I8z=dGjjTXk_a`?7#A(0==Lg+Bo=_wH5`72xH?nrhk#`f`rE znF#13LGR`3fLCvU(oKEfRp7a3TMSxw+q-@>&@#p{Ax|ZLU0NO)XR)-aM^;GPZ69!%=p?R1Yv^eWU z27b^_bRTgs0WX0wSc?+yQn=W1ni{m7`&_K_YW`98J5d9#O}xHSdcX_g;ivvuX@IMS z4SZ4u9U7Uyv~+Seen7oL&I))%d!^}2fc~0Mg%d!(63&pD171If-1u?ebu!zW^Z|7G zaQs_a&^OJv7RP~C(T&XJbHJ{_0A7zf+QvD7SFlrUQWj{6 zZ9~Gh&zks(?<1#7;3Y@ynnec$S0-1Fw1}K!fLChJ(zO-PlOfv`L_iD0%ydixuT%4NhZTWW-2Rg` zS#A;lJpo+6}04&&ptiCD=<^4Xg~1!`oQ5* zBj{wVn4sOD@3q;CwgE2z-3)1G;5CMfjoyH!C0yjT1-;dw1WN{9(OrTJQ^4zip9CQa z^pNbylii@r9k+YM0xF=qLt|uuDDc|y z$wg=*Xb$hub<&`(Ej_2B170c0D@?2T_cnS&j0JSkxdshc(5Iun?VJE!jk79gQoxJk z#TES-v_sA1b6Y_Bx(qmX0I$fA8!hs{>)_`nbpxQQQ;j3kKxZoNU@He+w7aLpS9!@* zuHV}L+W2LyAOaonX;$eK@OoB#hDINF`AhEeD+O)cT^qH^&pYF|%n{(#BR#(UHSm%% zU@&k5eRxY3%c@qSlI}kUymp>_v26%=g@-$L9sq3|MxWOYn&w*j1xw)d)N$^x4Db>q z+k2UU7Mtm~%nSMgUvs7r@H&teenn%Imwt|u2JMf}$OUy3< zUaw#3ZwmpvDE_&T0-E;64;@?Jb^c^IVU^d0JsYZ=KxdmBA6exorxR$n7kF)tTcg7U zyzDkRbgc5T=P?Ol1YJ?RPMQ{YeK5Gq>jAtLPr+lG#=^L3v(;&4&{aN-=hT6hwB@kT zcHq@A>`<^BG$u9?s0Z5UnrWd3@XBClcG3Y}hqs-NmIB@A-?(-U==47EURL0>4POv2 z0$w4>OIvwCR|=S1wgl}hc)oH4czx+J9W)1C_Q#_hFv0#y^UzRib^j%}6}xQ8&<0*o z(h4sPfLHA`whsvOA;n0J)p~HVwULKB@Y=VnG06;gC00djQU;xzI=tl|=;y7y2ZVta z`<<@%)%vBt{K1gaYW^iy-8&AtxX*}}3wZr__=xQU@WLrwm4cw(aRpyJ4Z7hQ+D`-Z zi@&%h$7;Q@^W!@M4$ykF-x7mCk6Fj$O#?47o8*PFz)Md$;nsK1zH%)OE`lCmia%Kd zyjsV74OV~`n#j1`2Kug(t1U0=$A^YIm~H^CpeSZDM%XV->UWPm2i>o}Y_PijI=0#O zy8zDLJfY(=%r&Gb8g-m)0+MvgsJ}fK%FQwAYj&Z;+zR!8>eM^D)UTv2-8k>iioqTY=ZS(9nuK@Zue<{-^;uGG##gIA}u| z-EalqrQ@p_;03&*`IE}kK~Ha@Ph731d8K1@HUh7(dF7TM;N`f&wr{nbt}qXgzXJMF zu|Fd}@Ujk2n2rQq8ZxCSQlRIKHEc=*eV^NkumFm~6eN)iKV$*F&z)QjRla?~@Vtnf8G6Y)o*_%V= zpj~qJ^W6YmhX?zg9s*u|nZCZO^>c8PZ~Gn4pR|@wIsvb2LE}R+z)RVOdu9am{D;SO zuVMcFNB3_Z{Yuka|9|c=NhroXW0<&6Ot%=fwXzls1xt9ush(dp5cKLJ>N_nZh)g_<^9Vz{2 zPERgIgJv>`b=Ad)*=o<$XB?#{+0)OUb$uz);nef%vn@mJeiapN$I1|0SnuT3Cm1U0 zov;ti!VvE>{R`&ja3rQ5$#6apN5)dLXQG9PXv%koVvjfxnX9rfI+Ka$a(%+oVZ(mreK$a1>^#c03bQAw&m!vc2)R08K zjoOtc$n3;6+MY^eDHeXeroR%Iw6D#-6 znUf1kFBT#47pa<@`eIZhc%0Dswisz$>sDbBEJcg0Y#;B5m!h@z_3-KAW$^vFHX@Bz z8R~p)C2N_7Atul1Q)49JL}BARShHhe2b_spMeCaEmREAd}!_@X%u1i(u_T+lu8c7Oz zSn%Oew-g0k8IS*#Qb9ophCI`O6biEN>JKn?u0}VQ(?qFHRHI3swLJ9O3sF1oXGzta zh3M-hR>hR4A{0}0qEjZe2+7@iR7>AkjB4J}hcEONBc)p@LmQ<^(RhJg@OAl8g!c=$ zjQf_MoL3lY}8XkAb2V3WioZ+xPoj!qHs5mxW+7j?}7}t=u;g5sh+? zk%tlyy*%TxFw{&$`u?qZwAzWtn_d46*^h+I4T>%u2q2*iSqewzH;~b>ekYCxqGY6D zL!R))$jHXjcMD@V8O`+^->&0Oj(V04<=%2GM-6NG3jF9Q(3_1x?Pr-PP!`MGrY(;v z&`fkOO>24uqV=IPgd0?%yyVu0jYgFyt9<5o)MzDAbLmmOG+v2f97`%JBdgHax=AJ0 zD^;kca;9}ohJxmUM)cj~Dd;rg8X3rfQd~d7bYoW`5;0*ov+jBka^f1FVU90CjC3dH<=z#e2_e46r-zDBp~?h<(PHmD3$tZP!C4=+OvoE?!n>M)d%`5bXHV@Sa9;qtAkI1**LFXj@5 zqo@AnX4_PW=+daQ&n0ytq86~SG4CRx@==`y^*$n6i00AG3??B-GoECT3nX-ZVizBW zBpF${*~nwkWc0$Y;DTKh88w9$u++aHBO06hzN{1F=-9cP#W&r{k-U28iY`k95?McJ zd7ZriJyNf(-1oc!9W8w>ESgn;3b#CYK)<&V9m(>2=wVukJX=ao@I)m#KKeE7)>I|> z=)1A3@oE(^_0%V%$5tUmO{Fjf1q!myw(fkkg@QIu9>t{VDd<+wRlkP~6x3WX#j(Sq z8ttZB_f7DuM#A#N(d!Hfk;ncDQwBmI8v4jepPNvGRMcm+@LNSl()mDu$#5~+la&8D zV5AtSiM(ZQ!KWQZGdhtJ5q`M3f76C&&Xt^gdT(N8cb3-LCjrt$T@t z+U$h*L!wBiy5mWSf;<`NU5gMHRD?V|e(2~-JsB-naHwr*CZoKf<`1G?<*4det-F>F zJV)uEb1jLp0=bBcy9%wXKsGUgA8N9J?|^vu?8^#t+ICp@?Y>H6zx}Np*{l-LAImGL zo~}eF&`0+A49x$z?j5h4GRiVNQy*q9zQV=mv&UT9u1>GQ7`Gc+91lb#U7iB2B`y;BuYltW zy>RGt;My}vHt+|IE+^u}6xB+R6mLb0-lZ~>Rwhnc4qP{fP4;oLW5{$zHDz}zhQ5T} z3G@VxX^dB#ErBaV_aWZ_;HcXwFg*rbAGhLiJKq!0(8OdYdPhWu@1D=+k0zl*}pA(2wKyfTPD-(F;ky6|YgWIR+f_ zNYURv1J_WR$Jx?370A3OIG7!H#t{7krhub^!A^f>u)D8^4!;JD`-ThGeFd&-`&iXN zfaBx6jso4lHO7`8u?IM2EQo6t0oTh?88y636l6CSo}B3S+9tMuu^5d=EV85@6 zUJ}|}h+=fl_@2@$L@TmSO44ASCq91Rl?U^_dvVrN>y0%1wisq--J zzwTjN&V_mYlJ_ZL4CZ~@wE3Y=Fwfhozqek7d4JIMwKW}ZtUq$++zib7PX&svnSi6I z{YCGsF;(cn{i_%iaNJ@NZgB+W{l(+1-P177^VOUR&cVFzOOGAxhIxMci$&f=nD;V4 zJo`Cx3enVShKga>m+&VMMBca}G{dRbelxlVrOjQmTInlBe0tIcIlGF{x9y?tvo@Ea z{GGuKhSH^odeEpsB%lnPru%qK{!AI_4>1^0tijNNURPie8ABq6sG56WAKF-U%UAO< zjzYYzd+I3@QU2v7Sv>_J+RRdAF5W>zl~)gQ*R~LmY+FIHaUcnKxKUE_&yrB?TzAGD zF*3T7X+Ja}NJeiFsdEYTp|aC%C{yu! z=_`;7TX%FqMg>~Q+5K%6_9f?#k9xs-D$!ww>_JsRC2H$tt(N>$iKNdtEy<2mB1+Wu zk9S}na#ZZYM_^yFWqz+T1^dveyr&}PWGTq^O31uy9R*F9hMpgTeQ8Qgvckf>8f6HL zMD09TjRf=~D4`mKNIE1fM_#oM&CG0e9=}wCGD5vu)x(MqJI~s2zP4iY_@TYVN<%SX zD^f3Blz{v^!|8iMs1zMuQ(Jl#eji|O$>Y?Mu4PE@>A)xdG7Me++|&H=HHI7>>hJT6 zz!80KP=r$`j)dL{cI=QNA~}Xlkzb&0x{#K@bfb}oI?SJ);;SX1u;u7}o6{s@w_?>- z;zdH=EA%SE_{pe2Ci8YDH|#%T!V+UC%==vlA6*N`h=WQ>dkX5L_d4IM_}YR^_LL^j zRG>S$Y`Ik5%h8WJ;l)x9E0F&p6-g+i0+|f8eoxy~i5TbIux_ZEV!ocIYl1q-H2#gQ z(fdlop_KF(gF0zb#pnHg7pjm2x9-tEX$mS@vFrEPNI|3KB{kgD6r|keJfc}iL4{1B zk}*&xEnZc=_8#h{fc&g@YQ;jtv)?+uMh1QdV$%<*>cApYd50};yI&E?)c9nt_of&< zP+QSIP+p9VdzfxN!dHs?b6b?TxJuD!ng?ENPQX*-@DJ-lWr(Zs@Wrh!F+^vROc6-O zQ1v@n%_BiLYS+%@kno4wG0MLCiSR^*)O^*Jh!PU-!DB!YHT;c@U-dt(Bl4-u=~QlSC*;%HXC8Z8-&n z^YP7>KwZx_^5lCR)bYl*bS!hAu4kp%_Um^_>8ZHkwy*CAcc<0MzJKp`0)O3Z{&i(2 Ry_1USA3wYLOaI?9{|nmw3q1e; literal 22648 zcmeHvd00*Dzy6M7h!kF9B4vz5Wb9KKWQar|nUa*LR1!%Ok|GL)QqgQ^(yX4{-pzxO zsX>zv3ZV><&a>9@toQo8?|aU5&h
    -t^id-R8A-+QlCYv1ehdG7mjUwiY`*U=j? zXzU;!^D|%oPn0LgKIH!Cr9Vy^$TJ+~Uf)I^h|u*EU1R!?2O_-wJm$60^m$2o{)v6b z6SnmEbw;|nJbWZ~er}~N^9;qvqXFC#zt{g&4Cw1>88UB}mf@dOT14xNBg4{q=UxkY zrjp2yKKH`^mw-)Q}C zG#sQ4xP921Vy<52aR>ev*R`uw>o66})!jK-x7_*uUdR8%)BdkJOoS)K?gQi>W}hJn zDk4IdoK1P<*1q_4<`t=l@aO(BTbwt$mmb#-Kg+z{N2CvXf&ujSzJ2<{(0j$4io0#R zVe}z$utZzm9Xy`X|JQrK>~~CEalU|C!^ZskmsR~&8^hyx`tuF``@i=4+20KOX5cpi zzZv+=z;6crg&3Hs*xD4fRfw*ZghxzMOh=28GeiU2h46WNw7>7syPzC!&uIMgN^tI4 zgcME)VMxy>z1-c&=$4^bn`hZ=lvof|H&*m9G|cfSeEFdS&h{MIVbD+wFSGmE$oN*m z)BNw=N|v!G(LgobYk35Utt=gWaL_Y2?4f>1T;?Gtn)$3MiE4nOP472;)_M#B)>wS= ze;9@w>eY|f-3mmee&Q2Nw=}`Hw9Z-eKc2vht1s?q$u+}|HlNhbI@O?Mt?~TxML#sH za;A#^iUUYyqOkK`a5FqT^0D`uuYpIGcIn0H9RRIQ4GDDt ze(+GB`?LTb|F!3JM+iRNVU_=ab@==Xzh9IZs#^^vhb*J*Zw7)*m|gdxhhcE>tVxUm zz8?3c?C&B!pMXjJ<8Ona8$h<-V0GywkHJLZOy|bs5ddkwdiq+%g12u(YZiXLCJBMr zYWV$1j(&Xh3g+kii!r5u`Eil^=2umA8}783Ii&7RhM>+hL$_i+V{eXb`lF)+=FYC^ zvB3PxIIaJwiuu3utlRUvVmf#ZIvx9Is}Oz&B{bjS`_eP=ViU;ivk$Cb0gk1Xc# zPTeoqZi?6H*8vxnX7hL&a$~}5@yZXbk6X#C7p{z%k6)MEAb(L4uZ8k^Q+F=rxm~eX zEPVju-N6gX4e^@v{NbI=c;zX@$t}pnzn7yYGXbw|FMMQm@$*YIO-shV$Lm~>kz|2a zi`vrS!<%4&EskMm6$EU=lx+&Y~zsqAv?zD$ol)}?V-W#ZLaw!>)kp2mhHJ{ z*T|i{^7FC~-!#g$UM~$zIlQdZIw=Y58p?~69T$&|U#R>et~(NGCnW6Yv z%ekRJOE3BeWZa-oOF4ML&Ja*um~B%1Jrd5V584(tHXdBtjCG8Xl3<{*;-{kGt;Z+L9TKtY3*55C`-gtQ|x_I(zKMRDkWOyCnAWKnStitqb?+cEPZYhEnK zEN!&gx-$vR^WHbE-H-vha$MwY)nvo5B}2AnNC=>S?;VvAod#ChyS?V{C1}qtAnf1LL~Wl=$&#j;(MW7V~R6!_Z#|VPv0=4WxHWVsjLeM z=)Si8z@?jL;@#(^j{e1HPuKtp>DBpY%hwSq#@qLztg~WUl0PJ($tsR|54Y_E55;%I zdH3%kUdPz0b6*S4$G9eM@h*fGNKbb=q(+(pvIq;I;ohQK@d*Ky@*(@;oN!`Z%TZopVq`2kBjQ&HA%(XyW1 zw~&^8qhr^gb7;!!JrxIHqhPUH+SZ$1Dd>zw?DUk@Xc!Q%ahco56jULe=|55I9DJ*n z7L&KW1=r2De|(9LHxGYw%myES`^-9-GJO71TF)-O_DDrry-(V;``?A`>l11eu4lpA z&iW@K+(W=OaoWaUd_4iCXWJ5T10eivMc!k4{k;;kS@XLPOe@saw0|uCW4*i`db3YJ za8|xu8-Bmtzjn{Nv12C;cwv2LM8#dSAgZus*Xn#&b16Wx!oL_^e)nst+O`+H&1v_M z>_|knla*G#ndE||+wYE_8h8_}kGXck^iDZUx%w+?dvFz`&Qy!)#P>_pDr3|~e82v_ zf6tsNee~zeN;Gh#-xk5pQWX1#TX|4uAzB%4G9q3o56uwWnEYat03|GSY-rt$(0L`R z%)K>P$bXU6Nb%DtsHk9>=jC&EkYd{V0P8#Fk&KakKyb`?sGITWh}YRWP?hGO7Vez_ zLa_ju!!=ouB_*+C=^g}A{1xVt?*za?$RH62P1>5a;@hI)gAY0wL z{%CoHRpb-X0_553EN#|Og_@(jXpT`YK=M%;SB>6QA&5VdWqjKo5)E}VrvJJF9eupA zlJN0$+bZOZ@bSkS7ds!o=O3A^drBRj|4hI2Z_L{6!kq8>)>a2(K;GQA<$m~jnzR;G z@59%ln=`fSyUpkHH*v?KV^aQ<^Vsfc?d$>!k_QJ@7G30eg0AWe%0&F z9ZbgjM8&i`D#83XMaCwq9heI@SB7lZmm-AZ=z9kyVm{X%8#VeF=JVO}7k3OX|4KGX z;^u72Ms2&}>ff0Zz`(29qwSPRA>+sgZ7Y1gZ0C)7D8TpY@Aar$krTd6RftlL%iZ>T zh*1A0xud_VxQBFQy{>y!q@%|wH*~UwrlI)9_a}eTPeE4c_F74*chQMNL0d~xB9iUd z|8|-17Fud2FPU`HA2sy$&*Me-gR!n;!sgst5O8F~c7^6dxVNtN{SdXgF#d7;t)t6Q zz*X{{&q~QO*!i->kyn`xFUH&N+Q0H1ye)E)&dR5^+>QOYZ|1s_QWb|XO z(tu;W2}m@;;N8|aN$9idvKfM$WOU!H*UrA2k4A^C?iasGfTBi)X^jcEje`5gwa-+F zLsNMcf{S*ANN@JGi#HOBk;js6qqi8ABctQTuJLD9qO{F?;|GB;NXx_2f2>Y4YA?4Q ze`DnXq~AR&Jw)drs%^XS@<`%CbUNv0X!!gm$b3#q^{>(>q_V1Co9?$rq!%vdJfY$- zTG-uQohDO*H2fM(Gmbt*3cvcL>VK|7bI(pmDGrK2Z~Wg?h*Nji(FN=cwp{Dvq_P16*zf1?XTIS>X_4j(@@eEae6sur=xhUsq+Zw!1s8UqigMFq^YG^p) zKQA_Qs9>zw`g7L$Js2B{jd63v>yh;1^2ZkQG`nWqd+vtu#+_wi7`uhe4;9#9?6Oe) zh&f)*X|5J-!0Ywe6yCaQ9?$lex}+M$>-IHIGQun0%fOKFxB7i_oh|X&(l#Uj`(?K& z?bE+lVf@#hE9=CVdIZreicCF1`g~G{k)ot7)7KcXb$D?Jsmnd8p-er(ZVy?e9wDb~ zD_e)}Kks1c(m#dN;esn)*}6237BKY)rw(mq>JhdSlRC7qBy~AClGLH%ysKAz4&eS8U>eIp0Bb?ACb=dob)MeyYQio@j zlDhQUbf2k5^zH#$kEnS3Yqk!ngrqL5W|2DN<&e6Z-Obh`0tdDp(df;j4lj=-b$K#~ z)Zs1Vv20z6so!Ag5lyXP>k&B~BXxN6IH}8}Dx?n0Mq9FVdB%mUM>O4WEmM!ESBKPL zi3O?4HKR!#o?gD0t;>BwrZe@3687F^>Jc^0B6T>wQiZL{zusLQuVyp5KZ2LaBxZjE zl@aXz2s?65GW#R&mb3dK{E=S4?2ph-=)mlc&^uJf?2iyW+J@O5K{;b0vp>RxmxGx7 z5#5N8V)jRr^OW5m(LGOge?%SI5}5rFWf!vhBl>cU-5*hlvjej~qLz>B{)m!Bvisvd z`sa!QtJe^m&+0Y2Z9Ks0^vc;OtZsKD5uI+|PIS9Mfz@j;c+cuJoU|r7Z8)CjcB2Z9 z)oCeFqT9z$vw98Th}CPj7WR$R>2F$xS>2xXp6K)=PompH9U8vWqT8YxM5iB?65XyXWAz$UzYAvc8p-`Jn$_uHwM4fw z9&55XeL`bXGl>dN>>nBv6x zN4Qh9f$@*<-25QpA7R=uOU6Hf@?K#4BNRF7F#Zv)U3$&%k@b(rvHuCiKmMa0`wTnG z%ro%)F+0yd-+Bplj>$P8VCR}lH!{a)`4+&X5Sfi+?o>>P8f!J3_Gq#8}x zIcEK^Xm+l--I2x2GswGxooCP#IWor>P1?cEH5WhCv2%=4(0X>R(P}=y%rnSj4m;1F ze2b~<9CJgvnVoCQHdnH9OsvcqcCK;tt6}CD6duXWGpPS`GRHi$EMVuFs~KdD@f|gr zook$)v-1qHGGyl&q`mPSJI7e79bxAhsSCTXcU@=m9%#MF<~`u5GKbB1X6~|f(`=&IrxliO=lJh*fNOGTFIhjn}0}A&ZF?kOh4D zWC+Q99?zV~<~(bqNbd9Ot2dMPK(RTS_rQ`}{n(sGT!!R6@+V2o)8$QapL+``n7jup zo7lVu{NqW^(=eChK7CC|&LbL0a-TTsolM??nm-sYc@I*Zx|q#*$_5G8+()gO0h9M2uTv!F z@eL)pPxiXwY|gWGC&_)Jy4bu2nN_fP51Kl0Ae-|j1d-gQ^Rg?O^JE<+xzDS4Y~F)J zAKzv29yDUsS~llFd0lMo^Y?n3{e{gRp|Fb0A7NH9n?HhJl?9VO0-wj`k8pF)awdO- z-P&ya2*29c{1H61CNudXeA&N?$sgg4haZzaqFxmxCVxbakKJSPNAx^^%^%UMwI)pd zh-y68{1MHqV)I8Nm&@jl=*2NMe?(H2XPEr)ANBaA^%axX!?R^T- zScS>!;g%(v(<6fvHn&G{g=|ib#CEZ{J!)LQ=JlYFKaa`lp~s!g>Ct=-Hn&GF`uQ?B zJ=zt_=Jx1^mnM_f!yDIyOkNMNT5L{_uFJByJreq^WO91c*M!aOk(~V=Ca*_bPEAZ+ zkF-VJFgZOsnaAe#$XDSKlhdOYhHP$+Bt}^?c|Cf%&YsEZkybgI)1w`y*xVkyJ}A!Q z^yr%}o7jBDUl^OyBgG4BZjbyr*qk2K^gLp6do=6L7bdSqq5*7P zkG4sVVRCvjajX%O+avkaY)+5veW+$~d$jxso7bb7Y;Pv7M=>XJn4BIV!=X%W|Bw4i zJ}1O{YK8o`Y5l!o|BAVfbMO`Z`uyJ!_xX;X{(pbvUladlpA$OAe%c35cXa)4d`_rx zCZUiy6-V>`n9m8h4gL@A|D3-*&++u<8~p3bztL|7elzf!f!_@LX5cpi{{js3i~b>| z(xcD+Iw~g2pjw}AzqPz+CdP+7mUS$__^H)AZ}sc?{IbCti!Igk`8_*oYiDAd{i`*3 zBgU>iU61D$F6I9b9;ogtv6TO{|4{Sk7=HnNDtq<#3G*jk^R?FH?|D08 z?<9;hU(8$Kg0cMgNg~tE>+nB>yEt6Y(dN6|6?%-u*nj-6=ch6LRDQj3?A#^%MDcBF zB0@F!g&H|)#4-MQV4d_ujNRA0wH}_T!4Igs{d`^50)Aygc2rl227kBOY4>>;m&Tm( ze%q_g_p0tvHod9FZ|i}hMmdI*7|8H7i+Zj5W1O|2iXtOTHr#4zDOvHQ@T`5d4$9#>p zowH#C#t+7RiZ{l%WnIRY(U>pG=}j?qn6HF{QIUf%o)s~q#Dc+g-ep5DU#De*72Pmj zHF7_n^`6t=tEr7#XNR#%+w*iD=BwPjf5~agSNw@55})TR;UA5$KIejQn8V1xX3W=z zFB{_~W4=TdZxS!UII-lss079X90T@cW4>yw7Bs40zWT>G@PBc4P*JF$V0a=U&|vUHw?sl9Z%wWoX6O$9my59YhX}?RN8b1q$MJlFCCfBBQ{~amZO6HU3>KSayudn zS7H35-z1qYm@kv`>9?#fU*<*=mcGGQN#?@$O&Djn>Yaas`Fd@rV(WnUO4Zvoz8>RE zijf{x824Uj@lM5jDOo5;{_0%7zc7DcSt!PxDpPK~!FbfHx|$%&SNaFxz#YEyl5VDNmPTd?~BI*c0<*mKf+Gi}_0a zu=JNb#zSw8D^SDuw$9j{6EI)N6KCYwx zeD!T>p4cDboN>!#r(vA#yiiUl)Jf z-L8xIvhSl~{2b$tBkEe`W4yJ!N4*L2HLd-|Pb17%!=wo9M;N~u{;X3E<0M@r!$+8} z@oS&RZNq#;9g6lU!uV%R-;u@`qjim<=P+NVUy0S#VZNTd(!A%4aY=C9>k^D_K6AD? zg8ABbQ$FKIoEksrn9GR0825dZEAtWK$G@&!vch~_=tiS8F<)MTMlRZhanNDaAa#tx zuI1_P!F*Xb6?kf5zKZRvOSfa3GfCp&DvZZx%zHf!^W`q7aB3pv%RMzyg2C79u6N2{ z9QxAauqfusA~M2B9`n`t%KKs;j71f6p2%XXGXBNl516lZ|Lmw4n6HC6F$(P%j~rN( zI0fTTZb3TFFkd+}9=@uWuj;-oK2I_3?0;2ZCdQlRi;0zCz7B4vb=JUq1=hV!5Mo>n zj#t$%Ht1`zkdOJg@HoWY0Q1$q{*ZSH#u|^>Zs}rdG5LDhe$3a`>~$abm@m`V%bRy# zTv6LGF#}_M_u1|pn6JZaY8Gvnub0-Zo@~Y#mdmTQVjLQ=v|F3Kmq@VlHezm0ix~Yxyvu_*eXG7YrV`@pi0`2Faw4WVY zNxy!k{ko9$bI425&r!5r)nAf+RjwobJSLX(^REGdv^lLKh z*W^p2pZ#n}KifDEy^uB{dXch%=t<~Cq9Gtmov zB+-+%8bnXBbBJD4niIXyrFyc7>d8o|7n)QrP0lwlfzU`c2K=o zLiJ)C)e~e+^klX?(Th{(h+ar35IxaaP4q-BljwyBB6=}Ajp)hw7NRHVtwb*_hY`J) zm`?QM*=?dHt?P+jIip4V%1Y{IUN0nmX3Rw5SNKDTU(q~3{LDV;XQtQ?zp{z?6<_LS zHeMxu#)A5lnbfcN8xub>Y6kH$U#VYdQzd@oEA=xv)X&IJzv3HB{K^pOXI{DyKeJ~c z@heV6#IL-yCw}Hj7x6Q)QN*uI?IeB$?1`V*V@dqXk;BBV7}gQL@_q{OGbzJ~pXo>> ze&x$c;#bC`5I^%Wj`$h#EyOR+aVLIxI`z}yqluqhYEJyJV-NAmrPNP9pnkenjQHhi z)GvoqKW#z%v?}$>qo`kAo=E((*aG6GKO_^s94t-zvKjT$Kd7I+znJ*tOk3iYKTtoN zM*Z}a8^kX!QYC&lnfhs0>Zf~DiC-S|gZO1l>ZjwL5Zw7ud@SB0(4Ezf)aCLCW=*iPFV7k+eE;FSJaJbRncGfWy znx6JwHO?gy%qEUd4-;mA)yC2?pOP#HEBaZrN-`UY#^0ZDZB#b6Xp0Z@G0TB3*7YHp z7CBJcKG;+{20?0$xYU*e1j8oBSSB6j!`{R8$3D35LEiX4zxOh^(09ZiGmYhPVR!jK zmr`La$hxaFtS!og{j(caYFi0F7-%x8d9MH@*6eXI94v%yrDNi|B!#d;y8=GO3qko+ z$0)zMLU_twvUBR%JUE&BvMGLj9xV0IdGoa^4<5u{p1ZX>57O-QR_qSK-@7qfeIYtH zAI^&34Dnh}0I#1MQ@f~907D!1d2M`F0Anpx0%kWAz|ZT?ec!kg!kiIT-OsxhLXOA| zE6oW-5Ho$QYwDyTIH+1xv6Ei}105%wx+N@v1NCrp>-J)}QGN5KvqdqCn%HG=O)Ucg zC**$VoR>Im6 zE*reXwIvg*bHHRtfNi#I4*YN#pC_A&KX2LMl9QN?V0(?!+-^5M9Jn7+-1h_@`XA!g z9iNg5IZXyHN@wJP_*26J;ib9wdyKsDah3S{Zhe!rm)HtGPV=UGqk{mllyp-!j}SuP z)0#(OV}v04dBv)ZG$9;baeeWbdqSwxxA5vQ&4Ww+b9)bM%>xM~IT62}JSfRnt^7nJ zADZ=g_9lhq!-ujurF%E>;k{*@$|0=+a2pq54!Q;KEc&8-`||=&nKyQ${i^~{%JS)$ zaJ&#~4kll=@+t)16)$YYPcDK!2Jh`=PAh`qBO*=j^NZk}#@*e;MMa>k(z+|yvKXdM z?Z{qgT@0@_IvqK*C=0irmqfed>$cHnpu59s-$_MnU)^JW-J{(URy1jmB0kry?Bu`yY0HTL_Q|n&i z?|u46ynEbM04>Xg+za(Cgt+SB@U^~$aAjSbx%!ME2%GfW>dwp}7@n#&?RH5KeAp$( z{7_Z|@s~PY4%kx+9b0nsWcC$9YQ&jy3fdW9`O4~Gjcx`QEMDPg>7NN6iC6VET+D=f z<{RX_nzLY+@d2CC7gL|x$#K%?AlzAmOg*< zT0<@vJ&iV-?1sNT>QcA$;&B0pd_11wBQJ#eJBq6oP8Y&lxm{5Y3x&|DA5f)TCIkh2 zHSw{#^B^$y?6M#G@*t<$)%luuK1{P1DE>wwA7)+YkqEt&4+{5(Pah& zaBxuh(c7yFV2-Bd%ZzupF29Mv{SO7OFjvW1^lTx#IbCF4<5vhN%9eh!6^p=b`b_>4 zr6RDtZ?w?peh~~hS1wmtQ3L_fck;~l7sK4QN^u!BxITwldnqr?fYE;6#>wbsKzG(z z-G`SlA-8ydZ*5>E{Iq`h!0uHR?EZ2>{Yq;VwAL;dJa}$4L@jQ)Xr`VGO_}!Nubjw% z@UV&jvR*k*vr@5X#A5_4pJEbTJwf2+Qy4tgpAU9@8pKTk`0&T(uAjN{b0IloXMgpD zc>nBG_V3@E3wM=<)=RbI!jQx%s>Y;nJjQW;@mmj_>GPU~MTnGdnHbxW2E%ZDqMt|os-$cNs? z8jFS|> z5DK!LEaQ}mV2a8aX$_SkI2V4_?^IoH%*E5vIE0ouPlt*vMYZm2Ehw?g^^2qVZ@pzWeV#)5-jnt@*Y}pR&$DUYXVO0B`o4trIoJ1GpL2b`llFNW?R&1zwQ1jvqkYcx zeH!g^uJ6lfpYv$nbA7H(`~DK`bFS~XKEFr%ekSennY8ciX`i2@eXmISoa_77w9mP| zcc6WqO8b5w?epEV@3}tb`d*6a5T{F=4#iVlil#dBlak{jO>X0+lB~FKOsV-$x9pZF})1ebo zmpZ5pak>;sb?7D4r5jX-u2NlkPj!gXrH@pHI9&>&K4}c~O`J~}L46bFlQ`exPJNOo z^-Y{l;(XIo>XUk?Z{mE?HtL%Y^+^M$Z{mCs=bI)`pTzkl&L`cYzKQcmlc;YBq(12( z^-ZPJCk>{)={@yHjnp@JP@nXH`XA`g+dCJ5pc2i~9I2)YlhNA0JA6{a5PaU8%3npgz8a`ucs;$8)}(^YN3ZujhO` z=j#R3$8)~EoBDW3>gzclZ%%!^7xnR+ujhO`=j&HdAJ6%E&d2*vU(fk?3F_*nkN3< KDSiC^cKJ6VOgxGJ diff --git a/integration-tests/data/moist_compressible_chkpt.h5 b/integration-tests/data/moist_compressible_chkpt.h5 index 22ed1b5bbb73fece643226430e654b991fa33cfa..0f62ecd6b82d1059b39233d5f369d98919b90194 100644 GIT binary patch literal 34128 zcmeI530O^S+xMI2c@mA9=Xs)GU(I_r4~im5qb9UVXcP^GNE)P2LK8(riF8$DNQy$J zM4FH$(m=j*KlFUxb9>+CJ&y1Fj`w?q_I(`v_iwMY*WT;5&g)v|UTxh6iDnk;EUQ?k zC^sf1Dmp5zzujH>@$>i1A@SEEOT4sv%kMkXzcb}`7SsH?LrujZL%m$tu@zZX+EF#Nj5wDiDo{*N+1G&gpj?C^%)d%aYp_kY*fk_^9X zzQI0ZN+tce5tkPHUnckOdi`r%vr(zRSfRWQBNZDJ6IJMruhoI7b-Tli? zXFGc{N(KKK15{LLe|i32>-aBEU&h}CsHqtKHV^#zL75i>|F(ph`qz;-|NDR}B`qwc zJoK*}KYBPKl;t$kG=F<0)9>|ZX=sXnA1lhHmJS;g&+ogGmFqU#@TQ{L{kPZu|Mi11 z-%;vnsV`VcarS@t9d-YuJeNHEeFy*Lr!7CaY+%{IvVmm-%LbMWEE`xhuxw!2z_Nj5 z1Iq@M4J;d2Hn419*}$@a|MCWYU8kp975;TQ$@tHA$>I0nwC(?VJ-OeDk=Q?9Pwn^O zeC>a}-qP={;+KAb@y}QP4`2Epc%9xs{MX+7x4KRrzf>^4kKFQg`v33}UVg!41Iq@M z4J;d2Ht-*AAiyt_dlfqwc~o7Kq~v7<#trQJWVBwuDuI%nH?|1TiICCH#+Y-ITy>N` zc2SCqE;<}xppqt|IVp1WcO^2~{3X_fN|}r@ZacqO)FPu#GQ(yJ+GJF}?+Y!xA#k?r zN*>5SPSaryb6Ef5Y9UF;EccF2LbhZib!Z?=7jmC9qobk=8R-|F)wY355SVl|@E{}8 z^i7^ibpf0x&lfj?tOEK z;xSd{dm#s8+~9~Ay$Bf%QSBlzK}M+bPEt#gk)rRo3kzgt$}tIgWimSB$j~7OsW4XH z%%x36l^=6V^&mGix7$h?l93@{6I~FbY=H0=f;ky&nRdK%8M4NGG0fAJjCMz9kiSAc zesw}J#)XVNpL0x7cO|2W!M55H9StWD91S#ffB;V;{6Q1gZS& z@O*bF85!}O^!o_87@cM{mq|u;O)cY8S!A>`kIzyZJSH1n8JB<*zsPrn;?Y35+}{=w zu_YT&JWg-XI!no0hYqe1Vu$TpvfBzn);sIn;^ZfzduyqAc_GcIezdWQkdgoS&u{r5 z*;aTm@JW*q{fCf6708=upW9TG;eGaA5fqa%8E_gc6?9qq~W zN1jlZjU=Z$FF3AOsy}M_lToYVQJ27N@V@DO41WQS53NhqKG;DgP#L@eeE5zw(Xm zOr7BM7}cZ+E@Z^G?R#s$Mi~FAh2rVJqm*JP&!9IPhv|U!+u$)h+NSw<5R7MX>@oH2 zuzk&=w@-vY`n;{kjDR|(cc?SjLq;ikJ6R*5U_Gb)H&pRt6yPnF61We>;nr0a9LZ!P zC!T%$738^UH>()%cv^>a`W+Om=0V%*`rt9P zb++9MQgAodr*!bRDL&;{DI|yets!R>nD-t&XF3e|IxdN=3Or7lwHzh)KQK-8Xy6g> zn7dbOc3U7B`F;DImK_XrRzsCny%Xj=6K~r)WHMr2^g?A&hn{vf-fabsmsoky=A+^L z-&vhy5DUkRcaGXO0miNUpuP5fGCFS3G3=iL^*PUE)(4r-bKsiXF{rO_=5cRG>(Y&n z4}!<3YtL#5A&XE6zbkl5<&>HBgxutjqP-bBnj4HC_JZ^u_G;V@9#`u}zd8WPBHqm= z0UrA~1ZgQgBTv|*NrFec^let+kgGns@>_yOYH|MzLrBT`yERw9<8#`WH`$N}4iG+O zgU6$u0_Qy;pI-hLtPdXbY2MSff>)EnyEN*-V|i4S)+~4x$k$m%nJ48&p4Z(4ugQ~F z((Pbei|!t*x(Z%1PLGquz@y7LE%`<8`iA%A*gSaLqIl-^7kY&hGU_=&d(IU+l8ihDvcYRl^v9JG;IT0* zOfM1AsCB1)FnClNf0?ceDP8kvjOsE|AW^Qk8~%h z1Uw;Cs~@r@g2z^OYFdg<$E~rQlyPpJUHO_a?mO=8WETOCNVD2=1tis(53QQuk$k^( zKmu|zSEgtzc$Bg{neGUga@Ai|3OuefEvCJ#4D-gMmVz^QjMx?aCQ%3G*;?^0lHf5S zs|eqKasQXklZ#PjkF_xJ5EpLFQ!JAr3j|J^2ryI}dr?XsFsub(2Cc*bFzZ@#v& z^*VTVrL8?xyAQ^_9F?v$jC(N`o-GXEky+c7K@Y~e)?P7%*`shi$;)p)4rx>0Q5g;% z$2>oW5g~nCWjO7?BganzPeVwy(|Ji7!K1y89jz;5`R5B=lyQD2Jh_K5?oDR)AE%s$ zE+4OSp`4fY>GW4p&O?2g4GVr$S1wWTNS}$Nm>FE7i9ayM3Dw~6cbSOl7rkGzdF_l zJnqu&6}5*9+aYzFavnN8oV1Q|UK)ACe4WxKF;%!bQu2XYMGd7-s+FG@qV!FXWzD6O zK1nQEdYsZX?dB-l&k7!I7pB~&g{*A2JS_ws83Vi=nIVr06F3OqG4QefWkEki?&6yu5ZW?>}l%j)6$e_kGl0jd3`NcLF>pM$|EQTrgyZmFL)bmY&noRE6RGgLJIypbkWuqmT(sbk zt}}1fd?w5XNpqJa!Q3wuq5v5O>AI_+v^i9(D$yJm-DK%h^9VL^_9+9DV ztSxQlqoiVU(jJOO<0B(}lyrXJe~Qw_J6Q{wQnInUgPzjIGmZLlXM| z;za^oZFZ+`GA9z8-)*XH&5ly99y#-a#R_jSnXVTyw8Do6*o60RuRt=peHNxd_|cp1 zo0IZ#M3E;=&-X6}gi-74y_WhwE4=NdP4SUZE8HRHH#h1oi8u$sikuZCP;-B{b3wf< zvUD9i@+CD-tBpzmXIdESbx>8I|Tf;+aB1p4|^v zKRGhbgPMe2o^D+!LrcPnG4eLA=}7o-=#hv7dJ^XPxwp9)vQfnKXcuI~jLpwmkS~5Z zsO+e>!n}k-6`=!GSc+>8_r_J$n6&YRiP3&*EKEp$<~v}GM|IEEZ4b1;j$5vNaBH!_ zO+l*7=S6JsV-3Nf2P(EWJ~jDdBbzNwdpz-Ys@4XxY+ZBwwVVxJ>ly78U1E(JLjvM0 zXkk0w)E;N&!twNs4os7>!qErZeXkCaFym!kn=Pdz+>pG=eC!wr2cFILnT;XgzxtF6 zE|?jMw#JF47yZMf*Bo%8?uYpL4zc9vRLBS zY_Q*^BloXdvB9~*A$I%5cu=rD(#~&UK|Q@I=}!8uM04*>zPT*RgI+N?M{tsCF{x4L zF6(MLto~5oXK|P~N__4xyJ#th8f|Wi7eALm3I#ts@2ZQVmDwdTan*L1s!VN1l8HUe zZ4hW@>ybxv0tEX-A8B+F^>W@Ms-n&ye#Yi3MU-{O&%fK<9)I|1!)z{Nj|<1D!tR=C zpwF8|;%@P%qy6H8k7Luc(7?pU*P<#~$Y0YmcTcVzZsI-0QCMq>6N@WzZc6H)AeqBN z+1)y*`KXb_p4U34GAB6_*XyFqtnt>apKWmZiw7qUeYD213-!Hy$8=D>VctW=DqXb2 zX78GEYi+bQP0Y)htb_LUNz$3vS>tv0s-myvSYfdfkvp{1G|@{c>6Zuew2;sjU)uBy zYG_Qg!mD0J9p$R++}3Sug?GNRWYioX;cpVdYdhb_qaF6DGo5UTD1e2;WXC0e2qH=4 zF>Mk^&ZPa7b~6b_ZOgxMp@f8aU6Sv#J?BB$7QNTk9pn9Dym+{J3`DWo;F%<`n_g}< zc(YTv_32g{+J^eKK_k{>teJWzDM)kcD~FW_YAC?_9i%Bf*{w+`636bJ(WVZ ze$xS~TM&I`PdVVxuE<<^Z3q0w!HsQ7*&fGj|8`c5-40LFZ>8fZvB8&3{WiJ3vc}er zN2v~fu)?{D5eE$8tuS|?#Fh^nR@g12^VUXW(c=A4{4_Nwlwr_qg!8KVv7UZ20s$q?QmK8 zrR#k4KM4X~TvSFF`x{*M-cd@FWI%Q;?mI`{4K^Hcl0RC%#U5u#$uZK99dPx?oX5V0 zoXEocoPE%EZ-b*&F{4`G3iQgOj-PSQ0Kw`0E6?yb2kfRr+hiN(h#%*D(-31?jhfrU zQ)n3&QT)@EyD6=*h?iDyOKZVObhkUhD=ga)+dih-wLZiVKX5TC?X!V-@lA#Dke&>> zYxN>T%T*I?`I>Y1ftVT^Txp{xGU9-Bo;$H~N;}|NKiUJW^mGv8QE7vcgIb7RaaENu zqaL!AmHp;^Nf!}CZ9j9ovcq>ZPUxq|*y5Vm)zx2}_2Bmc-$IXe=%J6|Syf^LT_m@E z`-qUL9=f1ges8my4gSDUbGC-l8efWS*Gin^etDIXGgQ%Fto!F-I)JZ`e7<=Y)sLOKwH!D>6?Gnqj$y2H%Go$WB;1#8(%|Rrc31pKaG9Y zz%9k5o*z2X;K9f9WVUOjLBF3%(%qIB|2(STgg@F~UOzdP3JE)GVtZ6N(GqT5D#*z|y4mjX5fmTo1w{qfLyj>NO^pcp? zmxmA@z>29m6XUDiNa3(h&N6!;N8Ck}9X@o!5qDcn>2fg4HO!ggd1VLir(Z#w7$ckSu#x_XGgJ;p?lN)Ns2-S@Vn z&leEhYd6tm9kilx&9>bII>@Z>_#swlO>{XT{R-_FO=M=VhrsA&h2Pv} zXNm2Gey&x>V@gW}^^biZ@w2L;DW!GXdk3Wu*78qm-7k%9FmcdttRvyrq^nGlxg@-A z=EAUU5kFcRKmGD3QQ(hy<)r?<7C+YtqRP4Sz?UWt~p??Y`@NTFC6d;)5`1b^Bu5N5BmduIR`we zeW&@=TRU9ouURF$)fV?g8QAE}SYu{dd+Q(xYuwQ87iO+wh3zZk8VY7fSkqLsXL|;m zw?Ca2tWJRQwNOnBU(lcR_+?*l*B9th8wysLJq)(OCDgmqTGSRBI#>_MmdVgy;r` znof?Cv&vz|_Z1EK_pEX6BUwdag$>^L)@zI9EqVN2PvPJiIaSQ6)t=B5u8aq=l$E57 zG%)vd@m{7(7=)H>uuql7o~H&wbHdQhwI>vGa_k zex?<65-$n8xJMi5tIK`4uCIk8T4gT0Hqu4Tcik3lmFXZFo3pmzVpdqe(uqC*>gE1f zn^%uLb&%+lO!rl4I%pzoh0?%zP2{p)b`5)lCMxbe`&}oNgw>+gQ-8OB^Y6&LCgZiL z$gT5;&t`8`G~fNA|4tqBDT-YK<|ESRXvP7rP67$@&sER`YLYOE7|#*&1OY_tF&KNl z>d$&?eA4Dm;wRx%k{;V12$FDSl4yUX2t0p<7-@$b3GYxO-!)PrVS|IT*^`DO>}Ar( z*>6n3jxQ62MIc|_wcSvwN5Zty)dL(-B-|-;-pQVqgvnkN`F+s8rmqPXX`g}qnUFf| z0DW{^G-O?V>Wj_%%xZ~J7dF*x!e5xRwP(r z`lxT*s}?M2R@`#$96eXStn%WaZ0^{)?%O;4$`(aD?Mw5g%fVf?TMAa z@f=y=jRiv3pDVy|*AYfM+8;jfHFTKpzOc-xm(3h&XE+zGX*b6+y=@CoUj-2Z-SCI2 zEE4GMo?7;YrCg_i4_jd$x{ZK*apR6{DcS|Gl1!Mo* zTxk@yTXy4>aw2XIS6K1-77;%c-(2>oKme%(pS^#_|Id1K5G%4AduxGN?F@b6Gc0g` zsyyutDGO{+M!V*tiaG9nUCk%sVusbrYWG}xWr9mWJY`Ny8siv75e>6<2zya~cTG5n zaMS!HF8ePC3$W9k$Zs{qdbPbzE#ysc-oy#|wV7r(BVO8W>zp|*9or+@{Llh_-IHYB zb=A% zD94#H@U7u#)W~bnoxujnFF99Uzhl7`g{w|aJIpssJjgwELBSjk9qHimn=!>?J(W?< z7y*o170B%aWS}Q2TQU=i5Jc zP%`fTw?vT)db!kJ#)C>plzb$8PLFyO$hVVaKt#^Sv6{oh~L*>Z*afnbOC8 z7-^%-sFh=h9;SGE$faF}gv~MEQ=eUusygU4Z6Tj*m@Z;J(z>{Ehc4p(ZhWJ=SQim4 zcumX*Ti|txEQeXlE%AZy@mm{>b2VSU9Oi+ zQyxUTuKDYoh*bi}s#>ULPxPPtV;%eTkxfqr!7`3%It|Df&k=;xyP6d4Tr|QeD=O$#=o#Qp%Ts&la`dr5wWF;_xB|-64xKdiYAaPaktAm zhn{&$oINqFDk@IIh|kZJZ1QJ4mMX?X-8^TG3Dfs3%bJ+ucJ9QpN{$r`_GHNx%Z4EW zWA3*Ex8mmwX{l#Uo~xcF6qk?cOI4fUO~Wte?_`-`5&I%dllqZ{$UXVItH}$5^ZBoj zn7x{4=rA=q{F9Fs(X$@ zIsKlU$3F5QrKxl9Nr8DiffszDzFjPOuoOfKzxDMS#{9g29Vj5xnONX>6jL2vApecx4TpfY8> zwb|kbQ$6LAFsU@g^_yomRzKH7HK(Udch+hn55|Hy15RD!6MiO1qEr`sr`Bh7QZ~iU z3XuJWK{L#l?Mh{CuZM!i*jN1E(L)+mVa9xqbWx1Yd)|I7UF51cX|bUP>ej$kIO4kn z*1X36pY&@XH$`RHzI_^qY|iS-RHKe;BulFkf-LcebbsckPnOuY!=uJk zQ5ls+5f5J*RYY~Cj;YtHNTE8}#V*B-k|;LT*Y=qJ5#RS;NIs-T#CHpasLLn#P>a>n zX-y%%Kl;b}-M32=2F>w?P8#vR40BAcvir1zs5#a;m76J~V}`eORx)_Qc;ousd(%kD z1dAl3Q_G9=k;${AEEwq{|)QV;Z;PzLg-XYSp`VH^~@_U$(+$I!tg*?uVqd z@22?M`p@yJea&&ssnqLKS{7KZePF>&(h_eBWpCPY$P#<|;~rnAN4J9wS8hoXv5GCr zC-Z-y9%n5!Cw>aHz)BI(vJz?*_>N|T!YjII0=v7ee_if8A!xr~aNQYdv^(@R&5#;1 zYSG`=l={LPj}K|&+*C5h;vMVFKISna^VT-gW5GN~=?snTOf)xopQM<&r$!KEd0)7( z=_u4MzI)C}JQ}WQvp11VgIy&2ZzWNWZ_g z4hp;9ZLA;x^PY*ImRO-KI?zh2?;p`Y>tY}7d&+H&Tetks(-*VA8R?Oq+>h%blhPOK zD&g}9_cgQW0Vf@_R^sVpbvJD!UDsu?rNIK9SGL}aBP_9qHc4s3flK#AY0UAug6SOBOJ?|> zhny4VDMg8?1P2j<*hq&7X9!!0iur^GMTL;zH_;G0`EG z_~cgxMLcAQPXwJ>(FObU%IQbh!WJTCUFTuS7DdE~jm4Mqj{G@4#;jc3yqlJY0}kBZ z&p&R7i)pz&YWXfU3T0<3@ayMu#t+*pF#pr2bJ$xEMe^vVNz3->MXp%cPqPeY$Xv@cdnIii7AqD{L4vo!F9 z9eSI$C2HZk>uy339$I*j%v|@G3&yQR)rl`#EOGh8;T`uLYva#Ua;B39v~hQJ+qz*J z9sD)g7409>#>b+aNhKF7anuX?2lswhVlQF!QsE(OJafl+Z$_*(j-oN{K3kxL-4yXb zJ|it$&Vuj6Xc4h%rH9eYokZ+adghTpSzOc<{ZW#JT zZ_an7J$GS#42slf4&cEZ-46<$iShjL`QYFBe#_Fo2VBa3$M;*(+JP$0-1OLZQFmZ5FX$gPPvpC@tb#f?bsd>_UD>m0l%IB*O32q`+xms3IDC@;c35L zcq#uK*TeI7{MrS|rv84p_`k~a@OZZ02jDN!@t6Ni(UPaX@8Dn7U&dtv%LbMWEE`xh zux#MJzJX-p4X(-urwGLF;cNKnH+y{USlq$!X$0i6>`_6gndon@TWg{_N4_tVn zJ}w&A1yx%Q-80aS60q;2JoBRs+tw*D^1o#S3&P!qfUwuCU9`XZ= zt2hcSY?~q|o~G6g0AI(cZyH_!U)#$%1>Ar~PvyFu1|~BoHoAhZD}MK${rove2%PO> zLclkbzSqA4UOkC8KZ7q?(vy_=)=9#y*>`VSfrr+8VR;H%jd$f~fUhfN4TkPpAYr}5s0{&{auT}sjC}3B1>y19OdPwl0WH5_;b$ZQcvMF2Aobdk5Tda!sHgaGPaC(lzk4Mm%G0 z1o(Q|yuJ54aPsxlOaZ__gXDr9@D*pRb)gGX`2_@Wopb+o}f~y|e#g1h9t!LysEx@(ziN zdkVggeA7j4;G(wV6D7cVMz(Ib3BD}xzSKtWb)q8c79eCaPet8MtWGBY-Ven;iuT-l6eA#t!AGiR#-{igEe&7NbG26r7i$kfP zwh4R%lOn#j0N?1+FUtm&TT{KS9el+&+i1ChFVa)XXJx=7p|c$Nz*2$Z*R#M^gNK3k zCh)a&m_UvPeo+;hZ3WzI5W>^~zSbwP#5I7gxpuk8>%dY`G?Q0=KPVP-e+FO7%35+y z!PoXfJzx8QN4jG+H3NT>{??%gzFb<(mHffi=I@JhtAJ}g!p}MaQ+>aeLJhuHcMs~G z1Ye@!Dy+jIQ-lvASG*&D9mo1U>;_*>L9=7w;A_7{g|t2JK(EB0G4OV0ceeB3>vWm? zP7d(pv7XB)2>614;M0!mNkRq94PqMjQV4v(84SM8zG*af0&Zoi*VO>Fu~@ss2z-sk z-V@sgzIxO))UyDOQ{BF~9yt2YMV@T%6*#sKBM-iWuj>Ry01pl|UgrYN?{;jpX-5624H4S`waKT16ZU#qDO zT=oE8YmTnKkAc(UA_H82XYS91tAj7sq7X4*@O4)GBA*;^>@MCAF5s^V-p&YovCOQO zItRY?UUe;12lnqPqdo?loEnrP3BIaDSza)LFT)O%+w8#eC64zde4ySRb#n~aPZHk8 zdbR6=FM@>`>HuED(itibykpCX+n(U7(q-lbHTarkd!6h6d{Sr4*%@lsj?bY{NAM-& z7VRSgzWkpDsFHwx3{)(L0YB2QFUkU6Pu_PZp8#Jv#-HVn1KUy!D;xzD6W(Y>@wNH< zQQm6swfAOUNjGrkb5^}N;H|-udzio%F^hAT4){_kl4aupmjCFMYz)jDn*;ymdy4Q} zlHg$gz77P;;11y7nhGZ;;AT_R3$)s&%;T=<2#p(#ehYxT<%^3))`g2@uhl_ zkm9xZwk`OIULi!iisG|nqJRu6#e3GVecvSERpx8MdhoUSO4{LRVB+J(fF9sL2`!ah@5*f4!$^t+!F19$6a|N=76Uk47Hhq zuNx`4U#t3jx#WN)g3=i_0}t(~;i1$Y zZJt7lGx*wia8=MGaEF_J!fxR6D@2Qh!58%bo%1u`EA>^JI|+DU7AvYi{S_9S`pN;m znzcTs27s?H(=K~kVB9AddwTK>c+x2lGY&qaRy72?0}~i+wKO zYYT6?m=gHfM)n-p1}w+F!CD{KTc$QK4}5(XpGP_1OJ6Cyt^|1X!F|aWfDbd*`pto_ zg6S_cJ>aXFXZ95H!70LFdd1c;;0MWqH+8_*nD9G?o#2c4n7|bhu-T=>ohiV(UVkZO zhWRh(CQ8o%UyajG$EASfe;8NZ1TK4({B0-rvZMYmyajv}JPzp00scBdTb2fFeZEdN z34C#i{u~$uU;E7LHr@oTTm3~x3g$z)z6}Av;48ke#LoeIC2O2{Qw03x>(lUqz&8xl zuPDKM=TRj`9}T{^TD~l70ggUoR#F7~gh_jE4)}80+jLhQd~w~2mb?vYa#y&`7C7^u zwNe!LN)sONB7(11is((!1U3G!MSo*dJuA1|Cqkk;(}DOs7+THjbuEYU@MG}x>URIlYrsYHWeFdFM>idc>H%LlPts()!57`7p6g$MUlymh zMFMY(_|SR|eDU$CjrmdPZ87sf74X$5*Wz&C7n?M@+QHX5r_6`z^waixB%LY!@`iEI3`)Non*Ovg9JpMl zx%)P7o~DmHrC+9B|22=&FQ-d?Q}PF1r`dLa2Y3VTL;s% M*CbmKDE;*R0`b(X*8l(j literal 34104 zcmeI5c{tVG*Z+@s%9JVdEF$xi;%qV`V;T$%h6qK%HFhxhs8!y`!+8 zoFD`L;^${@GsNhdW511%UzTa3_l)IZ*8|864iYz$c$9PW1_y(Wfv-&<_luF|SK%7H zV+FZ?j+cW213hx=`Ntx;%<$wLy~RIv2Rr^-5t!j%>xtj64|&&PRa#5xYz)Sd@Q{dL z3RlwTuNh+({x73Asnt=Pa7Q8Cf}k>PRS%R~L616PHQUSAv$6%*_q5xz1eJoxWd z8I0NP9`?9`M|*(5$Rc$+_VfR&<92}!0O$pbr$c3)s zE8;55zs?_1A%^J)c^nVA7hmV(;G~~`pFEz6lk+I~tngEfeQb>JPTKy9sl9)0LSK?7xEB4~yGULb_*Opf++N24 zi_Ud%4$APrmkHx$dkDyY>hX(;wz)}x#hka(9k(fg9jAT2?+#D^5!VwJotENi@5#HS{BmmG_&U8bmGX(;W$x}vS@BxndibU6FT)~GX~96Kb&VLT z<1sUzHcJmI;HWF&bzh6fb$LF$TJMwQ@5SnSzOt#?>terr-y|(>8}w1e*7(N%_1|0LEvY>$hkw0{FR5u|<(-({cw22e__jo3x!2P1;BTdQ(&q9xgOgSF zZVWrUxiU~?xL@n$LVXS!6#;7_m+iw<$wFK*9vwn9Fr&U>$Y{hmStr5#mN!Rrq5hiB zBCozy6$(yw$z^g?7lvx_j!fsOW;910n+G{;p!A~YcSXNdG5F20-c*%ByGt&ozY{!3oeW`U~V3x=~~H6wZY6Rl6x)>-xIANJ&O+CUDqN&S!BS5rpm8X@GP%pV^M`NQ5H0t&U%2ZW=j;=RmDUVZz zx0~wqp1o3mhEMnFdnl;Fym_aY?MW(d0_Ud&enTZ_Ij9%axLOukML4#B*;25;e$8vg zVNv+b`g(_2hybiMDm!adnnT`iQQVMJ2$w<`mK%)5FDwvNlbG9z#1WlYKx0H>02m9LI)fF;cC zIhIc*z+G-4+pQ&)pwjnOi^g9Z4{VBN-EUSB1Ys8|cWu#-1#43d_ZeIh2bCLk@6d@; zg$#d2bj3RjSfaQjC(cy`v|L$tRN{#Y=+{fj_AS!@3oUzDtR7{M@bh%7r}adb8)P55 zIaUY8Klu6f(FGmAPk9IU8ES(5smThtx(47zLl(<#nJ%cY6P>-fUI+d-v$t08jW!gy zR^b1`(ikY5SR5S0Wdy36V&(Rxn*uLOx0Ne3O+obuiIpGsYr(Do4*s=$YOqOfdNH4} zIoK+;W7m4h9Eca3n>X)?Ip}vA_UXK00k-|z89XgY8E!ootygL)1Ivy(X_;&@2g{5< zpSV_T0qzSQnx*1o2B1ZEoYV?)uy1eT^@q0OVTxVL%TQeb==N&Lf<#>t;JZ@3LC6%n z4yRBao7sjyTt;F~pQ;f!_B6G6=n@w^Pcw4dX$U{OhGcTO|J0FNJHW-;Cr< zfU6!Q^SVpQ!_*oPP`?;_#3|tWJRuiZXx3_b- zFKse}hjT5tHY*yzUH50zly5hLKi=mK!Ml@Tr_b7^A*cC1? zS#O@mK^BAi;@3xtkSl(V+*z^r%$O`ivrh-JSVkEZWgXJrnAiO{vU}?9v9uc-jyEc* z!go*XSB>AS37==ITPc5Zm|3@I$=4~OZ&;ttT64P`=LS_3GL!U8KeFQcJFh%Fp#z_* zoe|-DU;sznF1joGTo~|cRx7?q?qEG{nAW2|O&Yw|vE#*vWHYnFwsrA(K4Vx@b6BcN z#S~7S#(2ABp$eF@)Ae(H7dPNZUN9|CdLme@BG%rLFabRE{84!z#S}(;kxF{sXad2{ zq4YHe^+7<$tJOF6XoBl!?LTx2m;k>cKc`R4FaS>Lf<*Pojo?!=R`Wu`$#8Y!(x&b~ zGjROlu~ly;n}VoqK(V~p0&L{B+y2_i0`N6z&H3Ro5zg*7apv?o73k~WvhI$NB`Dyp zSGL$-3F4oPKi{Wm0irW|KX>+804X!^$Z}U%82c;LJIGlKs!uLi5#MbJO576lYF?Ou z6%QYW23<1(2U2dYEXy?pN)umrd87)!J&&)=2t6VgF&4e}Bl?3fUhlG{X|!XiI{#bDbRJxuyPAkD8sWx8Eeo z!JUU04<j6mcWIysQ`Y)B1!`--hVllJvU8JQ1bCUv{mKw}9(kMn$ud%xQS^i91FPzEN!9-R1a(6%?}XimuoIYbf<))|*S;SwVY)Y66D@pz1l3C~Xd9 zXjLB|XgyB@=Jgrv{oud@Ju3E6ooA$=L=9DEaX@VydD{hl_#4FoUg7{l?&CQ}|5S|4v@} z1Rzb7nx3-kXAQoK+9}^Z5o8t^OmUEw23yzVoQX3thFgR`Z^AyN-CKv`jU zWPgYzh_P8@{{6%N>rnQ%i(9XLW%2mUD^|&o0?#A`GlKM_|5lHeww@Q?y+Ry{gf|xF zf0BT=I|RK1H_E`9-OD>QLzJP?1--m(J#~09_>ii&kq%T`=9+Y^O%JNi;*1Q)FoX;- zyhIy*-cV>-nrR4QK?VOjO?CK^+3EXuk{WEBcz~)7m4jumQ_h9v zYeMJk>&req=wbf6Q(jV7%?*Xa`71VAtH7DT*2-aTW#Pe-uXh@LR|EXHf{X0iWWg7? zo`tJ+83L;rg7K2Adf<~w7hm&S3Aj+@XU1V}LHPa7Z3Qk5GjQ$mvl4DzQ=mDsL8VX4 z0zBC;{cUBcIS8%!%E`^*gkEhGaY?%0SldH;Q=MJRfyW{TGY)BUkbA~4H*>cMcXda=6!P>9`{^tn_M$iD9vdh|jQ zIJ=f6JrHbX$+on^P%Qq(5dHg($zGlgo;6VAu=&neh2S|zdy6F93Glarq| z-|DYoOgIYO-M?(jIM5o#H^^zzCawNl)|R8Xt=X(fE`hsR<>1|f@?aqwMw=LAWc{fs z`OVT^Uy&heMozrg<^s-Ya4BSZU8q2H;RAunWo8p>DzhRhsvk>K|EV9>h`Ql=1g574 zpQCyNd{euw;W|v+T1(eu+g3}u4sWhHN!R7V6aKg!fuO-4u1Da}zLTy)m6@05y5!b8 zNY~->gDG@f{ybBQ>k$kOPs8;H5`XTe>o8K2CkNN%4vsXs4mpi^>AGwVqU#aF7<=J* z1X90l(slUx>1bVc1s2hDIMcL+uFKQoPT+b3N+&yTJ%WRUQ|USkc(tFd%Z|WH*z`Q1 z>+ofGEnSy0-#OBC_$u%&U6)jV02iu9AQMr9>yi1b*`BV$JJ-7Dx;(Nij;_O3muAp) zX?@!o*CX@r>{?upOgJHvu0u&5#!Ot7r97YLIxKM6N7v=d{#aa(OdEkLT#roAqt0|4 zKI&4T>+;XfcT&nld_IDCVHNm%1Z|#E@c9T9D z=%v$hTi^W_uT2Z&HIX;*Feqk1D@BQ`n)YYr|)Qg zr02F<(de8$=e>rW+q-hQ@Vo}WLG-)^+#{FhIsN#-e0pv_a!{h@v|ssBdT#so@4)jK zG+t=O^BTN9XHU=Rw#TpOx&1)qJ3Xf(Vm0Ww9sW%K&ud`)%@5COFyxj=&uN}*S@hhl z@7Yh!>7k|1>AC&5?GT>V;LDY6Jg>pdndj*_T{NeNp4*)z-{?6lZZ$f$snErEUW2&C z(|BGpFY4@|=k#6GB6@DW&6r2e>EIit^xWR%WQ*rD^Gb?3p4ZHSrQGzKE-ac$&+R`y zUwOZ$SU-aJRvD}xfsMBh){kJ<_eiWCL2&^K>qjs__7&ETAnoF7tRF%8>sqWI!KYdu ztRKO^oO@V5g1UBYtRKPWFGg5Df@S+mv3>;Zhp%G&2;8i^v3>;h7N@X&1l1g0sPty^d3{iTTJgYuM#wL@E)@`#*W@=bWV2ReFnVEslfXT zIIne?-edFv!sxx`^`|@Z9%J*+f!=FEm$l)22Fy7|?=xU8WYT+#@zl}1rm$r^y~kt+ z_|tpM!PEVCp8>xzB=9~1B79otJtn^M2ff$mH}cx!J*L9ZhTdxeGwFQ>tou=g_ZhHC za&(X3x)n+9HLwfNd(4)0CcW2`C~M+<228zt6z?j6nMgcv!p%?=xW7V;#N6l#U#v_nN_?I(m=!YWRZQYmOUy#QO|ThoyL*0arFn zr}r3;^pxIff@ZCw_n0H@yXn0qwP+&VXP6CZits+eELk^;-eclZbm+b2&w3OJtik&= zko&L=@7Ex${}U7K*PuqU9q-rR(xeo;UxOy=O1xi#lEg&3UjqRv2HvlMR@)T3Ujtti zFT7uaSGk+;ehv0~Nyhs%kgjRK`!#5|5Q+C|Q2OXO-mihryllK*13q!lnP|TTyvH)} zeho~PTHyVf`RkQ7-mm{wkLUfnvEKvE+Cl91fCr34w9hlkZ5i$RtiAM<_Iaw#Y@~gk zk>a1&?*Zz|g|OcPn*3~OpJ&6@CffJ;X}OE`c}jyueV@VPo!IXICv&*4-vg#}rqVvo z)cGN_?_+=CJMHs??^#FtJ`qi{-vjve(S8pwx#3IuJd^f*r+ptjy>!~=*}|Jd`#!s4 zg|OcPB)syl-vfG&1=Bvy1FHkH@AK5agZ6p2Tt|H$-QDA`-vh*#GqB$S_&#LOK2O^Q zK>I#6xeQJIX@Tsk3jihD)t{i zWC}+c@*lxzg+%N>0wDz!_8&p1eh~H_!NHL8*nb2^9PeQN5is3%VgC{AQZ2>)Bgoe7 z!2Tn+b4DBck6^pPbnHKZmor1L{|NMXZe#xu%>B9(`;Wj$O&|M@%vN7L>_0N4MBkSq z|B-oN>YQ-oKmM&Ai!|3`zn&?;tBU=4a8heN_UTyw${+9qiM? z<@t-SZx5etE5JTIwBM`A4vC4?lj|h5dSFlEe}0 z*E1U}4`81jR(V!o-yYue)xkbJj5$mD_Hgzu!N`>905V z`(<{T`9OZAuFFj?up@+G$soz9;Ne}a)aLhYnlnXO9&f`aUj+uMfOV1u=jZ8XI z$h3!PQSsx!J(LH%l-|HV`SU@yi=xUsOzvmBAGdJ!FnOd?Qg~5*EFo?829z)JO1XXQ zb2qbnf`VJh{cfg7y6@g9lvfTH+pk5rLb@9B-IZ=;Tq zcjtJ1RYkmfbAE~QBVHN`9|NUO{^U?+>0o#-)54&8dnn?S?K$InDdKf?h0q}%lqVN$ zE8d6lP@S4Ha}ciyUkhc2MtYe35^H|hpnT)e?T_0~?knqh`aR+mFFUWLx2}iz&9Q%= z7UdG(_2TZLymn|#f)3(!?LbGyUc_sL&i-o*l)syyup$cOt(W&pDXb0apjPg15XJbQA4x;3JUm#v< zWxSR&Azn$$u&MV@{wvo^`8LXj;Y(9N#7kRY_RTKDD@~r`+mFq?OoMr+*1kmfYv~9# z7sTt<)C7%fh}Tp~r`-bO)^l%nQYa6zKRZJg@e;c5@Zl}Q%Vc%Eo(RfIcIR(9gmP}9 zjYm%)UP=Pi`Nt72&Grj66(v>$cDAx})&V6>ShpFrC!*WEtl$KWK6{39pvva5U zQNCDDFa0*+RhKyN$q~dWbgs$%ttjt&bzdX`Okg2i;oxoKIh2do_@|bmyfxR+vJ3H=VaLUDAMtY1bB}35 zxxVwp6*VZ|@6*$ygLs`PeC8U2c+GrUHeVLy#kr0dZYZA{_W8?q#a^cC(UN(a5U)ce zy0=OaK{>dnx-uNKxIO3Icg*ki#@w&9u(Q`k_qrX}n%|rRhgyZj; z5w8%=OR<*_FUzvBq=zVvH4Cgci*mcw5_826uaJOcrgn(er(NPL0w}-tEKb}CbocKBVJAOTvth^@a49>j~1s@?f1xQF>EUL@diN)NMZ z&2+BEC|}v>_`DJEsu+6NuZMWKS|!C@MtPkT=XM2@w{_^epNn|0IQKRTar7|1I;U)( zf$|pfMdQDs{HObFT@S=7Zd>kLO~k9?;*t<2l>cmvj8I0o->(>n-H6xSC0=iKAYMwB zhTmqPe6O#T@OG3>CcI(@uSCLg3E??Fcs(Y(0O1)%c*+r8)r8kA!n2m}d_j1vCcMUY z<`JHGgqI%Sr9^lZ5T0YaJ`i4e3D4Dp=NPXB!Yi8aJVSW?B)pmluaAUh9pU+k@ERe! z1_{qd!c&Iu+Cg}^5uOc%=M2)X(@4KQBmL|}`q`QE>k-ngM@Tmt(6 zvq(Rukba#@`n8bsvm@ze4bra)q+eO2pW{eBzaafuK>Bq*>1Scm&()+~?MT1okbdqY z{k)6xD?jPivw!sST+**`q+c(Sehwr3e4X@b4Cz;A($CqXpEr?tp-JWiC#GaxRFHWQMdpblnJ0#1Ua-i#5cp%BsFQi&P3DCOnI~pso-85r zqL0jrKr&CRk$F;2=EXiTFSd|*!dX8$PcD#ov6#$@1!SJYk$ECb<^>;_7lD7wlf`6S z2$Oj+j?5D`GEXKEz4Dmom42dUf{C8FLG;Q_qE{SC(KAX!&&U$J;!gC+Wuj-c5AOTv|9Adv5%F!u;{U+E&9WQa*YH#QWpDm(`Tsmsg#Ytp6VCko z{xME}y}{ovvokgUHUTyPHUTyPHi3Ttfy}bD@p8fxsIJn0wO=2B;$B` z<>JZQiYU&yc?8F!Uq?>;(x8A2*AE$nCIw>aAoHu8H?K8%oY8z;2^7=o`ksq9P$2$V@@EYcy9V{wE4fiX^Q*6v6^a(#YaR6G zqQ}oZvd;s>c~TzU%!L%V7gN<1grZWX^Q>7R6cBvmb}AZ0xd}aogCi-BAAlbVT&jb zeAqI-EQA97oJ(}WmQWzj?sa}eBn8ezaIJM&PJv%y+gmGQ(Br%VT-0JIkUT?B@j@a6 zE{+^e{DtDUmi&j;(kKw5ZvFlviZ*pWJf3AxAbd;y)6XbQct7v%j~oiPosH(;+)aTs z{#Gkg5J&#}iyrDIe*RQmj5)pzxU$I&#o+X%_LyT4b@&*Lt90*}%L`MW&_69w0mW5e zev)ER6i~gOHD3b7XoJxE;}j{7!+c*ajiRRuWJzmMV7&GH6eASR8hUb>Po_X#tiQ}E z6jw>LAE(SHa4ktytOvy~YjC2;ngS+1{MRdND4@J0sn*L8eGW5E-9G6=fg53J%bO8L zC#NjoGoGkkI6}|ad!yHj>T10&dL8qft>XPDV5@uW;LAV?4DOFF4-2Kh!V8NQ9Sx^I zThV$!l_&~$B}(fKub{w+Mf^@xh+}F=^%~hY3M9s5wFf3qVBHz(J+|v8uvqcoWA%*` zFfA;TJcQzo^`RaD+bPhoueitwMK9}brNW5gnabUMj4@);;+!^ zjyRq#u#@pZaf00uvn0gv+N|#p>rm7jpL|Raactq4u~!*IfppIXX~gm1Lw;`s6oq{( z3{??F?L48&MkorEUyyV|9GfgBR>Yy`9($+k4B~jgqqb8@#k=^o2Ba8xvSaoyzmQf&dTz*nE;<)?#_Mnbv^!mcC%XP${ zzT$p0F)1FsZo7wHOV?0fb5z_*tz-%qoqbT2nS%Nt@$!gE1A&$;L(Sr^sT7{Q) z^En^8Tk-YwOriRy8Vf=cFXREP8BOg~jGMt$A9^W&MN z6wn=1UtbqNf$RlZoIh8h=Lzu^@1rR2J!-k}+tujYK6>kCI^uY8q~wEUBI>8MPdj~* zD8TEkWuCI00vjWj)H`lK=kLQtp}5WHoQ>@-d57W+7vqP9nWzsMY`wxz>@X8roPjv5 zvKxp#g5qzxW*EWdkNP7!ZDL@T|_wg60SOg zV*ue=NH}gMT;&MIVZwC>;dp{@-9$K!xQ%kvBpjs)S3|sX(U_5E0%kM+F(>GNOSqkaFH^tlM>`+U;pfb{(n(&xWO-^-FduOodg zLi${a^!-?$_eqTQ{dm&n1*GrSkUlpheZP?O`2o`R`lQddk-irpeg2a4y*}ylJ*4lK zkv>-Zt zLvmy;`I0%rC>Whf?qm+VA#-VL4lN~fDU8ga4P-7&A#-RxnM)tZ96C(qQW=>;bI4rM zBy%X2%%y5FhZ4zLx=!Yh44F%5L?^Wn-PA>Nk{Z!XCPXKN65Vu~=%h5Fn{0_riX^(} z712p!x@k-&O(D9;hv=kKqMO8tPHH5&NtWoOGekEj5}kCN=q4YclLCltsw6tekLac) zL?;aq-BeC=QU=jYAw(zf5Z&ZTbP^!CX%^8*vP3rt5}mY`=%!YplLm-xIzn_(Akj^b z==cnx>v@Qdw!XN{HzT_K3eoXVMAwTE9WP0A{U)O0)rhX2M0C6+(e*Qlj@Ku;eu(J!7NYC5 zh>m|vbp1M_kx_bbS%g@mWOI{|o))ZlBS8@*m|ds|5bOB(weH z-+8bzHUTyPHUTyPHi3U4fi;q)POlPDne52sz^(&!AF$5>_Bp^l2iWHT`y61O1MG8v zeGdG)p9B9wf7x^NUm#-qqx|LM)xU4eY=8N89_)-wfK7l+fK7l+fK7l+fK7l+fK7l+ zfK7l+fK7l+fK7l+fKA|^OJE%O8=|?<+v8yHFeK4mTd^uUG6cWd|LWg4Cx2ZA0iq2U L3{P_V|84jW6S5pP diff --git a/integration-tests/equations/test_advection_diffusion.py b/integration-tests/equations/test_advection_diffusion.py index 6c91c1e35..20e7076b5 100644 --- a/integration-tests/equations/test_advection_diffusion.py +++ b/integration-tests/equations/test_advection_diffusion.py @@ -23,7 +23,7 @@ def run_advection_diffusion(tmpdir): # Equation diffusion_params = DiffusionParameters(kappa=0.75, mu=5) - V = domain.spaces("DG", "DG", 1) + V = domain.spaces("DG") Vu = VectorFunctionSpace(mesh, "CG", 1) equation = AdvectionDiffusionEquation(domain, V, "f", Vu=Vu, diff --git a/integration-tests/equations/test_dry_compressible.py b/integration-tests/equations/test_dry_compressible.py index 5a016daae..edb06e6f9 100644 --- a/integration-tests/equations/test_dry_compressible.py +++ b/integration-tests/equations/test_dry_compressible.py @@ -7,7 +7,7 @@ from gusto import * from gusto import thermodynamics as tde from firedrake import (SpatialCoordinate, PeriodicIntervalMesh, exp, - sqrt, ExtrudedMesh, norm) + sqrt, ExtrudedMesh, norm, as_vector) def run_dry_compressible(tmpdir): @@ -57,6 +57,7 @@ def run_dry_compressible(tmpdir): rho0 = stepper.fields("rho") theta0 = stepper.fields("theta") + u0 = stepper.fields("u") # Approximate hydrostatic balance x, z = SpatialCoordinate(mesh) @@ -66,6 +67,9 @@ def run_dry_compressible(tmpdir): theta0.interpolate(tde.theta(parameters, T, p)) rho0.interpolate(p / (R_d * T)) + # Add horizontal translation to ensure some transport happens + u0.project(as_vector([0.5, 0.0])) + stepper.set_reference_profiles([('rho', rho0), ('theta', theta0)]) # Add perturbation diff --git a/integration-tests/equations/test_incompressible.py b/integration-tests/equations/test_incompressible.py index 51c6ac9f4..9bdf9f358 100644 --- a/integration-tests/equations/test_incompressible.py +++ b/integration-tests/equations/test_incompressible.py @@ -6,7 +6,7 @@ from os.path import join, abspath, dirname from gusto import * from firedrake import (SpatialCoordinate, PeriodicIntervalMesh, exp, - sqrt, ExtrudedMesh, Function, norm) + sqrt, ExtrudedMesh, Function, norm, as_vector) def run_incompressible(tmpdir): @@ -53,6 +53,10 @@ def run_incompressible(tmpdir): p0 = stepper.fields("p") b0 = stepper.fields("b") + u0 = stepper.fields("u") + + # Add horizontal translation to ensure some transport happens + u0.project(as_vector([0.5, 0.0])) # z.grad(bref) = N**2 x, z = SpatialCoordinate(mesh) diff --git a/integration-tests/equations/test_moist_compressible.py b/integration-tests/equations/test_moist_compressible.py index c51852e5e..91d55d1d8 100644 --- a/integration-tests/equations/test_moist_compressible.py +++ b/integration-tests/equations/test_moist_compressible.py @@ -7,7 +7,7 @@ from gusto import * import gusto.thermodynamics as tde from firedrake import (SpatialCoordinate, PeriodicIntervalMesh, exp, - sqrt, ExtrudedMesh, norm) + sqrt, ExtrudedMesh, norm, as_vector) def run_moist_compressible(tmpdir): @@ -60,6 +60,10 @@ def run_moist_compressible(tmpdir): rho0 = stepper.fields("rho") theta0 = stepper.fields("theta") m_v0 = stepper.fields("vapour_mixing_ratio") + u0 = stepper.fields("u") + + # Add horizontal translation to ensure some transport happens + u0.project(as_vector([0.5, 0.0])) # Approximate hydrostatic balance x, z = SpatialCoordinate(mesh) From 9b4b2194692724bc5a874305c5a7289df68658f9 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Mon, 16 Jan 2023 22:06:34 +0000 Subject: [PATCH 11/12] few final tidyings --- .../compressible/dcmip_3_1_meanflow_quads.py | 7 +++---- gusto/equations.py | 2 +- gusto/function_spaces.py | 2 +- gusto/timeloop.py | 17 ----------------- 4 files changed, 5 insertions(+), 23 deletions(-) diff --git a/examples/compressible/dcmip_3_1_meanflow_quads.py b/examples/compressible/dcmip_3_1_meanflow_quads.py index 7b7caa20d..40cbd20fd 100644 --- a/examples/compressible/dcmip_3_1_meanflow_quads.py +++ b/examples/compressible/dcmip_3_1_meanflow_quads.py @@ -99,9 +99,9 @@ # Initial conditions # ---------------------------------------------------------------------------- # -u0 = stepper.fields.u -theta0 = stepper.fields.theta -rho0 = stepper.fields.rho +u0 = stepper.fields('u') +theta0 = stepper.fields('theta') +rho0 = stepper.fields('rho') # spaces Vu = domain.spaces("HDiv") @@ -151,7 +151,6 @@ theta0 += theta_b rho0.assign(rho_b) -stepper.initialise([('u', u0), ('rho', rho0), ('theta', theta0)]) stepper.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) # ---------------------------------------------------------------------------- # diff --git a/gusto/equations.py b/gusto/equations.py index 1dc1281a0..c4f540a11 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -954,7 +954,7 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, muexpr = conditional(z <= zc, 0.0, mubar*sin((pi/2.)*(z-zc)/(H-zc))**2) - self.mu = Function(W_DG).interpolate(muexpr) + self.mu = self.prescribed_fields("coriolis", W_DG).interpolate(muexpr) residual += name(subject(prognostic( self.mu*inner(w, domain.k)*inner(u, domain.k)*dx, "u"), self.X), "sponge") diff --git a/gusto/function_spaces.py b/gusto/function_spaces.py index e7ee2687d..0e3a82e70 100644 --- a/gusto/function_spaces.py +++ b/gusto/function_spaces.py @@ -312,7 +312,7 @@ def check_degree_args(name, mesh, degree, horizontal_degree, vertical_degree): if degree is None and horizontal_degree is None: raise ValueError(f'Either "degree" or "horizontal_degree" must be passed to {name}') if extruded_mesh and degree is None and vertical_degree is None: - raise ValueError(f'For extruded meshes, either degree or "vertical_degree" must be passed to {name}') + raise ValueError(f'For extruded meshes, either "degree" or "vertical_degree" must be passed to {name}') if degree is not None and horizontal_degree is not None: raise ValueError(f'Cannot pass both "degree" and "horizontal_degree" to {name}') if extruded_mesh and degree is not None and vertical_degree is not None: diff --git a/gusto/timeloop.py b/gusto/timeloop.py index fab64ae67..6de21b00b 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -103,7 +103,6 @@ def set_reference_profiles(self, reference_profiles): profile field expr is the :class:`ufl.Expr` whose value is used to set the reference field. """ - # TODO: come back and consider all aspects of this for field_name, profile in reference_profiles: if field_name+'_bar' in self.fields: # For reference profiles already added to state, allow @@ -128,22 +127,6 @@ def set_reference_profiles(self, reference_profiles): self.reference_profiles_initialised = True - # TODO: do we need this interface? If so, should we use it in all examples? - def initialise(self, initial_conditions): - """ - Initialise the state's fields. - - Args: - initial_conditions (list): an iterable of pairs: (field_name, expr), - where 'field_name' is the string giving the name of the - prognostic field and expr is the :class:`ufl.Expr` whose value - is used to set the initial field. - """ - for field_name, ic in initial_conditions: - f_init = getattr(self.fields, field_name) - f_init.assign(ic) - f_init.rename(field_name) - class Timestepper(BaseTimestepper): """ From c7ea2d6380c5e4001668a747e30246e6a57a301d Mon Sep 17 00:00:00 2001 From: tommbendall Date: Tue, 17 Jan 2023 10:40:53 +0000 Subject: [PATCH 12/12] Correct name of sponge function --- gusto/equations.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/equations.py b/gusto/equations.py index c4f540a11..da5ce89b5 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -954,7 +954,7 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, muexpr = conditional(z <= zc, 0.0, mubar*sin((pi/2.)*(z-zc)/(H-zc))**2) - self.mu = self.prescribed_fields("coriolis", W_DG).interpolate(muexpr) + self.mu = self.prescribed_fields("sponge", W_DG).interpolate(muexpr) residual += name(subject(prognostic( self.mu*inner(w, domain.k)*inner(u, domain.k)*dx, "u"), self.X), "sponge")