diff --git a/gusto/core/kernels.py b/gusto/core/kernels.py index 5a644914c..2f2116ceb 100644 --- a/gusto/core/kernels.py +++ b/gusto/core/kernels.py @@ -10,7 +10,7 @@ """ from firedrake import dx -from firedrake.parloops import par_loop, READ, WRITE, INC, MIN, MAX, op2 +from firedrake.parloops import par_loop, READ, WRITE, RW, MIN, MAX, op2 import numpy as np @@ -158,7 +158,7 @@ def apply(self, lamda, mX_field, mean_field): lives in the continuous target space. """ par_loop(self._kernel, dx, - {"lamda": (lamda, INC), + {"lamda": (lamda, RW), "mX_field": (mX_field, READ), "mean_field": (mean_field, READ)}) diff --git a/gusto/spatial_methods/augmentation.py b/gusto/spatial_methods/augmentation.py index afa0ce559..d53efc82f 100644 --- a/gusto/spatial_methods/augmentation.py +++ b/gusto/spatial_methods/augmentation.py @@ -11,13 +11,10 @@ transpose, nabla_grad, outer, dS, dS_h, dS_v, sign, jump, div, Constant, sqrt, cross, curl, FunctionSpace, assemble, DirichletBC ) -from firedrake.fml import ( - subject, all_terms, replace_subject, replace_test_function, - drop, Term, LabelledForm -) +from firedrake.fml import subject from gusto import ( time_derivative, transport, transporting_velocity, TransportEquationType, - logger, prognostic, mass_weighted + logger ) from gusto.spatial_methods.limiters import MeanLimiter from gusto.core.conservative_projection import ConservativeProjector @@ -257,9 +254,10 @@ def update(self, x_in_mixed): 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. + This augments a transport problem involving a k=1 mixing ratio, by adding + a mean mixing ratio field. This enables positivity to be + ensured after each conservative transport step by blending the k=1 and + mean fields. Args: domain (:class:`Domain`): The domain object. @@ -278,20 +276,18 @@ def __init__( # Store information about original equation set self.field_names = [] + self.orig_spaces = [] for i in np.arange(len(eqns.field_names)): self.field_names.append(eqns.field_names[i]) + self.orig_spaces.append(eqns.spaces[i]) - self.eqn_orig = eqns self.domain = domain - exist_spaces = eqns.spaces - self.idx_orig = len(exist_spaces) + self.idx_orig = len(self.orig_spaces) DG0 = FunctionSpace(domain.mesh, "DG", 0) DG1 = FunctionSpace(domain.mesh, "DG", 1) # Set up fields and names for each mixing ratio - self.mean_names = [] - self.mean_idxs = [] self.mX_idxs = [] mX_spaces = [] mean_spaces = [] @@ -299,12 +295,7 @@ def __init__( for i in range(self.mX_num): mX_name = mX_names[i] - self.mean_names.append('mean_'+mX_name) - self.field_names.append(self.mean_names[-1]) 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) @@ -313,13 +304,14 @@ def __init__( # Determine if this is a conservatively transported tracer. # If so, extract the corresponding density name, if not - # set this to None. + # raise an error as this limiter shouldn't be used. for tracer in eqns.active_tracers: if tracer.name == mX_name: if tracer.density_name is not None: self.rho_idxs.append(eqns.field_names.index(tracer.density_name)) else: - self.rho_idxs.append('None') + raise ValueError("Mean mixing ratio can't be used for non-conservatively " + + f"transported mixing ratio {mX_name}") # Define a limiter using the mean mixing ratios self.limiters = MeanLimiter(mX_spaces) @@ -328,162 +320,36 @@ def __init__( self.DG1_field = Function(DG1) self.rho_field = Function(DG1) self.DG0_field = Function(DG0) - self.compute_mean_mX = ConservativeProjector(self.rho_field, self.rho_field, self.DG1_field, self.DG0_field) - # Create the new mixed function space - self.fs = MixedFunctionSpace(exist_spaces) + # Construct a conservative, consistent, projection operator for computing + # the mean mixing ratios. + self.compute_mean_mX = ConservativeProjector(self.rho_field, self.rho_field, self.DG1_field, self.DG0_field, subtract_mean=True) - self.X = Function(self.fs) - self.tests = TestFunctions(self.fs) + # Set up the function space and residual as + # per the original equation set + self.residual = eqns.residual + self.X = eqns.X + self.tests = eqns.tests + self.fs = MixedFunctionSpace(self.orig_spaces) self.x_in = Function(self.fs) self.x_out = Function(self.fs) self.bcs = None + # Make mean fields + mean_fs = MixedFunctionSpace(mean_spaces) + self.mean_fields = Function(mean_fs) + def setup_residual(self, equation): """ - Create a new residual for the augmented equation set, - using the larger mixed function space that includes the mean - mixing ratios, and new residual terms for the mean mixing ratio - equations + Copy the residual to the augmentation Args: equation (:class:`PrognosticEquationSet`): The overarching equation set. Note, this does not include the mean mixing ratios. """ - # Copy the existing residual - new_residual = equation.residual - - # Replace test and trial functions of original residual with - # those from the new (larger) mixed function space. - # The indices of the original fields - # are the same in the new mixed space. - for idx in range(self.idx_orig): - new_residual = new_residual.label_map( - all_terms, - replace_subject(self.X, old_idx=idx, new_idx=idx) - ) - new_residual = new_residual.label_map( - all_terms, - replace_test_function(self.tests, old_idx=idx, new_idx=idx) - ) - - # Loop over each mean mixing ratio, - # copy the residual terms relating to the original mixing ratio, - # update the test and trial functions, then add to the new residual. - for i in range(self.mX_num): - mean_residual = new_residual.label_map( - lambda t: t.get(prognostic) == self.mX_names[i], - map_if_false=drop - ) - - # Replace all instances of original mixing ratios with - # the mean versions - for j in range(self.mX_num): - mean_residual = mean_residual.label_map( - all_terms, - replace_subject(self.X, old_idx=self.mX_idxs[j], new_idx=self.mean_idxs[j]) - ) - - 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]) - ) - - # Update the name to be that of the mean mixing ratio - mean_residual = mean_residual.label_map( - all_terms, - lambda t: prognostic.update_value(t, self.mean_names[i]) - ) - - # Append to the new residual - new_residual += mean_residual - - self.residual = subject(new_residual, self.X) - - # For any mass_weighted (conservative transport) terms, - # replace the subject and test functions for - # the mass-weighted form, and update the label to - # point to the new form. - for term in self.residual: - if term.has_label(mass_weighted): - field = term.get(prognostic) - mass_term = term.get(mass_weighted) - - # Extract the previous labels for the - # mass-weighted term - if term.has_label(transport): - old_mass_weighted_labels = mass_term.labels - - # Transport terms are Terms not LabelledForms, - # so this change this to use the label_map - mass_term = LabelledForm(mass_term) - else: - old_mass_weighted_labels = mass_term.terms[0].labels - - if field in self.mX_names: - list_idx = self.mX_names.index(field) - mX_idx = self.mX_idxs[list_idx] - rho_idx = self.rho_idxs[list_idx] - - # Replace mixing ratio - mass_term_new = mass_term.label_map( - all_terms, - replace_subject(self.X, old_idx=mX_idx, new_idx=mX_idx) - ) - - # Replace original density - mass_term_new = mass_term_new.label_map( - all_terms, - replace_subject(self.X, old_idx=rho_idx, new_idx=rho_idx) - ) - - # Replace test function - mass_term_new = mass_term_new.label_map( - all_terms, - replace_test_function(self.tests, old_idx=mX_idx, new_idx=mX_idx) - ) - - elif field in self.mean_names: - list_idx = self.mean_names.index(field) - mX_idx = self.mX_idxs[list_idx] - mean_idx = self.mean_idxs[list_idx] - rho_idx = self.rho_idxs[list_idx] - - # Replace mixing ratio - mass_term_new = mass_term.label_map( - all_terms, - replace_subject(self.X, old_idx=mX_idx, new_idx=mean_idx) - ) - - # Replace density - mass_term_new = mass_term_new.label_map( - all_terms, - replace_subject(self.X, old_idx=rho_idx, new_idx=rho_idx) - ) - - # Replace test function - mass_term_new = mass_term_new.label_map( - all_terms, - replace_test_function(self.tests, old_idx=mX_idx, new_idx=mean_idx) - ) - - # Create a new mass-weighted term, which has the correct labels - mass_term_new = Term(mass_term_new.form, old_mass_weighted_labels) - mass_term_new = subject(mass_term_new, self.X) - - # Make a new term, that links to the new mass-weighted term - new_term = Term(term.form, term.labels) - new_term = mass_weighted.update_value(new_term, mass_term_new) - - # Put this new term back in the residual: - self.residual = self.residual.label_map( - lambda t: t == term, - map_if_true=lambda t: new_term - ) - - self.residual = subject(self.residual, self.X) + self.residual = equation.residual def pre_apply(self, x_in): """ @@ -492,9 +358,7 @@ def pre_apply(self, x_in): Args: x_in (:class:`Function`): The input fields """ - - for idx in range(self.idx_orig): - self.x_in.subfunctions[idx].assign(x_in.subfunctions[idx]) + self.x_in.assign(x_in) def post_apply(self, x_out): """ @@ -503,13 +367,11 @@ def post_apply(self, x_out): 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]) + x_out.assign(self.x_out) def update(self, x_in_mixed): """ - Compute the mean mixing ratio field by conservative projection, + Compute the mean mixing ratio fields by conservative projection, where both the target and source density are in the higher-order space. @@ -517,28 +379,12 @@ def update(self, x_in_mixed): x_in_mixed (:class:`Function`): The mixed function, containing mean fields to update. """ - - # Update the density field for the mean mixing ratios - self.rho_field.assign(x_in_mixed.subfunctions[0]) - - # Update the mean mixing ratios: - for i in range(self.mX_num): - # Extract the reference density: - self.rho_field.assign(x_in_mixed.subfunctions[self.rho_idxs[i]]) - - # Compute the mean mixing ratio with conservative projection - self.DG1_field.assign(x_in_mixed.subfunctions[self.mX_idxs[i]]) - self.compute_mean_mX.project() - - # Clip any minuscule negative values in the mean mixing ratio - self.limiters._clip_means_kernel.apply(self.DG0_field, self.DG0_field) - - x_in_mixed.subfunctions[self.mean_idxs[i]].assign(self.DG0_field) + pass def limit(self, x_in_mixed): """ - Limit the mixing ratios using a blended limiter with - the mean mixing ratios + Limit k=1 mixing ratios using a limiter that blends the + k=1 field and its mean field. Args: x_in_mixed (:class:`Function`): The mixed function, containing @@ -549,12 +395,27 @@ def limit(self, x_in_mixed): mX_pre = [] means = [] + # Compute the new mean mixing ratio for i in range(self.mX_num): + self.rho_field.assign(x_in_mixed.subfunctions[self.rho_idxs[i]]) + + # Compute the mean mixing ratio with conservative projection + self.DG1_field.assign(x_in_mixed.subfunctions[self.mX_idxs[i]]) + self.compute_mean_mX.project() + + # Clip any minuscule negative values in the mean mixing ratio + # that arise from numerical error + self.limiters._clip_means_kernel.apply(self.DG0_field, self.DG0_field) + + self.mean_fields.subfunctions[i].assign(self.DG0_field) + mX_pre.append(x_in_mixed.subfunctions[self.mX_idxs[i]]) - means.append(x_in_mixed.subfunctions[self.mean_idxs[i]]) + means.append(self.mean_fields.subfunctions[i]) + # Limit all mixing ratios at the same time to have + # the same blending weights. self.limiters.apply(mX_pre, means) + # Update the mixing ratios with the limited version for i in range(self.mX_num): - self.limiters._clip_DG1_field.apply(mX_pre[i], mX_pre[i]) x_in_mixed.subfunctions[self.mX_idxs[i]].assign(mX_pre[i]) diff --git a/gusto/spatial_methods/limiters.py b/gusto/spatial_methods/limiters.py index 3d5e5170a..936b7426f 100644 --- a/gusto/spatial_methods/limiters.py +++ b/gusto/spatial_methods/limiters.py @@ -303,7 +303,6 @@ def __init__(self, spaces): DG1_equispaced = FunctionSpace(mesh, DG1_element) DG0 = FunctionSpace(mesh, 'DG', 0) - DG1 = FunctionSpace(mesh, 'DG', 1) self.lamda = Function(DG0) self.mX_field = Function(DG1_equispaced) @@ -312,10 +311,9 @@ def __init__(self, spaces): self._lamda_kernel = MeanMixingRatioWeights(DG1_equispaced) - # Also construct kernels to clip any very small negatives - # that arise from numerical error. These are used in the - # mean mixing ratio augmentation limit routine. - self._clip_DG1_field = ClipZero(DG1) + # Also construct a kernels to clip any very small negatives + # that arise from numerical error when computing the + # mean mixing ratio. self._clip_means_kernel = ClipZero(DG0) def apply(self, mX_fields, mean_fields): diff --git a/integration-tests/model/test_mean_mixing_ratio.py b/integration-tests/model/test_mean_mixing_ratio.py index 842913439..e9d7f7879 100644 --- a/integration-tests/model/test_mean_mixing_ratio.py +++ b/integration-tests/model/test_mean_mixing_ratio.py @@ -1,31 +1,27 @@ """ Tests the mean mixing ratio augmentation, which is used for non-negativity limiting in a conservative transport scheme. -A few timesteps are taken in the terminator toy test, with -non-negativity and mass conservation checked. +This uses a transport test on the sphere, with non-negativity +and mass conservation checked after a couple of timesteps. """ from gusto import * from firedrake import ( - exp, cos, sin, SpatialCoordinate, - pi, max_value, assemble, dx + cos, sin, SpatialCoordinate, pi, assemble, dx ) def setup_mean_mixing_ratio(dirname): - dt = 450. - tau = 12.*24.*60.*60. # time period of reversible wind, in s + # Parameters radius = 6371220. # radius of the sphere, in m - theta_cr = pi/9. # central latitude of first reaction rate, in rad - lamda_cr = -pi/3. # central longitude of first reaction rate, in rad - theta_c1 = 0. # central latitude of first chemical blob, in rad - theta_c2 = 0. # central latitude of second chemical blob, in rad - lamda_c1 = -pi/4. # central longitude of first chemical blob, in rad - lamda_c2 = pi/4. # central longitude of second chemical blob, in rad - rho_b = 1 # Base dry density - g_max = 0.5 # Maximum amplitude of Gaussian density perturbations - b0 = 5 # Controls the width of the chemical blobs + dt = 900. # Timestep size + tmax = 12*24*60*60. # Twelve days + theta_c1 = 0.0 # latitude of first cylinder, in rad + theta_c2 = 0.0 # latitude of second cylinder, in rad + lamda_c1 = -pi/4 # longitude of first cylinder, in rad + lamda_c2 = pi/4 # longitude of second cylinder, in rad + rho_b = 1. mesh = GeneralCubedSphereMesh(radius, 12, degree=2) xyz = SpatialCoordinate(mesh) @@ -36,25 +32,18 @@ def setup_mean_mixing_ratio(dirname): # get lat lon coordinates lamda, theta, _ = lonlatr_from_xyz(xyz[0], xyz[1], xyz[2]) - # Define co-located tracers of the dry density and the two species - rho_d = ActiveTracer(name='rho_d', space='DG', + tracer_space = 'DG' + + rho_d = ActiveTracer(name='rho_d', space=tracer_space, variable_type=TracerVariableType.density, transport_eqn=TransportEquationType.conservative) - X_tracer = ActiveTracer(name='X_tracer', space='DG', - variable_type=TracerVariableType.mixing_ratio, - transport_eqn=TransportEquationType.tracer_conservative, - density_name='rho_d') - - X2_tracer = ActiveTracer(name='X2_tracer', space='DG', - variable_type=TracerVariableType.mixing_ratio, - transport_eqn=TransportEquationType.tracer_conservative, - density_name='rho_d') + m_X = ActiveTracer(name='m_X', space=tracer_space, + variable_type=TracerVariableType.mixing_ratio, + transport_eqn=TransportEquationType.tracer_conservative, + density_name='rho_d') - # Define the mixing ratios first to test that the tracers will be - # automatically re-ordered such that the density field - # is indexed before the mixing ratio. - tracers = [X_tracer, X2_tracer, rho_d] + tracers = [rho_d, m_X] # Equation V = domain.spaces("HDiv") @@ -63,71 +52,54 @@ def setup_mean_mixing_ratio(dirname): output = OutputParameters(dirname=dirname) io = IO(domain, output) - k1 = max_value(0, sin(theta)*sin(theta_cr) + cos(theta)*cos(theta_cr)*cos(lamda-lamda_cr)) - k2 = 1 - - mixed_phys_limiter = MixedFSLimiter( - eqn, - {'rho_d': ZeroLimiter(domain.spaces('DG')), - 'X_tracer': ZeroLimiter(domain.spaces('DG')), - 'X2_tracer': ZeroLimiter(domain.spaces('DG'))} - ) - - # Using the analytical forcing from Appendix D of Lauritzen et. al. - physics_schemes = [(TerminatorToy(eqn, k1=k1, k2=k2, species1_name='X_tracer', - species2_name='X2_tracer', analytical_formulation=True), - ForwardEuler(domain, limiter=mixed_phys_limiter))] - - X, Y, Z = xyz - X1, Y1, Z1 = xyz_from_lonlatr(lamda_c1, theta_c1, radius) - X2, Y2, Z2 = xyz_from_lonlatr(lamda_c2, theta_c2, radius) - - g1 = g_max*exp(-(b0/(radius**2))*((X-X1)**2 + (Y-Y1)**2 + (Z-Z1)**2)) - g2 = g_max*exp(-(b0/(radius**2))*((X-X2)**2 + (Y-Y2)**2 + (Z-Z2)**2)) - - rho_expr = rho_b + g1 + g2 + augmentation = MeanMixingRatio(domain, eqn, ['m_X']) + transport_scheme = SSPRK3(domain, augmentation=augmentation, rk_formulation=RungeKuttaFormulation.predictor) - X_T_0 = 4e-6 - r = k1/(4*k2) - D_val = sqrt(r**2 + 2*X_T_0*r) + # Details of transport + transport_methods = [DGUpwind(eqn, 'rho_d'), DGUpwind(eqn, 'm_X')] - # Initial condition for each species - X_0 = D_val - r - X2_0 = 0.5*(X_T_0 - D_val + r) + time_varying_velocity = True + tau = tmax def u_t(t): - k = 10*radius/tau + k = 5.*radius/tau + u_background = 2*pi*radius/tau + lamda_prime = lamda - 2*pi*t/tau u_zonal = ( - k*(sin(lamda - 2*pi*t/tau)**2)*sin(2*theta)*cos(pi*t/tau) - + ((2*pi*radius)/tau)*cos(theta) + u_background*cos(theta) + - k*(sin(lamda_prime/2)**2)*sin(2*theta)*(cos(theta)**2)*cos(pi*t/tau) ) - u_merid = k*sin(2*(lamda - 2*pi*t/tau))*cos(theta)*cos(pi*t/tau) + u_merid = 0.5*k*sin(lamda_prime)*(cos(theta)**3)*cos(pi*t/tau) return xyz_vector_from_lonlatr(u_zonal, u_merid, Constant(0.0), xyz) - augmentation = MeanMixingRatio(domain, eqn, ['X_tracer', 'X2_tracer']) - transport_scheme = SSPRK3(domain, augmentation=augmentation, rk_formulation=RungeKuttaFormulation.predictor) - transport_method = [DGUpwind(eqn, 'rho_d'), DGUpwind(eqn, 'X_tracer'), DGUpwind(eqn, 'X2_tracer')] - - time_varying_velocity = True - stepper = SplitPrescribedTransport(eqn, transport_scheme, io, - time_varying_velocity, - spatial_methods=transport_method, - physics_schemes=physics_schemes) + stepper = PrescribedTransport( + eqn, transport_scheme, io, time_varying_velocity, transport_methods + ) stepper.setup_prescribed_expr(u_t) + rho_d_0 = rho_b + 0.5*cos(theta) + + # Slotted cylinders + m_X_0 = conditional(great_arc_angle(lamda, theta, lamda_c1, theta_c1) < 0.5, + conditional(abs(lamda - lamda_c1) < 1./12., + conditional(theta - theta_c1 < -5./24., 1.0, 0.0), + 1.0), + conditional(great_arc_angle(lamda, theta, lamda_c2, theta_c2) < 0.5, + conditional(abs(lamda - lamda_c2) < 1./12., + conditional(theta - theta_c2 > 5./24., 1.0, 0.0), + 1.0), + 0.0)) + # Initial conditions - stepper.fields("rho_d").interpolate(rho_expr) - stepper.fields("X_tracer").interpolate(X_0) - stepper.fields("X2_tracer").interpolate(X2_0) + stepper.fields("m_X").interpolate(m_X_0) + stepper.fields("rho_d").interpolate(rho_d_0) - X_sum = assemble(stepper.fields("rho_d")*stepper.fields("X_tracer")*dx) - X2_sum = assemble(stepper.fields("rho_d")*stepper.fields("X2_tracer")*dx) - XT_init = X_sum + 2*X2_sum + rho_X_0 = assemble(stepper.fields("rho_d")*stepper.fields("m_X")*dx) - return stepper, XT_init + return stepper, dt, rho_X_0 def test_mean_mixing_ratio(tmpdir): @@ -135,22 +107,19 @@ def test_mean_mixing_ratio(tmpdir): # Setup and run dirname = str(tmpdir) - stepper, XT_init = setup_mean_mixing_ratio(dirname) + stepper, dt, rho_X_0 = setup_mean_mixing_ratio(dirname) # Run for four timesteps - stepper.run(t=0, tmax=1800.) + stepper.run(t=0, tmax=4*dt) rho_d = stepper.fields("rho_d") - X_tracer = stepper.fields("X_tracer") - X2_tracer = stepper.fields("X2_tracer") + m_X = stepper.fields("m_X") - X_sum = assemble(rho_d*X_tracer*dx) - X2_sum = assemble(rho_d*X2_tracer*dx) - XT_sum = X_sum + 2*X2_sum - td_err = np.abs(XT_init - XT_sum)/XT_init + rho_X = assemble(rho_d*m_X*dx) + rho_X_err = np.abs(rho_X - rho_X_0)/rho_X - # Check that all the fields are non-negative - assert all(X_tracer.dat.data >= 0.0) and all(X2_tracer.dat.data >= 0.0), \ + # Check that the mixing ratio is non-negative throughout the domain + assert assemble((abs(m_X) - m_X) * dx) < 1e-14, \ "mean mixing ratio field has not ensured non-negativity" # Confirm mass conservation to a certain tolerance - assert td_err < 1e-14, "mean mixing ratio field has not ensured mass conservation" + assert rho_X_err < 1e-14, "mean mixing ratio field has not ensured mass conservation"