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 1c577c877..e271061f0 100644 --- a/gusto/spatial_methods/augmentation.py +++ b/gusto/spatial_methods/augmentation.py @@ -20,9 +20,9 @@ time_derivative, transport, transporting_velocity, TransportEquationType, logger, prognostic, mass_weighted ) -from gusto.spatial_methods.transport_methods import DGUpwind +from gusto.spatial_methods.limiters import MeanLimiter, MixedFSLimiter import copy - +import numpy as np class Augmentation(object, metaclass=ABCMeta): """ @@ -55,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): """ @@ -255,52 +262,64 @@ class MeanMixingRatio(Augmentation): Args: domain (:class:`Domain`): The domain object. eqns (:class:`PrognosticEquationSet`): The overarching equation set. - mixing_ratio (:class: list): A list of mixing ratios that - are to have augmented mean mixing ratio fields. - OR, keep as a single mixing ratio, but define - multiple augmentations? - START with just a single. + mX_names (:class: list): A list of mixing ratios that + require augmented mean mixing ratio fields. """ def __init__( - self, domain, eqns, mX_name + self, domain, eqns, mX_names ): self.name = 'mean_mixing_ratio' - - self.mX_name = mX_name - self.mean_name = 'mean_'+mX_name + self.mX_names = mX_names + self.mX_num = len(mX_names) + # Store information about original equation set self.eqn_orig = eqns - self.domain=domain + self.domain = domain exist_spaces = eqns.spaces self.idx_orig = len(exist_spaces) - # Change this to adjust for a list - mean_idx = self.idx_orig - self.mean_idx = mean_idx - - # Extract the mixing ratio in question: - mX_idx = eqns.field_names.index(mX_name) - self.mX_idx = mX_idx - # Define the mean mixing ratio on the DG0 space DG0 = FunctionSpace(domain.mesh, "DG", 0) - mX_space = eqns.spaces[mX_idx] - - # But, what if the mixing ratios are in - # different function spaces??? - - self.mX_in = Function(mX_space) - self.mean = Function(DG0) - self.compute_mean = Projector(self.mX_in, self.mean) - - # Set up the scheme for the mean mixing ratio - - mean_mX = Function(DG0, name=self.mean_name) - mean_space = DG0 - exist_spaces.append(mean_space) + # 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) @@ -315,14 +334,24 @@ def setup_residual(self, spatial_methods, equation): # Copy spatial method for the mixing ratio onto the # mean mixing ratio. - old_residual = equation.residual + orig_residual = equation.residual # Copy the mean mixing ratio residual terms: - mean_residual = old_residual.label_map( - lambda t: t.get(prognostic) == self.mX_name, - map_if_false=drop - ) - mean_residual = prognostic.update_value(mean_residual, self.mean_name) + 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') @@ -334,7 +363,7 @@ def setup_residual(self, spatial_methods, equation): #print('\n', idx) #print(field) - prog = split(self.X)[idx] + #prog = split(self.X)[idx] #print('\n residual term before change') #print(old_residual.label_map( @@ -342,11 +371,11 @@ def setup_residual(self, spatial_methods, equation): # map_if_false=drop #).form) - old_residual = old_residual.label_map( + 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) ) - old_residual = old_residual.label_map( + 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) ) @@ -361,19 +390,25 @@ def setup_residual(self, spatial_methods, equation): #print('\n mean mX residual after change') #print(mean_residual.form) - mean_residual = mean_residual.label_map( - all_terms, - replace_subject(self.X, old_idx=self.mX_idx, new_idx=self.mean_idx) - ) - mean_residual = mean_residual.label_map( - all_terms, - replace_test_function(self.tests, old_idx=self.mX_idx, new_idx=self.mean_idx) - ) + + # 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 = old_residual + mean_residual + residual = orig_residual + mean_residual self.residual = subject(residual, self.X) #Check these two forms @@ -383,6 +418,7 @@ def setup_residual(self, spatial_methods, equation): + def setup_transport_old(self, spatial_methods): mX_spatial_method = next(method for method in spatial_methods if method.variable == self.mX_name) @@ -427,6 +463,7 @@ def pre_apply(self, x_in): def post_apply(self, x_out): """ + Apply the limiters. Sets the output fields, i.e. not the mean mixing ratios Args: @@ -438,18 +475,28 @@ def post_apply(self, x_out): 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. """ - - # Compute the mean mixing ratio - # How do I do this? - self.mX_in.assign(x_in_mixed.subfunctions[self.mX_idx]) - - # Project this into the lowest order space: - self.compute_mean.project() - - self.x_in.subfunctions[self.mean_idx].assign(self.mean) - pass \ No newline at end of file + 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)