diff --git a/gusto/core/labels.py b/gusto/core/labels.py index 06ce4fa8b..5f928d6d2 100644 --- a/gusto/core/labels.py +++ b/gusto/core/labels.py @@ -95,6 +95,7 @@ def __call__(self, target, value=None): transporting_velocity = Label("transporting_velocity", validator=lambda value: type(value) in [Function, ufl.tensors.ListTensor, ufl.indexed.Indexed]) prognostic = Label("prognostic", validator=lambda value: type(value) == str) linearisation = Label("linearisation", validator=lambda value: type(value) in [LabelledForm, Term]) +mass_weighted = Label("mass_weighted", validator=lambda value: type(value) in [LabelledForm, Term]) ibp_label = Label("ibp", validator=lambda value: type(value) == IntegrateByParts) # labels for terms in the equations diff --git a/gusto/equations/prognostic_equations.py b/gusto/equations/prognostic_equations.py index 9f88bf31c..bf8dd7845 100644 --- a/gusto/equations/prognostic_equations.py +++ b/gusto/equations/prognostic_equations.py @@ -10,7 +10,7 @@ replace_subject, replace_trial_function ) from gusto.core import PrescribedFields -from gusto.core.labels import time_derivative, prognostic, linearisation +from gusto.core.labels import time_derivative, prognostic, linearisation, mass_weighted from gusto.equations.common_forms import ( advection_form, continuity_form, tracer_conservative_form ) @@ -142,9 +142,31 @@ def generate_mass_terms(self): :class:`LabelledForm`: a labelled form containing the mass terms. """ + if self.active_tracers is None: + tracer_names = [] + else: + tracer_names = [tracer.name for tracer in self.active_tracers] + for i, (test, field_name) in enumerate(zip(self.tests, self.field_names)): prog = split(self.X)[i] mass = subject(prognostic(inner(prog, test)*dx, field_name), self.X) + + # Check if the field is a conservatively transported tracer. If so, + # create a mass-weighted mass form and store this and the original + # mass form in a mass-weighted label + for j, tracer_name in enumerate(tracer_names): + if field_name == tracer_name: + if self.active_tracers[j].transport_eqn == TransportEquationType.tracer_conservative: + standard_mass_form = mass + + # The mass-weighted mass form is multiplied by the reference density + ref_density_idx = self.field_names.index(self.active_tracers[j].density_name) + ref_density = split(self.X)[ref_density_idx] + q = prog*ref_density + mass_weighted_form = time_derivative(subject(prognostic(inner(q, test)*dx, + field_name), self.X)) + + mass = mass_weighted(standard_mass_form, mass_weighted_form) if i == 0: mass_form = time_derivative(mass) else: @@ -304,43 +326,6 @@ def add_tracers_to_prognostics(self, domain, active_tracers): else: raise TypeError(f'Tracers must be ActiveTracer objects, not {type(tracer)}') - def generate_tracer_mass_terms(self, active_tracers): - """ - Adds the mass forms for the active tracers to the equation set. - - Args: - active_tracers (list): A list of :class:`ActiveTracer` objects that - encode the metadata for the active tracers. - - Returns: - :class:`LabelledForm`: a labelled form containing the mass - terms for the active tracers. This is the usual mass form - unless using tracer_conservative, where it is multiplied - by the reference density. - """ - - for i, tracer in enumerate(active_tracers): - idx = self.field_names.index(tracer.name) - tracer_prog = split(self.X)[idx] - tracer_test = self.tests[idx] - - if tracer.transport_eqn == TransportEquationType.tracer_conservative: - ref_density_idx = self.field_names.index(tracer.density_name) - ref_density = split(self.X)[ref_density_idx] - q = tracer_prog*ref_density - mass = subject(prognostic(inner(q, tracer_test)*dx, - self.field_names[idx]), self.X) - else: - mass = subject(prognostic(inner(tracer_prog, tracer_test)*dx, - self.field_names[idx]), self.X) - - if i == 0: - mass_form = time_derivative(mass) - else: - mass_form += time_derivative(mass) - - return mass_form - def generate_tracer_transport_terms(self, active_tracers): """ Adds the transport forms for the active tracers to the equation set. @@ -377,29 +362,35 @@ def generate_tracer_transport_terms(self, active_tracers): tracer_prog = split(self.X)[idx] tracer_test = self.tests[idx] if tracer.transport_eqn == TransportEquationType.advective: - tracer_adv = prognostic( + tracer_adv = subject(prognostic( advection_form(tracer_test, tracer_prog, u), - tracer.name) + tracer.name), self.X) elif tracer.transport_eqn == TransportEquationType.conservative: - tracer_adv = prognostic( + tracer_adv = subject(prognostic( continuity_form(tracer_test, tracer_prog, u), - tracer.name) + tracer.name), self.X) elif tracer.transport_eqn == TransportEquationType.tracer_conservative: + default_adv_form = subject(prognostic( + advection_form(tracer_test, tracer_prog, u), + tracer.name), self.X) + ref_density_idx = self.field_names.index(tracer.density_name) ref_density = split(self.X)[ref_density_idx] - tracer_adv = prognostic( - tracer_conservative_form(tracer_test, tracer_prog, - ref_density, u), tracer.name) + mass_weighted_tracer_adv = subject(prognostic( + tracer_conservative_form(tracer_test, tracer_prog, ref_density, u), + tracer.name), self.X) + # Store the conservative transport form in the mass_weighted label, + # but by default use an advective form. + tracer_adv = mass_weighted(default_adv_form, mass_weighted_tracer_adv) else: raise ValueError(f'Transport eqn {tracer.transport_eqn} not recognised') - if no_tracer_transported: # We arrive here for the first tracer to be transported - adv_form = subject(tracer_adv, self.X) + adv_form = tracer_adv no_tracer_transported = False else: - adv_form += subject(tracer_adv, self.X) + adv_form += tracer_adv return adv_form diff --git a/gusto/equations/transport_equations.py b/gusto/equations/transport_equations.py index 8be8b7f12..c6d6b48c6 100644 --- a/gusto/equations/transport_equations.py +++ b/gusto/equations/transport_equations.py @@ -122,7 +122,7 @@ def __init__(self, domain, active_tracers, Vu=None): # Add mass forms for the tracers, which will use # mass*density for any tracer_conservative terms - self.residual = self.generate_tracer_mass_terms(active_tracers) + self.residual = self.generate_mass_terms() # Add transport of tracers self.residual += self.generate_tracer_transport_terms(active_tracers) diff --git a/gusto/spatial_methods/transport_methods.py b/gusto/spatial_methods/transport_methods.py index c2b625882..dc2f9d6e2 100644 --- a/gusto/spatial_methods/transport_methods.py +++ b/gusto/spatial_methods/transport_methods.py @@ -8,7 +8,8 @@ ) from firedrake.fml import Term, keep, drop from gusto.core.configuration import IntegrateByParts, TransportEquationType -from gusto.core.labels import prognostic, transport, transporting_velocity, ibp_label +from gusto.core.labels import (prognostic, transport, transporting_velocity, ibp_label, + mass_weighted) from gusto.core.logging import logger from gusto.spatial_methods.spatial_methods import SpatialMethod @@ -34,7 +35,12 @@ def __init__(self, equation, variable): # Inherited init method extracts original term to be replaced super().__init__(equation, variable, transport) - self.transport_equation_type = self.original_form.terms[0].get(transport) + # If this is term has a mass_weighted label, then we need to + # use the tracer_conservative version of the transport method. + if self.original_form.terms[0].has_label(mass_weighted): + self.transport_equation_type = TransportEquationType.tracer_conservative + else: + self.transport_equation_type = self.original_form.terms[0].get(transport) if self.transport_equation_type == TransportEquationType.tracer_conservative: # Extract associated density of the variable @@ -77,6 +83,22 @@ def replace_form(self, equation): # Create new term new_term = Term(self.form.form, original_term.labels) + # Check if this is a conservative transport + if original_term.has_label(mass_weighted): + # Extract the original and discretised mass_weighted terms + original_mass_weighted_term = original_term.get(mass_weighted).terms[0] + new_mass_weighted = self.form.terms[0].get(mass_weighted) + + # Ensure the correct labels for the new mass weighted term + new_mass_weighted_term = Term(new_mass_weighted.form, original_mass_weighted_term.labels) + # Update the mass weighted transporting velocity + new_mass_weighted_transporting_velocity = new_mass_weighted.terms[0].get(transporting_velocity) + new_mass_weighted_term = transporting_velocity.update_value(new_mass_weighted_term, new_mass_weighted_transporting_velocity) + + # Add the discretised mass weighted transport term as the + # new mass weighted label. + new_term = mass_weighted.update_value(new_term, new_mass_weighted_term) + # Replace original term with new term equation.residual = equation.residual.label_map( lambda t: t.has_label(transport) and t.get(prognostic) == self.variable, @@ -199,14 +221,20 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, self.field, ibp=ibp) elif self.transport_equation_type == TransportEquationType.tracer_conservative: - form = upwind_tracer_conservative_form(self.domain, self.test, + mass_weighted_form = upwind_tracer_conservative_form(self.domain, self.test, + self.field, + self.conservative_density, + ibp=ibp) + advective_form = upwind_advection_form(self.domain, self.test, self.field, - self.conservative_density, ibp=ibp) + + # Store the conservative transport form in the mass_weighted label, + # but by default use an advective form. + form = mass_weighted(advective_form, mass_weighted_form) else: raise NotImplementedError('Upwind transport scheme has not been ' + 'implemented for this transport equation type') - self.form = form diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py index d1e27c425..c0438295c 100644 --- a/gusto/time_discretisation/time_discretisation.py +++ b/gusto/time_discretisation/time_discretisation.py @@ -17,7 +17,7 @@ from firedrake.utils import cached_property from gusto.core.configuration import EmbeddedDGOptions, RecoveryOptions -from gusto.core.labels import time_derivative, prognostic, physics_label +from gusto.core.labels import time_derivative, prognostic, physics_label, mass_weighted from gusto.core.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.time_discretisation.wrappers import * @@ -161,6 +161,28 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.evaluate_source.append(t.labels[physics_name]) self.physics_names.append(t.labels[physics_name]) + # Check if there are any mass-weighted terms: + if len(self.residual.label_map(lambda t: t.has_label(mass_weighted), map_if_false=drop)) > 0: + for field in equation.field_names: + + # Check if the mass term for this prognostic is mass-weighted + if len(self.residual.label_map(lambda t: t.get(prognostic) == field and t.has_label(time_derivative) and t.has_label(mass_weighted), map_if_false=drop)) == 1: + field_terms = self.residual.label_map(lambda t: t.get(prognostic) == field and not t.has_label(time_derivative), map_if_false=drop) + + # Check that the equation for this prognostic does not involve + # both mass-weighted and non-mass-weighted terms; if so, a split + # timestepper should be used instead. + if len(field_terms.label_map(lambda t: t.has_label(mass_weighted), map_if_false=drop)) > 0: + if len(field_terms.label_map(lambda t: not t.has_label(mass_weighted), map_if_false=drop)) > 0: + raise ValueError(f"Mass-weighted and non-mass-weighted terms are present in a timestepping equation for {field}. As these terms cannot be solved for simultaneously, a split timestepping method should be used instead.") + else: + # Replace the terms with a mass_weighted label with the + # mass_weighted form. It is important that the labels from + # this new form are used. + self.residual = self.residual.label_map( + lambda t: t.get(prognostic) == field and t.has_label(mass_weighted), + map_if_true=lambda t: t.get(mass_weighted)) + # -------------------------------------------------------------------- # # Set up Wrappers # -------------------------------------------------------------------- # diff --git a/integration-tests/model/test_conservative_transport_with_physics.py b/integration-tests/model/test_conservative_transport_with_physics.py new file mode 100644 index 000000000..a73ddb60f --- /dev/null +++ b/integration-tests/model/test_conservative_transport_with_physics.py @@ -0,0 +1,97 @@ +""" +This tests that conservative transport is correctly working with a physics +scheme. The conservative transport equations require the tracer to be multiplied +by the density through the 'mass_weighted' label, whilst the physics terms +do not. Here, we test that we correctly replace the transport terms with the +mass_weighted counterpart but leave the physics terms unchanged. +""" + +from gusto import * +from firedrake import (as_vector, PeriodicSquareMesh, SpatialCoordinate, + assemble, Constant, conditional, sin, pi) + + +def run_conservative_transport_with_physics(dirname): + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # set up mesh and domain + L = 10 + nx = 10 + mesh = PeriodicSquareMesh(nx, nx, L, quadrilateral=True) + dt = 0.1 + tmax = 5*dt + domain = Domain(mesh, dt, "RTCF", 1) + x, y = SpatialCoordinate(mesh) + + rho_d_space = 'DG' + ash_space = 'DG' + + ash = ActiveTracer(name='ash', space=ash_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) + + tracers = [ash, rho_d] + + eqn = CoupledTransportEquation(domain, active_tracers=tracers) + + # I/O + output = OutputParameters(dirname=dirname+"/conservative_transport_with_physics", + dumpfreq=1) + diagnostic_fields = [CourantNumber()] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + transport_method = [DGUpwind(eqn, "rho_d"), DGUpwind(eqn, "ash")] + + # Physics scheme --------------------------------------------------------- # + # Source is a constant, but constrained to a box in the bottom left corner + # of size 1-by-1, such that the total ash value + # should be equal to tmax = 0.5. + basic_expression = conditional(x < 1.0, conditional(y < 1.0, -Constant(1.0), Constant(0.0)), Constant(0.0)) + + physics_schemes = [(SourceSink(eqn, 'ash', basic_expression), SSPRK3(domain))] + + # Time stepper + stepper = SplitPrescribedTransport(eqn, SSPRK3(domain, increment_form=False), + io, transport_method, + physics_schemes=physics_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + rho0 = stepper.fields("rho_d") + ash0 = stepper.fields("ash") + + # Set a spatially varying density field and no ash + rho0.interpolate(1000.0*sin(pi*x/L)*sin(pi*y/L)+1000.0) + ash0.interpolate(Constant(0.0)) + + # Constant wind + u = stepper.fields("u") + u.project(as_vector([0.5, 0.5])) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=tmax) + return stepper + + +def test_conservative_transport_with_physics(tmpdir): + dirname = str(tmpdir) + stepper = run_conservative_transport_with_physics(dirname) + final_ash = stepper.fields("ash") + + final_total_ash = assemble(final_ash*dx) + + tol = 1e-3 + assert np.abs(final_total_ash - 0.5) < tol, \ + "Conservative transport did not correctly implement the Source physics"