diff --git a/gusto/core/kernels.py b/gusto/core/kernels.py index 6b78691dc..422b7456b 100644 --- a/gusto/core/kernels.py +++ b/gusto/core/kernels.py @@ -113,6 +113,54 @@ def apply(self, field, field_in): {"field": (field, WRITE), "field_in": (field_in, READ)}) +class MeanMixingRatioWeights(): + """ + Finds the lambda values for blending a mixing ratio and its + mean DG0 field in the MeanMixingRatioLimiter. + + First, each cell is looped over and the minimum value is computed + """ + + def __init__(self, V): + """ + Args: + V (:class:`FunctionSpace`): The space of the field to be clipped. + """ + # Using DG1-equispaced, with 4 DOFs per cell + shapes = {'nDOFs': V.finat_element.space_dimension(), + 'nDOFs_base': int(V.finat_element.space_dimension() / 4)} + domain = "{{[i]: 0 <= i < {nDOFs_base}}}".format(**shapes) + + instrs = (""" + min_value = 0.0 + for i + min_value = fmin(mX_field[i*4], mX_field[i*4+1]) + min_value = fmin(min_value, mX_field[i*4+2]) + min_value = fmin(min_value, mX_field[i*4+3]) + if min_value < 0.0 + lamda[i] = -min_value/(mean_field[i] - min_value) + else + lamda[i] = 0.0 + end + end + """) + + self._kernel = (domain, instrs) + + def apply(self, lamda, mX_field, mean_field): + """ + Performs the par loop. + + Args: + w (:class:`Function`): the field in which to store the weights. This + lives in the continuous target space. + """ + par_loop(self._kernel, dx, + {"lamda": (lamda, WRITE), + "mX_field": (mX_field, READ), + "mean_field": (mean_field, READ)}) + + class MinKernel(): """Finds the minimum DoF value of a field.""" diff --git a/gusto/spatial_methods/augmentation.py b/gusto/spatial_methods/augmentation.py index 42aedaa64..e271061f0 100644 --- a/gusto/spatial_methods/augmentation.py +++ b/gusto/spatial_methods/augmentation.py @@ -9,14 +9,20 @@ LinearVariationalProblem, LinearVariationalSolver, lhs, rhs, dot, ds_b, ds_v, ds_t, ds, FacetNormal, TestFunction, TrialFunction, transpose, nabla_grad, outer, dS, dS_h, dS_v, sign, jump, div, - Constant, sqrt, cross, curl, FunctionSpace, assemble, DirichletBC + Constant, sqrt, cross, curl, FunctionSpace, assemble, DirichletBC, + Projector +) +from firedrake.fml import ( + subject, all_terms, replace_subject, keep, replace_test_function, + replace_trial_function, drop ) -from firedrake.fml import subject from gusto import ( time_derivative, transport, transporting_velocity, TransportEquationType, - logger + logger, prognostic, mass_weighted ) - +from gusto.spatial_methods.limiters import MeanLimiter, MixedFSLimiter +import copy +import numpy as np class Augmentation(object, metaclass=ABCMeta): """ @@ -49,6 +55,13 @@ def update(self, x_in_mixed): pass + def limit(self, x_in_mixed): + """ + Apply any special limiting as part of the augmentation + """ + + pass + class VorticityTransport(Augmentation): """ @@ -73,6 +86,8 @@ def __init__( self, domain, eqns, transpose_commutator=True, supg=False ): + self.name = 'vorticity' + V_vel = domain.spaces('HDiv') V_vort = domain.spaces('H1') @@ -236,3 +251,252 @@ def update(self, x_in_mixed): logger.debug('Vorticity solve') self.solver.solve() self.x_in.subfunctions[1].assign(self.Z_in) + + +class MeanMixingRatio(Augmentation): + """ + This augments a transport problem involving a mixing ratio to + include a mean mixing ratio field. This enables posivity to be + ensured during conservative transport. + + Args: + domain (:class:`Domain`): The domain object. + eqns (:class:`PrognosticEquationSet`): The overarching equation set. + mX_names (:class: list): A list of mixing ratios that + require augmented mean mixing ratio fields. + """ + + def __init__( + self, domain, eqns, mX_names + ): + + self.name = 'mean_mixing_ratio' + self.mX_names = mX_names + self.mX_num = len(mX_names) + + # Store information about original equation set + self.eqn_orig = eqns + self.domain = domain + exist_spaces = eqns.spaces + self.idx_orig = len(exist_spaces) + + # Define the mean mixing ratio on the DG0 space + DG0 = FunctionSpace(domain.mesh, "DG", 0) + + # Set up fields and names for each mixing ratio + self.mean_names = [] + self.mean_idxs = [] + self.mX_idxs = [] + mX_spaces = [] + mean_spaces = [] + self.limiters = [] + #self.sublimiters = {} + + for i in range(self.mX_num): + mX_name = mX_names[i] + print(mX_names) + print(mX_name) + self.mean_names.append('mean_'+mX_name) + mean_spaces.append(DG0) + exist_spaces.append(DG0) + self.mean_idxs.append(self.idx_orig + i) + + # Extract the mixing ratio in question: + mX_idx = eqns.field_names.index(mX_name) + mX_spaces.append(eqns.spaces[mX_idx]) + self.mX_idxs.append(mX_idx) + + # Define a limiter + self.limiters.append(MeanLimiter(eqns.spaces[mX_idx])) + #self.sublimiters.update({mX_name: MeanLimiter(eqns.spaces[mX_idx])}) + + #self.limiter = MixedFSLimiter(self.eqn_orig, self.sublimiters) + #self.limiter = MixedFSLimiter(sublimiters) + + # Contruct projectors for computing the mean mixing ratios + self.mX_ins = [Function(mX_spaces[i]) for i in range(self.mX_num)] + self.mean_outs = [Function(mean_spaces[i]) for i in range(self.mX_num)] + self.compute_means = [Projector(self.mX_ins[i], self.mean_outs[i]) \ + for i in range(self.mX_num)] + + # Create the new mixed function space: + self.fs = MixedFunctionSpace(exist_spaces) + self.X = Function(self.fs) + self.tests = TestFunctions(self.fs) + + self.bcs = None + + self.x_in = Function(self.fs) + self.x_out = Function(self.fs) + + + def setup_residual(self, spatial_methods, equation): + # Copy spatial method for the mixing ratio onto the + # mean mixing ratio. + + orig_residual = equation.residual + + # Copy the mean mixing ratio residual terms: + for i in range(self.mX_num): + if i == 0: + mean_residual = orig_residual.label_map( + lambda t: t.get(prognostic) == self.mX_names[i], + map_if_false=drop + ) + mean_residual = prognostic.update_value(mean_residual, self.mean_names[i]) + else: + mean_residual_term = orig_residual.label_map( + lambda t: t.get(prognostic) == self.mX_names[i], + map_if_false=drop + ) + mean_residual_term = prognostic.update_value(mean_residual_term,\ + self.mean_names[i]) + mean_residual = mean_residual + mean_residual_term + + print('\n in setup_transport') + + # Replace the tests and trial functions for all terms + # of the fields in the original equation + for idx in range(self.idx_orig): + field = self.eqn_orig.field_names[idx] + # Seperate logic if mass-weighted or not? + #print('\n', idx) + #print(field) + + #prog = split(self.X)[idx] + + #print('\n residual term before change') + #print(old_residual.label_map( + # lambda t: t.get(prognostic) == field, + # map_if_false=drop + #).form) + + orig_residual = orig_residual.label_map( + lambda t: t.get(prognostic) == field, + map_if_true=replace_subject(self.X, old_idx=idx, new_idx = idx) + ) + orig_residual = orig_residual.label_map( + lambda t: t.get(prognostic) == field, + map_if_true=replace_test_function(self.tests, old_idx=idx, new_idx=idx) + ) + #print('\n residual term after change') + #print(old_residual.label_map( + # lambda t: t.get(prognostic) == field, + # map_if_false=drop + #).form) + + #print('\n now setting up mean mixing ratio residual terms') + + + #print('\n mean mX residual after change') + #print(mean_residual.form) + + # Update the subject and test functions for the + # mean mixing ratios + for i in range(self.mX_num): + mean_residual = mean_residual.label_map( + all_terms, + replace_subject(self.X, old_idx=self.mX_idxs[i], new_idx=self.mean_idxs[i]) + ) + mean_residual = mean_residual.label_map( + all_terms, + replace_test_function(self.tests, old_idx=self.mX_idxs[i], new_idx=self.mean_idxs[i]) + ) + + + #print('\n mean mX residual after change') + #print(mean_residual.form) + + # Form the new residual + residual = orig_residual + mean_residual + self.residual = subject(residual, self.X) + + #Check these two forms + #print('\n Original equation with residual of length, ', len(equation.residual)) + #print('\n Augmented equation with residual of length, ', len(self.residual)) + + + + + + def setup_transport_old(self, spatial_methods): + mX_spatial_method = next(method for method in spatial_methods if method.variable == self.mX_name) + + mean_spatial_method = copy.copy(mX_spatial_method) + mean_spatial_method.variable = self.mean_name + self.spatial_methods = copy.copy(spatial_methods) + self.spatial_methods.append(mean_spatial_method) + for method in self.spatial_methods: + print(method.variable) + method.equation.residual = self.residual + print(method.form.form) + print(len(method.equation.residual)) + + # Alternatively, redo all the spatial methods + # using the new mixed function space. + # So, want to make a new list of spatial methods + new_spatial_methods = [] + for method in self.spatial_methods: + # Determine the tye of transport method: + new_method = DGUpwind(self, method.variable) + new_spatial_methods.append(new_method) + + def pre_apply(self, x_in): + """ + Sets the original fields, i.e. not the mean mixing ratios + + Args: + x_in (:class:`Function`): The input fields + """ + + #for idx, field in enumerate(self.eqn.field_names): + # self.x_in.subfunctions[idx].assign(x_in.subfunctions[idx]) + for idx in range(self.idx_orig): + self.x_in.subfunctions[idx].assign(x_in.subfunctions[idx]) + + # Set the mean mixing ratio to be zero, just because + #DG0 = FunctionSpace(self.domain.mesh, "DG", 0) + #mean_mX = Function(DG0, name=self.mean_name) + + #self.x_in.subfunctions[self.mean_idx].assign(mean_mX) + + + def post_apply(self, x_out): + """ + Apply the limiters. + Sets the output fields, i.e. not the mean mixing ratios + + Args: + x_out (:class:`Function`): The output fields + """ + + for idx in range(self.idx_orig): + x_out.subfunctions[idx].assign(self.x_out.subfunctions[idx]) + + def update(self, x_in_mixed): + """ + Compute the mean mixing ratio field by projecting the mixing + ratio from DG1 into DG0. + + SHouldn't this be a conservative projection??!! + This requires density fields... + + Args: + x_in_mixed (:class:`Function`): The mixed function to update. + """ + print('in update') + for i in range(self.mX_num): + self.mX_ins[i].assign(x_in_mixed.subfunctions[self.mX_idxs[i]]) + self.compute_means[i].project() + self.x_in.subfunctions[self.mean_idxs[i]].assign(self.mean_outs[i]) + + def limit(self, x_in_mixed): + # Ensure non-negativity by applying the blended limiter + for i in range(self.mX_num): + print('limiting within the augmentation') + mX_field = x_in_mixed.subfunctions[self.mX_idxs[i]] + mean_field = x_in_mixed.subfunctions[self.mean_idxs[i]] + print(np.min(x_in_mixed.subfunctions[self.mX_idxs[i]].dat.data)) + self.limiters[i].apply(mX_field, mean_field) + print(np.min(x_in_mixed.subfunctions[self.mX_idxs[i]].dat.data)) + #x_in_mixed.subfunctions[self.mX_idxs[i]].assign(mX_field) diff --git a/gusto/spatial_methods/limiters.py b/gusto/spatial_methods/limiters.py index 8a782a054..c04311540 100644 --- a/gusto/spatial_methods/limiters.py +++ b/gusto/spatial_methods/limiters.py @@ -6,14 +6,14 @@ """ from firedrake import (BrokenElement, Function, FunctionSpace, interval, - FiniteElement, TensorProductElement) + FiniteElement, TensorProductElement, Constant) from firedrake.slope_limiter.vertex_based_limiter import VertexBasedLimiter -from gusto.core.kernels import LimitMidpoints, ClipZero +from gusto.core.kernels import LimitMidpoints, ClipZero, MeanMixingRatioWeights import numpy as np __all__ = ["DG1Limiter", "ThetaLimiter", "NoLimiter", "ZeroLimiter", - "MixedFSLimiter"] + "MixedFSLimiter", "MeanLimiter"] class DG1Limiter(object): @@ -261,3 +261,75 @@ def apply(self, fields): for field, sublimiter in self.sublimiters.items(): field = fields.subfunctions[self.field_idxs[field]] sublimiter.apply(field) + +class MeanLimiter(object): + """ + A limiter for a mixing ratio that ensures non-negativity by blending + the mixing ratio field with that of a mean mixing ratio. The blending + factor is given by the DG0 function lamda. + """ + + def __init__(self, space): + """ + Args: + space: The mixed function space for the equation set + mX_field: The mixing ratio field + mean_field: The mean mixing ratio field + Raises: + ValueError: If the space is not appropriate for the limiter. + """ + + self.space = space + mesh = space.mesh() + + # check that space is DG1 + degree = space.ufl_element().degree() + if (space.ufl_element().sobolev_space.name != 'L2' + or ((type(degree) is tuple and np.any([deg != 1 for deg in degree])) + and degree != 1)): + raise NotImplementedError('MeanLimiter only implemented for mixing' + + 'ratios in the DG1 space') + + # Create equispaced DG1 space needed for limiting + if space.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_element = TensorProductElement(DG1_hori_elt, DG1_vert_elt) + else: + cell = mesh.ufl_cell().cellname() + DG1_element = FiniteElement("DG", cell, 1, variant="equispaced") + + DG1_equispaced = FunctionSpace(mesh, DG1_element) + print(DG1_equispaced.finat_element.space_dimension()) + + DG0 = FunctionSpace(mesh, 'DG', 0) + + self.lamda = Function(DG1_equispaced) + self.mX_field = Function(DG1_equispaced) + self.mean_field = Function(DG0) + + self._kernel = MeanMixingRatioWeights(self.space) + + def apply(self, mX_field, mean_field): + """ + The application of the limiter to the field. + + Args: + field (:class:`Function`): the field to apply the limiter to. + + Raises: + AssertionError: If the field is not in the correct space. + """ + + #self.field_old.interpolate(mX_field) + #self.mean_field.interpolate(mean_field) + + # Compute the weights, lamda: + self._kernel.apply(self.lamda, mX_field, mean_field) + + print('applying limiter') + #print(self.lamda.dat.data) + + # Compute the blended field + mX_field.interpolate((Constant(1.0) - self.lamda)*mX_field + self.lamda*mean_field) diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py index 16732bd50..bbd37e565 100644 --- a/gusto/time_discretisation/explicit_runge_kutta.py +++ b/gusto/time_discretisation/explicit_runge_kutta.py @@ -477,6 +477,8 @@ def apply_cycle(self, x_out, x_in): for i in range(self.nStages): self.solve_stage(x_in, i) + if self.augmentation is not None: + self.augmentation.limit(self.x1) x_out.assign(self.x1) diff --git a/gusto/timestepping/split_timestepper.py b/gusto/timestepping/split_timestepper.py index 2006c73ac..4738688d1 100644 --- a/gusto/timestepping/split_timestepper.py +++ b/gusto/timestepping/split_timestepper.py @@ -4,7 +4,7 @@ from firedrake.fml import Label, drop from pyop2.profiling import timed_stage from gusto.core import TimeLevelFields, StateFields -from gusto.core.labels import time_derivative, physics_label +from gusto.core.labels import time_derivative, physics_label, dynamics_label from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation from gusto.timestepping.timestepper import BaseTimestepper, Timestepper from numpy import ones @@ -296,7 +296,25 @@ def transporting_velocity(self): return self.fields('u') def setup_scheme(self): + print('Setting up base equation') self.setup_equation(self.equation) + + # If there is an augmentation, set up the residual now + if self.scheme.augmentation is not None: + if self.scheme.augmentation.name == 'mean_mixing_ratio': + self.scheme.augmentation.setup_residual(self.spatial_methods, self.equation) + print('Setting up augmented equation') + # Go through and label all non-physics terms with a "dynamics" label + dynamics = Label('dynamics') + self.scheme.augmentation.residual = self.scheme.augmentation.residual.label_map( + lambda t: not any(t.has_label(time_derivative, physics_label)), + map_if_true=lambda t: dynamics(t) + ) + print(len(self.scheme.augmentation.residual.label_map( + lambda t: t.has_label(dynamics), + map_if_false=drop + ))) + # Go through and label all non-physics terms with a "dynamics" label dynamics = Label('dynamics') self.equation.label_terms(lambda t: not any(t.has_label(time_derivative, physics_label)), dynamics) @@ -306,6 +324,7 @@ def setup_scheme(self): if self.io.output.log_courant: self.scheme.courant_max = self.io.courant_max + def setup_prescribed_expr(self, expr_func): """ Sets up the prescribed transporting velocity, through a python function diff --git a/integration-tests/model/test_mean_mixing_ratio.py b/integration-tests/model/test_mean_mixing_ratio.py new file mode 100644 index 000000000..35bdf2c75 --- /dev/null +++ b/integration-tests/model/test_mean_mixing_ratio.py @@ -0,0 +1,175 @@ +""" +Tests the mean mixing ratio augmentation. Here a mixing ratio is transported +conservatively along with the dry density. The mean mixing ratio augmentation +should ensure that the mixing ratio remains non-negative. To test this, +a physics scheme of a sink is applied, which without the mean mixing ratio +will lead to negative values. Hence, a check is performed that the +mixing ratio has no negative values after a couple of time steps. +""" + +from gusto import * +from firedrake import ( + PeriodicIntervalMesh, ExtrudedMesh, exp, cos, sin, SpatialCoordinate, + assemble, dx, FunctionSpace, pi, min_value, as_vector, BrokenElement, + errornorm +) +import pytest + + +def setup_mean_mixing_ratio(dirname, pair_of_spaces): + + # Domain + Lx = 2000. + Hz = 2000. + + # Time parameters + dt = 2. + tmax = 2000. + + nlayers = 10. # horizontal layers + columns = 10. # number of columns + + # Define the spaces for the tracers + if pair_of_spaces == 'same': + rho_d_space = 'DG' + m_X_space = 'DG' + else: + rho_d_space = 'DG' + m_X_space = 'theta' + + period_mesh = PeriodicIntervalMesh(columns, Lx) + mesh = ExtrudedMesh(period_mesh, layers=nlayers, layer_height=Hz/nlayers) + domain = Domain(mesh, dt, "CG", 1) + x, z = SpatialCoordinate(mesh) + + V_rho = domain.spaces(rho_d_space) + V_m_X = domain.spaces(m_X_space) + + m_X = ActiveTracer(name='m_X', space=m_X_space, + variable_type=TracerVariableType.mixing_ratio, + transport_eqn=TransportEquationType.tracer_conservative, + density_name='rho_d') + + rho_d = ActiveTracer(name='rho_d', space=rho_d_space, + variable_type=TracerVariableType.density, + transport_eqn=TransportEquationType.conservative) + + # Define m_X first to test that the tracers will be + # automatically re-ordered such that the density field + # is indexed before the mixing ratio. + tracers = [m_X, rho_d] + + # Equation + V = domain.spaces("HDiv") + eqn = CoupledTransportEquation(domain, active_tracers=tracers, Vu=V) + + # IO + output = OutputParameters(dirname=dirname) + io = IO(domain, output) + + + if pair_of_spaces == 'diff': + Vt_brok = FunctionSpace(mesh, BrokenElement(V_m_X.ufl_element())) + suboptions = { + 'rho_d': EmbeddedDGOptions(embedding_space=Vt_brok), + 'm_X': ConservativeEmbeddedDGOptions( + rho_name='rho_d', + orig_rho_space=V_rho + ) + } + else: + suboptions = {} + + opts = MixedFSOptions(suboptions=suboptions) + + transport_scheme = SSPRK3( + domain, options=opts, rk_formulation=RungeKuttaFormulation.predictor + ) + transport_methods = [DGUpwind(eqn, "m_X"), DGUpwind(eqn, "rho_d")] + + #physics_schemes = [(SourceSink(eqn, 'm_X', -Constant(0.1)), SSPRK3(domain))] + + # Timestepper + time_varying = True + #stepper = SplitPrescribedTransport( + # eqn, transport_scheme, + # io, time_varying, transport_methods, + # physics_schemes=physics_schemes + #) + stepper = PrescribedTransport( + eqn, transport_scheme, io, time_varying, transport_methods + ) + + # Initial Conditions + # Specify locations of the two Gaussians + xc1 = 5.*Lx/8. + zc1 = Hz/2. + + xc2 = 3.*Lx/8. + zc2 = Hz/2. + + def l2_dist(xc, zc): + return min_value(abs(x-xc), Lx-abs(x-xc))**2 + (z-zc)**2 + + lc = 2.*Lx/25. + m0 = 0.02 + + # Set the initial state with two Gaussians for the density + # and a linear variation in mixing ratio. + + f0 = 0.5 + rho_b = 0.5 + + g1 = f0*exp(-l2_dist(xc1, zc1)/(lc**2)) + g2 = f0*exp(-l2_dist(xc2, zc2)/(lc**2)) + + rho_d_0 = rho_b + g1 + g2 + + # Constant mass field, starting with no mixing + # ratio at z=0 and m=0.5 at the model top + m_X_0 = m0 + 0.5*z/Hz + + # Set up the divergent, time-varying, velocity field + U = Lx/tmax + W = U/10. + + def u_t(t): + xd = x - U*t + u = U - (W*pi*Lx/Hz)*cos(pi*t/tmax)*cos(2*pi*xd/Lx)*cos(pi*z/Hz) + w = 2*pi*W*cos(pi*t/tmax)*sin(2*pi*xd/Lx)*sin(pi*z/Hz) + + u_expr = as_vector((u, w)) + + return u_expr + + stepper.setup_prescribed_expr(u_t) + + stepper.fields("m_X").interpolate(m_X_0) + stepper.fields("rho_d").interpolate(rho_d_0) + stepper.fields("u").project(u_t(0)) + + m_X_init = Function(V_m_X) + rho_d_init = Function(V_rho) + + m_X_init.assign(stepper.fields("m_X")) + rho_d_init.assign(stepper.fields("rho_d")) + + return stepper, m_X_init, rho_d_init + + +@pytest.mark.parametrize("pair_of_spaces", ["same", "diff"]) +def test_mean_mixing_ratio(tmpdir, pair_of_spaces): + + # Setup and run + dirname = str(tmpdir) + + stepper, m_X_0, rho_d_0 = \ + setup_mean_mixing_ratio(dirname, pair_of_spaces) + + # Run for two timesteps + stepper.run(t=0, tmax=4) + m_X = stepper.fields("m_X") + rho_d = stepper.fields("rho_d") + + # Check that the mixing ratio has no negative values + assert all(m_X_0 >= 0.0), "mean mixing ratio field has not ensured non-negativity" \ No newline at end of file