From 052a0438d7fb7589630c82b5fd104b28d3d2c968 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sat, 21 Sep 2024 12:53:43 +0100 Subject: [PATCH 1/3] put in structure and tests for advective-then-flux formulation --- gusto/spatial_methods/transport_methods.py | 75 +++++-- .../explicit_runge_kutta.py | 190 ++++++++++++------ .../time_discretisation.py | 2 +- .../equations/test_coupled_transport.py | 9 +- ...est_conservative_transport_with_physics.py | 6 +- .../model/test_time_discretisation.py | 15 +- .../transport/test_dg_transport.py | 13 +- 7 files changed, 216 insertions(+), 94 deletions(-) diff --git a/gusto/spatial_methods/transport_methods.py b/gusto/spatial_methods/transport_methods.py index dc2f9d6e..29fa1d4b 100644 --- a/gusto/spatial_methods/transport_methods.py +++ b/gusto/spatial_methods/transport_methods.py @@ -151,7 +151,8 @@ class DGUpwind(TransportMethod): transported variable at facets. """ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, - vector_manifold_correction=False, outflow=False): + vector_manifold_correction=False, outflow=False, + advective_then_flux=False): """ Args: equation (:class:`PrognosticEquation`): the equation, which includes @@ -163,6 +164,15 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, vector manifold correction term. Defaults to False. outflow (bool, optional): whether to include outflow at the domain boundaries, through exterior facet terms. Defaults to False. + advective_then_flux (bool, optional): whether to use the advective- + then-flux formulation. This uses the advective form of the + transport equation for all but the last steps of some + (potentially subcycled) Runge-Kutta scheme, before using the + conservative form for the final step to deliver a mass- + conserving increment. This optiona only makes sense to use with + Runge-Kutta, and should be used with the "linear" Runge-Kutta + formulation. Defaults to False, in which case the conservative + form is used for every step. """ super().__init__(equation, variable) @@ -170,6 +180,13 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, self.vector_manifold_correction = vector_manifold_correction self.outflow = outflow + if (advective_then_flux + and self.transport_equation_type != TransportEquationType.conservative): + raise ValueError( + 'DG Upwind: advective_then_flux form can only be used with ' + + 'the conservative form of the transport equation' + ) + # -------------------------------------------------------------------- # # Determine appropriate form to use # -------------------------------------------------------------------- # @@ -177,36 +194,52 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, if equation.domain.mesh.topological_dimension() == 1 and len(equation.domain.spaces("HDiv").shape) == 0: assert not vector_manifold_correction if self.transport_equation_type == TransportEquationType.advective: - form = upwind_advection_form_1d(self.domain, self.test, - self.field, - ibp=ibp, outflow=outflow) + form = upwind_advection_form_1d( + self.domain, self.test, self.field, ibp=ibp, + outflow=outflow + ) elif self.transport_equation_type == TransportEquationType.conservative: - form = upwind_continuity_form_1d(self.domain, self.test, - self.field, - ibp=ibp, outflow=outflow) + form = upwind_continuity_form_1d( + self.domain, self.test, self.field, ibp=ibp, + outflow=outflow + ) else: if self.transport_equation_type == TransportEquationType.advective: if vector_manifold_correction: - form = vector_manifold_advection_form(self.domain, - self.test, - self.field, ibp=ibp, - outflow=outflow) + form = vector_manifold_advection_form( + self.domain, self.test, self.field, ibp=ibp, + outflow=outflow + ) else: - form = upwind_advection_form(self.domain, self.test, - self.field, - ibp=ibp, outflow=outflow) + form = upwind_advection_form( + self.domain, self.test, self.field, ibp=ibp, + outflow=outflow + ) elif self.transport_equation_type == TransportEquationType.conservative: if vector_manifold_correction: - form = vector_manifold_continuity_form(self.domain, - self.test, - self.field, ibp=ibp, - outflow=outflow) + form = vector_manifold_continuity_form( + self.domain, self.test, self.field, ibp=ibp, + outflow=outflow + ) else: - form = upwind_continuity_form(self.domain, self.test, - self.field, - ibp=ibp, outflow=outflow) + form = upwind_continuity_form( + self.domain, self.test, self.field, ibp=ibp, + outflow=outflow + ) + + if advective_then_flux and vector_manifold_correction: + self.form_for_early_stages = vector_manifold_advection_form( + self.domain, self.test, self.field, ibp=ibp, + outflow=outflow + ) + + elif advective_then_flux: + self.form_for_early_stages = upwind_advection_form( + self.domain, self.test, self.field, ibp=ibp, + outflow=outflow + ) elif self.transport_equation_type == TransportEquationType.circulation: if outflow: diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py index 64e9545e..f34879bf 100644 --- a/gusto/time_discretisation/explicit_runge_kutta.py +++ b/gusto/time_discretisation/explicit_runge_kutta.py @@ -2,6 +2,7 @@ import numpy as np +from enum import Enum from firedrake import (Function, Constant, NonlinearVariationalProblem, NonlinearVariationalSolver) from firedrake.fml import replace_subject, all_terms, drop, keep @@ -12,7 +13,25 @@ from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation -__all__ = ["ForwardEuler", "ExplicitRungeKutta", "SSPRK3", "RK4", "Heun"] +__all__ = [ + "ForwardEuler", "ExplicitRungeKutta", "SSPRK3", "RK4", "Heun", + "RungeKuttaFormulation" +] + + +class RungeKuttaFormulation(Enum): + """ + Enumerator to describe the formulation of a Runge-Kutta scheme. + + The following Runge-Kutta methods are encoded here: + - `increment`: + - `predictor`: + - `linear`: + """ + + increment = 1595712 + predictor = 8234900 + linear = 269207 class ExplicitRungeKutta(ExplicitTimeDiscretisation): @@ -59,14 +78,14 @@ class ExplicitRungeKutta(ExplicitTimeDiscretisation): def __init__(self, domain, butcher_matrix, field_name=None, fixed_subcycles=None, subcycle_by_courant=None, - increment_form=True, solver_parameters=None, - limiter=None, options=None): + rk_formulation=RungeKuttaFormulation.increment, + solver_parameters=None, limiter=None, options=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the mesh and the compatible function spaces. - butcher_matrix (numpy array): A matrix containing the coefficients of - a butcher tableau defining a given Runge Kutta time discretisation. + butcher_matrix (numpy array): A matrix containing the coefficients + of a butcher tableau defining a given Runge Kutta scheme. field_name (str, optional): name of the field to be evolved. Defaults to None. fixed_subcycles (int, optional): the fixed number of sub-steps to @@ -76,11 +95,11 @@ def __init__(self, domain, butcher_matrix, field_name=None, make the scheme perform adaptive sub-cycling based on the Courant number. The specified argument is the maximum Courant for one sub-cycle. Defaults to None, in which case adaptive - sub-cycling is not used. This option cannot be specified with the - `fixed_subcycles` argument. - increment_form (bool, optional): whether to write the RK scheme in - "increment form", solving for increments rather than updated - fields. Defaults to True. + sub-cycling is not used. This option cannot be specified with + the `fixed_subcycles` argument. + rk_formulation (:class:`RungeKuttaFormulation`, optional): + an enumerator object, describing the formulation of the Runge- + Kutta scheme. Defaults to the increment form. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. limiter (:class:`Limiter` object, optional): a limiter to apply to @@ -97,7 +116,7 @@ def __init__(self, domain, butcher_matrix, field_name=None, limiter=limiter, options=options) self.butcher_matrix = butcher_matrix self.nbutcher = int(np.shape(self.butcher_matrix)[0]) - self.increment_form = increment_form + self.rk_formulation = rk_formulation @property def nStages(self): @@ -114,16 +133,21 @@ def setup(self, equation, apply_bcs=True, *active_labels): """ super().setup(equation, apply_bcs, *active_labels) - if not self.increment_form: + if self.rk_formulation == RungeKuttaFormulation.predictor: self.field_i = [Function(self.fs) for i in range(self.nStages+1)] - else: + elif self.rk_formulation == RungeKuttaFormulation.increment: self.k = [Function(self.fs) for i in range(self.nStages)] + else: + raise NotImplementedError( + 'Runge-Kutta formulation is not implemented' + ) @cached_property def solver(self): - if self.increment_form: + if self.rk_formulation == RungeKuttaFormulation.increment: return super().solver - else: + + elif self.rk_formulation == RungeKuttaFormulation.predictor: # In this case, don't set snes_type to ksp only, as we do want the # outer Newton iteration solver_list = [] @@ -140,11 +164,16 @@ def solver(self): solver_list.append(solver) return solver_list + else: + raise NotImplementedError( + 'Runge-Kutta formulation is not implemented' + ) + @cached_property def lhs(self): """Set up the discretisation's left hand side (the time derivative).""" - if self.increment_form: + if self.rk_formulation == RungeKuttaFormulation.increment: l = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.x_out, self.idx), @@ -152,7 +181,7 @@ def lhs(self): return l.form - else: + elif self.rk_formulation == RungeKuttaFormulation.predictor: lhs_list = [] for stage in range(self.nStages): l = self.residual.label_map( @@ -163,11 +192,16 @@ def lhs(self): return lhs_list + else: + raise NotImplementedError( + 'Runge-Kutta formulation is not implemented' + ) + @cached_property def rhs(self): """Set up the time discretisation's right hand side.""" - if self.increment_form: + if self.rk_formulation == RungeKuttaFormulation.increment: r = self.residual.label_map( all_terms, map_if_true=replace_subject(self.x1, old_idx=self.idx)) @@ -179,7 +213,7 @@ def rhs(self): # If there are no active labels, we may have no terms at this point # So that we can still do xnp1 = xn, put in a zero term here - if self.increment_form and len(r.terms) == 0: + if len(r.terms) == 0: logger.warning('No terms detected for RHS of explicit problem. ' + 'Adding a zero term to avoid failure.') null_term = Constant(0.0)*self.residual.label_map( @@ -191,7 +225,7 @@ def rhs(self): return r.form - else: + elif self.rk_formulation == RungeKuttaFormulation.predictor: rhs_list = [] for stage in range(self.nStages): @@ -217,9 +251,14 @@ def rhs(self): return rhs_list + else: + raise NotImplementedError( + 'Runge-Kutta formulation is not implemented' + ) + def solve_stage(self, x0, stage): - if self.increment_form: + if self.rk_formulation == RungeKuttaFormulation.increment: self.x1.assign(x0) for i in range(stage): @@ -241,7 +280,7 @@ def solve_stage(self, x0, stage): if self.limiter is not None: self.limiter.apply(self.x1) - else: + elif self.rk_formulation == RungeKuttaFormulation.predictor: # Set initial field if stage == 0: self.field_i[0].assign(x0) @@ -267,6 +306,11 @@ def solve_stage(self, x0, stage): if self.limiter is not None: self.limiter.apply(self.x1) + else: + raise NotImplementedError( + 'Runge-Kutta formulation is not implemented' + ) + def apply_cycle(self, x_out, x_in): """ Apply the time discretisation through a single sub-step. @@ -294,9 +338,12 @@ class ForwardEuler(ExplicitRungeKutta): k0 = F[y^n] \n y^(n+1) = y^n + dt*k0 \n """ - def __init__(self, domain, field_name=None, fixed_subcycles=None, - subcycle_by_courant=None, increment_form=True, - solver_parameters=None, limiter=None, options=None): + def __init__( + self, domain, field_name=None, + fixed_subcycles=None, subcycle_by_courant=None, + rk_formulation=RungeKuttaFormulation.increment, + solver_parameters=None, limiter=None, options=None + ): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -310,11 +357,11 @@ def __init__(self, domain, field_name=None, fixed_subcycles=None, make the scheme perform adaptive sub-cycling based on the Courant number. The specified argument is the maximum Courant for one sub-cycle. Defaults to None, in which case adaptive - sub-cycling is not used. This option cannot be specified with the - `fixed_subcycles` argument. - increment_form (bool, optional): whether to write the RK scheme in - "increment form", solving for increments rather than updated - fields. Defaults to True. + sub-cycling is not used. This option cannot be specified with + the `fixed_subcycles` argument. + rk_formulation (:class:`RungeKuttaFormulation`, optional): + an enumerator object, describing the formulation of the Runge- + Kutta scheme. Defaults to the increment form. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. limiter (:class:`Limiter` object, optional): a limiter to apply to @@ -324,11 +371,13 @@ def __init__(self, domain, field_name=None, fixed_subcycles=None, to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. """ + butcher_matrix = np.array([1.]).reshape(1, 1) + super().__init__(domain, butcher_matrix, field_name=field_name, fixed_subcycles=fixed_subcycles, subcycle_by_courant=subcycle_by_courant, - increment_form=increment_form, + rk_formulation=rk_formulation, solver_parameters=solver_parameters, limiter=limiter, options=options) @@ -343,9 +392,12 @@ class SSPRK3(ExplicitRungeKutta): k2 = F[y^n + (1/4)*dt*(k0+k1)] \n y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + 4*k2) \n """ - def __init__(self, domain, field_name=None, fixed_subcycles=None, - subcycle_by_courant=None, increment_form=True, - solver_parameters=None, limiter=None, options=None): + def __init__( + self, domain, field_name=None, + fixed_subcycles=None, subcycle_by_courant=None, + rk_formulation=RungeKuttaFormulation.increment, + solver_parameters=None, limiter=None, options=None + ): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -359,11 +411,11 @@ def __init__(self, domain, field_name=None, fixed_subcycles=None, make the scheme perform adaptive sub-cycling based on the Courant number. The specified argument is the maximum Courant for one sub-cycle. Defaults to None, in which case adaptive - sub-cycling is not used. This option cannot be specified with the - `fixed_subcycles` argument. - increment_form (bool, optional): whether to write the RK scheme in - "increment form", solving for increments rather than updated - fields. Defaults to True. + sub-cycling is not used. This option cannot be specified with + the `fixed_subcycles` argument. + rk_formulation (:class:`RungeKuttaFormulation`, optional): + an enumerator object, describing the formulation of the Runge- + Kutta scheme. Defaults to the increment form. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. limiter (:class:`Limiter` object, optional): a limiter to apply to @@ -373,12 +425,16 @@ def __init__(self, domain, field_name=None, fixed_subcycles=None, to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. """ - butcher_matrix = np.array([[1., 0., 0.], [1./4., 1./4., 0.], [1./6., 1./6., 2./3.]]) + butcher_matrix = np.array([ + [1., 0., 0.], + [1./4., 1./4., 0.], + [1./6., 1./6., 2./3.] + ]) super().__init__(domain, butcher_matrix, field_name=field_name, fixed_subcycles=fixed_subcycles, subcycle_by_courant=subcycle_by_courant, - increment_form=increment_form, + rk_formulation=rk_formulation, solver_parameters=solver_parameters, limiter=limiter, options=options) @@ -398,10 +454,12 @@ class RK4(ExplicitRungeKutta): where superscripts indicate the time-level. \n """ - def __init__(self, domain, field_name=None, fixed_subcycles=None, - subcycle_by_courant=None, increment_form=True, - solver_parameters=None, - limiter=None, options=None): + def __init__( + self, domain, field_name=None, + fixed_subcycles=None, subcycle_by_courant=None, + rk_formulation=RungeKuttaFormulation.increment, + solver_parameters=None, limiter=None, options=None + ): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -415,11 +473,11 @@ def __init__(self, domain, field_name=None, fixed_subcycles=None, make the scheme perform adaptive sub-cycling based on the Courant number. The specified argument is the maximum Courant for one sub-cycle. Defaults to None, in which case adaptive - sub-cycling is not used. This option cannot be specified with the - `fixed_subcycles` argument. - increment_form (bool, optional): whether to write the RK scheme in - "increment form", solving for increments rather than updated - fields. Defaults to True. + sub-cycling is not used. This option cannot be specified with + the `fixed_subcycles` argument. + rk_formulation (:class:`RungeKuttaFormulation`, optional): + an enumerator object, describing the formulation of the Runge- + Kutta scheme. Defaults to the increment form. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. limiter (:class:`Limiter` object, optional): a limiter to apply to @@ -429,11 +487,16 @@ def __init__(self, domain, field_name=None, fixed_subcycles=None, to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. """ - butcher_matrix = np.array([[0.5, 0., 0., 0.], [0., 0.5, 0., 0.], [0., 0., 1., 0.], [1./6., 1./3., 1./3., 1./6.]]) + butcher_matrix = np.array([ + [0.5, 0., 0., 0.], + [0., 0.5, 0., 0.], + [0., 0., 1., 0.], + [1./6., 1./3., 1./3., 1./6.] + ]) super().__init__(domain, butcher_matrix, field_name=field_name, fixed_subcycles=fixed_subcycles, subcycle_by_courant=subcycle_by_courant, - increment_form=increment_form, + rk_formulation=rk_formulation, solver_parameters=solver_parameters, limiter=limiter, options=options) @@ -451,9 +514,12 @@ class Heun(ExplicitRungeKutta): where superscripts indicate the time-level and subscripts indicate the stage number. """ - def __init__(self, domain, field_name=None, fixed_subcycles=None, - subcycle_by_courant=None, increment_form=True, - solver_parameters=None, limiter=None, options=None): + def __init__( + self, domain, field_name=None, + fixed_subcycles=None, subcycle_by_courant=None, + rk_formulation=RungeKuttaFormulation.increment, + solver_parameters=None, limiter=None, options=None + ): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -469,9 +535,9 @@ def __init__(self, domain, field_name=None, fixed_subcycles=None, for one sub-cycle. Defaults to None, in which case adaptive sub-cycling is not used. This option cannot be specified with the `fixed_subcycles` argument. - increment_form (bool, optional): whether to write the RK scheme in - "increment form", solving for increments rather than updated - fields. Defaults to True. + rk_formulation (:class:`RungeKuttaFormulation`, optional): + an enumerator object, describing the formulation of the Runge- + Kutta scheme. Defaults to the increment form. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. limiter (:class:`Limiter` object, optional): a limiter to apply to @@ -481,10 +547,14 @@ def __init__(self, domain, field_name=None, fixed_subcycles=None, to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. """ - butcher_matrix = np.array([[1., 0.], [0.5, 0.5]]) + + butcher_matrix = np.array([ + [1., 0.], + [0.5, 0.5] + ]) super().__init__(domain, butcher_matrix, field_name=field_name, fixed_subcycles=fixed_subcycles, subcycle_by_courant=subcycle_by_courant, - increment_form=increment_form, + rk_formulation=rk_formulation, solver_parameters=solver_parameters, limiter=limiter, options=options) diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py index eddb7245..2e584516 100644 --- a/gusto/time_discretisation/time_discretisation.py +++ b/gusto/time_discretisation/time_discretisation.py @@ -94,7 +94,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None, 'Time discretisation: suboption SUPG is currently not implemented within MixedOptions') else: raise RuntimeError( - f'Time discretisation: suboption wrapper {wrapper_name} not implemented') + f'Time discretisation: suboption wrapper {self.wrapper_name} not implemented') elif self.wrapper_name == "embedded_dg": self.wrapper = EmbeddedDGWrapper(self, options) elif self.wrapper_name == "recovered": diff --git a/integration-tests/equations/test_coupled_transport.py b/integration-tests/equations/test_coupled_transport.py index 6103d008..7923f008 100644 --- a/integration-tests/equations/test_coupled_transport.py +++ b/integration-tests/equations/test_coupled_transport.py @@ -98,9 +98,14 @@ def test_conservative_coupled_transport(tmpdir, m_X_space, tracer_setup): 'f2': EmbeddedDGOptions()} opts = MixedFSOptions(suboptions=suboptions) - transport_scheme = SSPRK3(domain, options=opts, increment_form=False) + transport_scheme = SSPRK3( + domain, options=opts, + rk_formulation=RungeKuttaFormulation.predictor + ) else: - transport_scheme = SSPRK3(domain, increment_form=False) + transport_scheme = SSPRK3( + domain, rk_formulation=RungeKuttaFormulation.predictor + ) transport_method = [DGUpwind(eqn, 'f1'), DGUpwind(eqn, 'f2')] diff --git a/integration-tests/model/test_conservative_transport_with_physics.py b/integration-tests/model/test_conservative_transport_with_physics.py index 67860f3b..1833e221 100644 --- a/integration-tests/model/test_conservative_transport_with_physics.py +++ b/integration-tests/model/test_conservative_transport_with_physics.py @@ -60,8 +60,10 @@ def run_conservative_transport_with_physics(dirname): # Time stepper time_varying_velocity = False stepper = SplitPrescribedTransport( - eqn, SSPRK3(domain, increment_form=False), io, time_varying_velocity, - transport_method, physics_schemes=physics_schemes) + eqn, SSPRK3(domain, rk_formulation=RungeKuttaFormulation.predictor), + io, time_varying_velocity, transport_method, + physics_schemes=physics_schemes + ) # ------------------------------------------------------------------------ # # Initial conditions diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py index 5146107b..6d484bfd 100644 --- a/integration-tests/model/test_time_discretisation.py +++ b/integration-tests/model/test_time_discretisation.py @@ -9,9 +9,12 @@ def run(timestepper, tmax, f_end): @pytest.mark.parametrize( - "scheme", ["ssprk3_increment", "TrapeziumRule", "ImplicitMidpoint", - "QinZhang", "RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth", - "Leapfrog", "AdamsMoulton", "AdamsMoulton", "ssprk3_predictor"]) + "scheme", [ + "ssprk3_increment", "TrapeziumRule", "ImplicitMidpoint", "QinZhang", + "RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth", "Leapfrog", + "AdamsMoulton", "AdamsMoulton", "ssprk3_predictor", "ssprk3_linear" + ] +) def test_time_discretisation(tmpdir, scheme, tracer_setup): if (scheme == "AdamsBashforth"): # Tighter stability constraints @@ -28,9 +31,11 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup): eqn = AdvectionEquation(domain, V, "f") if scheme == "ssprk3_increment": - transport_scheme = SSPRK3(domain, increment_form=True) + transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.increment) elif scheme == "ssprk3_predictor": - transport_scheme = SSPRK3(domain, increment_form=False) + transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.predictor) + elif scheme == "ssprk3_linear": + transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.linear) elif scheme == "TrapeziumRule": transport_scheme = TrapeziumRule(domain) elif scheme == "ImplicitMidpoint": diff --git a/integration-tests/transport/test_dg_transport.py b/integration-tests/transport/test_dg_transport.py index f581e252..67428b78 100644 --- a/integration-tests/transport/test_dg_transport.py +++ b/integration-tests/transport/test_dg_transport.py @@ -14,7 +14,9 @@ def run(timestepper, tmax, f_end): @pytest.mark.parametrize("geometry", ["slice", "sphere"]) -@pytest.mark.parametrize("equation_form", ["advective", "continuity"]) +@pytest.mark.parametrize("equation_form", [ + "advective", "continuity", "advective_then_flux" +]) def test_dg_transport_scalar(tmpdir, geometry, equation_form, tracer_setup): setup = tracer_setup(tmpdir, geometry) domain = setup.domain @@ -25,8 +27,13 @@ def test_dg_transport_scalar(tmpdir, geometry, equation_form, tracer_setup): else: eqn = ContinuityEquation(domain, V, "f") - transport_scheme = SSPRK3(domain) - transport_method = DGUpwind(eqn, "f") + + if equation_form == "advective_then_flux": + transport_method = DGUpwind(eqn, "f", advective_then_flux=True) + transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.linear) + else: + transport_method = DGUpwind(eqn, "f") + transport_scheme = SSPRK3(domain) time_varying_velocity = False timestepper = PrescribedTransport( From 8f14a42605e622fa8bf3abe33db7a501b1797121 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sat, 21 Sep 2024 13:32:44 +0100 Subject: [PATCH 2/3] add new test and fix lint --- gusto/spatial_methods/transport_methods.py | 2 +- .../transport/test_advective_then_flux.py | 109 ++++++++++++++++++ .../transport/test_dg_transport.py | 1 - 3 files changed, 110 insertions(+), 2 deletions(-) create mode 100644 integration-tests/transport/test_advective_then_flux.py diff --git a/gusto/spatial_methods/transport_methods.py b/gusto/spatial_methods/transport_methods.py index 29fa1d4b..a0006096 100644 --- a/gusto/spatial_methods/transport_methods.py +++ b/gusto/spatial_methods/transport_methods.py @@ -181,7 +181,7 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, self.outflow = outflow if (advective_then_flux - and self.transport_equation_type != TransportEquationType.conservative): + and self.transport_equation_type != TransportEquationType.conservative): raise ValueError( 'DG Upwind: advective_then_flux form can only be used with ' + 'the conservative form of the transport equation' diff --git a/integration-tests/transport/test_advective_then_flux.py b/integration-tests/transport/test_advective_then_flux.py new file mode 100644 index 00000000..413a1d1e --- /dev/null +++ b/integration-tests/transport/test_advective_then_flux.py @@ -0,0 +1,109 @@ +""" +Tests transport using a Runge-Kutta scheme with an advective-then-flux approach. +This should yield increments that are linear in the divergence (and thus +preserve a constant in divergence-free flow). +""" + +from gusto import * +from firedrake import ( + PeriodicRectangleMesh, cos, sin, SpatialCoordinate, + assemble, dx, pi, as_vector, errornorm, Function +) +import pytest + + +def setup_advective_then_flux(dirname, desirable_property): + + # ------------------------------------------------------------------------ # + # Model set up + # ------------------------------------------------------------------------ # + + # Time parameters + dt = 2. + + # Domain + domain_width = 2000. + ncells_1d = 10. + mesh = PeriodicRectangleMesh( + ncells_1d, ncells_1d, domain_width, domain_width, quadrilateral=True + ) + domain = Domain(mesh, dt, "RTCF", 1) + + # Equation + V_DG = domain.spaces('DG') + V_HDiv = domain.spaces("HDiv") + eqn = ContinuityEquation(domain, V_DG, "rho", Vu=V_HDiv) + + # IO + output = OutputParameters(dirname=dirname) + io = IO(domain, output) + + # Transport method + transport_scheme = SSPRK3( + domain, rk_formulation=RungeKuttaFormulation.linear, fixed_subcycles=3 + ) + transport_method = DGUpwind(eqn, "rho", advective_then_flux=True) + + # Timestepper + time_varying = False + stepper = PrescribedTransport( + eqn, transport_scheme, io, time_varying, transport_method + ) + + # ------------------------------------------------------------------------ # + # Initial Conditions + # ------------------------------------------------------------------------ # + + x, y = SpatialCoordinate(mesh) + + # Density is initially constant for both tests + rho_0 = 10.0 + + # Set the initial state from the configuration choice + if desirable_property == 'constancy': + # Divergence free velocity + num_steps = 5 + psi = Function(domain.spaces('H1')) + psi_expr = cos(2*pi*x/domain_width)*sin(2*pi*y/domain_width) + psi.interpolate(psi_expr) + u_expr = domain.perp(grad(psi_expr)) + + elif desirable_property == 'divergence_linearity': + # Divergent velocity + num_steps = 1 + u_expr = as_vector([ + cos(2*pi*x/domain_width)*sin(4*pi*y/domain_width), + -pi*sin(2*pi*x*y/(domain_width)**2) + ]) + + stepper.fields("rho").assign(Constant(rho_0)) + stepper.fields("u").project(u_expr) + + rho_true = Function(V_DG) + rho_true.assign(stepper.fields("rho")) + + return stepper, rho_true, dt, num_steps + + +@pytest.mark.parametrize("desirable_property", ["constancy", "divergence_linearity"]) +def test_advective_then_flux(tmpdir, desirable_property): + + # Setup and run + dirname = str(tmpdir) + + stepper, rho_true, dt, num_steps = \ + setup_advective_then_flux(dirname, desirable_property) + + # Run for five timesteps + stepper.run(t=0, tmax=dt*num_steps) + rho = stepper.fields("rho") + + # Check for divergence-linearity/constancy + assert errornorm(rho, rho_true) < 2e-13, \ + "advective-then-flux form is not yielding the correct answer" + + # Check for conservation + mass_initial = assemble(rho_true*dx) + mass_final = assemble(rho*dx) + assert abs(mass_final - mass_initial) < 1e-14, \ + "advective-then-flux form is not conservative" diff --git a/integration-tests/transport/test_dg_transport.py b/integration-tests/transport/test_dg_transport.py index 67428b78..63a3898a 100644 --- a/integration-tests/transport/test_dg_transport.py +++ b/integration-tests/transport/test_dg_transport.py @@ -27,7 +27,6 @@ def test_dg_transport_scalar(tmpdir, geometry, equation_form, tracer_setup): else: eqn = ContinuityEquation(domain, V, "f") - if equation_form == "advective_then_flux": transport_method = DGUpwind(eqn, "f", advective_then_flux=True) transport_scheme = SSPRK3(domain, rk_formulation=RungeKuttaFormulation.linear) From 850006006fb08f3f3db0c9b98a20a4248292b4a5 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 22 Sep 2024 17:55:40 +0100 Subject: [PATCH 3/3] working implementation of advective-then-flux --- .../skamarock_klemp_nonhydrostatic.py | 11 +- .../moist_thermal_williamson_5.py | 26 ++- examples/shallow_water/williamson_5.py | 12 +- gusto/core/labels.py | 2 + gusto/spatial_methods/transport_methods.py | 14 +- .../explicit_runge_kutta.py | 171 ++++++++++++++++-- .../time_discretisation.py | 1 + gusto/timestepping/timestepper.py | 21 ++- .../transport/test_advective_then_flux.py | 8 +- 9 files changed, 233 insertions(+), 33 deletions(-) diff --git a/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py b/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py index 777134d6..5b65fb86 100644 --- a/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py +++ b/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py @@ -21,7 +21,8 @@ Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, TrapeziumRule, SUPGOptions, CourantNumber, Perturbation, Gradient, CompressibleParameters, CompressibleEulerEquations, CompressibleSolver, - compressible_hydrostatic_balance, logger, RichardsonNumber + compressible_hydrostatic_balance, logger, RichardsonNumber, + RungeKuttaFormulation ) skamarock_klemp_nonhydrostatic_defaults = { @@ -106,13 +107,13 @@ def skamarock_klemp_nonhydrostatic( # Transport schemes theta_opts = SUPGOptions() transported_fields = [ - TrapeziumRule(domain, "u"), - SSPRK3(domain, "rho"), - SSPRK3(domain, "theta", options=theta_opts) + SSPRK3(domain, "u", subcycle_by_courant=0.25), + SSPRK3(domain, "rho", subcycle_by_courant=0.25, rk_formulation=RungeKuttaFormulation.linear), + SSPRK3(domain, "theta", subcycle_by_courant=0.25, options=theta_opts) ] transport_methods = [ DGUpwind(eqns, "u"), - DGUpwind(eqns, "rho"), + DGUpwind(eqns, "rho", advective_then_flux=True), DGUpwind(eqns, "theta", ibp=theta_opts.ibp) ] diff --git a/examples/shallow_water/moist_thermal_williamson_5.py b/examples/shallow_water/moist_thermal_williamson_5.py index 9ce7da77..305437d8 100644 --- a/examples/shallow_water/moist_thermal_williamson_5.py +++ b/examples/shallow_water/moist_thermal_williamson_5.py @@ -24,7 +24,8 @@ ShallowWaterParameters, ShallowWaterEquations, Sum, lonlatr_from_xyz, InstantRain, SWSaturationAdjustment, WaterVapour, CloudWater, Rain, GeneralIcosahedralSphereMesh, RelativeVorticity, - ZonalComponent, MeridionalComponent + ZonalComponent, MeridionalComponent, RungeKuttaFormulation, SSPRK3, + SemiImplicitQuasiNewton, ThermalSWSolver ) moist_thermal_williamson_5_defaults = { @@ -142,13 +143,28 @@ def gamma_v(x_in): gamma_r=gamma_r ) + transported_fields = [ + SSPRK3(domain, "u", subcycle_by_courant=0.25), + SSPRK3(domain, "D", subcycle_by_courant=0.25, rk_formulation=RungeKuttaFormulation.linear), + SSPRK3(domain, "b", subcycle_by_courant=0.25), + SSPRK3(domain, "water_vapour", subcycle_by_courant=0.25), + SSPRK3(domain, "cloud_water", subcycle_by_courant=0.25), + ] transport_methods = [ - DGUpwind(eqns, field_name) for field_name in eqns.field_names + DGUpwind(eqns, "u"), + DGUpwind(eqns, "D", advective_then_flux=True), + DGUpwind(eqns, "b"), + DGUpwind(eqns, "water_vapour"), + DGUpwind(eqns, "cloud_water") ] - # Timestepper - stepper = Timestepper( - eqns, RK4(domain), io, spatial_methods=transport_methods + # Linear solver + linear_solver = ThermalSWSolver(eqns) + + # Time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods, + linear_solver=linear_solver ) # ------------------------------------------------------------------------ # diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index cdb2b58a..41daa057 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -14,7 +14,7 @@ Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations, Sum, lonlatr_from_xyz, GeneralIcosahedralSphereMesh, ZonalComponent, - MeridionalComponent, RelativeVorticity + MeridionalComponent, RelativeVorticity, RungeKuttaFormulation ) williamson_5_defaults = { @@ -83,8 +83,14 @@ def williamson_5( io = IO(domain, output, diagnostic_fields=diagnostic_fields) # Transport schemes - transported_fields = [TrapeziumRule(domain, "u"), SSPRK3(domain, "D")] - transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")] + transported_fields = [ + SSPRK3(domain, "u", subcycle_by_courant=0.25), + SSPRK3(domain, "D", subcycle_by_courant=0.25, rk_formulation=RungeKuttaFormulation.linear) + ] + transport_methods = [ + DGUpwind(eqns, "u"), + DGUpwind(eqns, "D", advective_then_flux=True) + ] # Time stepper stepper = SemiImplicitQuasiNewton( diff --git a/gusto/core/labels.py b/gusto/core/labels.py index 5f928d6d..7936926e 100644 --- a/gusto/core/labels.py +++ b/gusto/core/labels.py @@ -97,6 +97,8 @@ def __call__(self, target, value=None): 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) +all_but_last = Label("all_but_last", validator=lambda value: type(value) in [LabelledForm, Term]) + # labels for terms in the equations time_derivative = Label("time_derivative") diff --git a/gusto/spatial_methods/transport_methods.py b/gusto/spatial_methods/transport_methods.py index a0006096..dbbd56b1 100644 --- a/gusto/spatial_methods/transport_methods.py +++ b/gusto/spatial_methods/transport_methods.py @@ -8,8 +8,10 @@ ) 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, - mass_weighted) +from gusto.core.labels import ( + prognostic, transport, transporting_velocity, ibp_label, mass_weighted, + all_but_last +) from gusto.core.logging import logger from gusto.spatial_methods.spatial_methods import SpatialMethod @@ -83,6 +85,10 @@ def replace_form(self, equation): # Create new term new_term = Term(self.form.form, original_term.labels) + # Add all_but_last form + if hasattr(self, "all_but_last_form"): + new_term = all_but_last(new_term, self.all_but_last_form) + # Check if this is a conservative transport if original_term.has_label(mass_weighted): # Extract the original and discretised mass_weighted terms @@ -230,13 +236,13 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, ) if advective_then_flux and vector_manifold_correction: - self.form_for_early_stages = vector_manifold_advection_form( + self.all_but_last_form = vector_manifold_advection_form( self.domain, self.test, self.field, ibp=ibp, outflow=outflow ) elif advective_then_flux: - self.form_for_early_stages = upwind_advection_form( + self.all_but_last_form = upwind_advection_form( self.domain, self.test, self.field, ibp=ibp, outflow=outflow ) diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py index f34879bf..4e7886e6 100644 --- a/gusto/time_discretisation/explicit_runge_kutta.py +++ b/gusto/time_discretisation/explicit_runge_kutta.py @@ -5,10 +5,11 @@ from enum import Enum from firedrake import (Function, Constant, NonlinearVariationalProblem, NonlinearVariationalSolver) -from firedrake.fml import replace_subject, all_terms, drop, keep +from firedrake.fml import replace_subject, all_terms, drop, keep, Term from firedrake.utils import cached_property +from firedrake.formmanipulation import split_form -from gusto.core.labels import time_derivative +from gusto.core.labels import time_derivative, all_but_last from gusto.core.logging import logger from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation @@ -23,10 +24,19 @@ class RungeKuttaFormulation(Enum): """ Enumerator to describe the formulation of a Runge-Kutta scheme. - The following Runge-Kutta methods are encoded here: - - `increment`: + The following Runge-Kutta methods for solving dy/dt = F(y) are encoded here: + - `increment`: \n + k_0 = F[y^n] \n + k_m = F[y^n - dt*\sum_{i=0}^{m-1} a_{m,i} * k_i], for m = 1 to M - 1 \n + y^{n+1} = y^n - dt*\sum_{i=0}^{M-1} b_i*k_i \n - `predictor`: + y^0 = y^n \n + y^m = q^0 - dt*\sum_{i=0}^{m-1} a_{m,i} * F[y^i], for m = 1 to M - 1 \n + y^{n+1} = y^0 - dt*\sum_{i=0}^{m-1} b_i * F[y^i] \n - `linear`: + y^0 = y^n \n + y^m = q^0 - dt*F[\sum_{i=0}^{m-1} a_{m,i} * y^i], for m = 1 to M - 1 \n + y^{n+1} = y^0 - dt*F[\sum_{i=0}^{m-1} b_i * y^i] \n """ increment = 1595712 @@ -134,9 +144,11 @@ def setup(self, equation, apply_bcs=True, *active_labels): super().setup(equation, apply_bcs, *active_labels) if self.rk_formulation == RungeKuttaFormulation.predictor: - self.field_i = [Function(self.fs) for i in range(self.nStages+1)] + self.field_i = [Function(self.fs) for _ in range(self.nStages+1)] elif self.rk_formulation == RungeKuttaFormulation.increment: - self.k = [Function(self.fs) for i in range(self.nStages)] + self.k = [Function(self.fs) for _ in range(self.nStages)] + elif self.rk_formulation == RungeKuttaFormulation.linear: + self.field_rhs = Function(self.fs) else: raise NotImplementedError( 'Runge-Kutta formulation is not implemented' @@ -149,21 +161,49 @@ def solver(self): elif self.rk_formulation == RungeKuttaFormulation.predictor: # In this case, don't set snes_type to ksp only, as we do want the - # outer Newton iteration + # outer Newton iteration. This is achieved by not calling the + # "super" method, in which the default snes_type is set to ksp_only solver_list = [] for stage in range(self.nStages): # setup linear solver using lhs and rhs defined in derived class problem = NonlinearVariationalProblem( self.lhs[stage].form - self.rhs[stage].form, - self.field_i[stage+1], bcs=self.bcs) + self.field_i[stage+1], bcs=self.bcs + ) solver_name = self.field_name+self.__class__.__name__+str(stage) solver = NonlinearVariationalSolver( problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) + options_prefix=solver_name + ) solver_list.append(solver) return solver_list + elif self.rk_formulation == RungeKuttaFormulation.linear: + # In this case, don't set snes_type to ksp only, as we do want the + # outer Newton iteration. This is achieved by not calling the + # "super" method, in which the default snes_type is set to ksp_only + problem = NonlinearVariationalProblem( + self.lhs - self.rhs[0], self.x1, bcs=self.bcs + ) + solver_name = self.field_name+self.__class__.__name__ + solver = NonlinearVariationalSolver( + problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name + ) + + # Set up problem for final step + problem_last = NonlinearVariationalProblem( + self.lhs - self.rhs[1], self.x1, bcs=self.bcs + ) + solver_name = self.field_name+self.__class__.__name__+'_last' + solver_last = NonlinearVariationalSolver( + problem_last, solver_parameters=self.solver_parameters, + options_prefix=solver_name + ) + + return solver, solver_last + else: raise NotImplementedError( 'Runge-Kutta formulation is not implemented' @@ -192,6 +232,14 @@ def lhs(self): return lhs_list + if self.rk_formulation == RungeKuttaFormulation.linear: + l = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x1, self.idx), + map_if_false=drop) + + return l.form + else: raise NotImplementedError( 'Runge-Kutta formulation is not implemented' @@ -251,6 +299,49 @@ def rhs(self): return rhs_list + elif self.rk_formulation == RungeKuttaFormulation.linear: + + r = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x0, old_idx=self.idx), + map_if_false=replace_subject(self.field_rhs, old_idx=self.idx) + ) + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=keep, + map_if_false=lambda t: -self.dt*t + ) + + # Set up all-but-last RHS + if self.idx is not None: + # If original function is in mixed function space, then ensure + # correct test function in the all-but-last form + r_all_but_last = self.residual.label_map( + lambda t: t.has_label(all_but_last), + map_if_true=lambda t: + Term(split_form(t.get(all_but_last).form)[self.idx].form, + t.labels), + map_if_false=keep + ) + else: + r_all_but_last = self.residual.label_map( + lambda t: t.has_label(all_but_last), + map_if_true=lambda t: Term(t.get(all_but_last).form, t.labels), + map_if_false=keep + ) + r_all_but_last = r_all_but_last.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x0, old_idx=self.idx), + map_if_false=replace_subject(self.field_rhs, old_idx=self.idx) + ) + r_all_but_last = r_all_but_last.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=keep, + map_if_false=lambda t: -self.dt*t + ) + + return r_all_but_last.form, r.form + else: raise NotImplementedError( 'Runge-Kutta formulation is not implemented' @@ -293,8 +384,8 @@ def solve_stage(self, x0, stage): # TODO: not implemented! Here we need to evaluate the m-th term # in the i-th RHS with field_m raise NotImplementedError( - 'Physics not implemented with RK schemes that do not use ' - + 'the increment form') + 'Physics not implemented with RK schemes that use the ' + + 'predictor form') if self.limiter is not None: self.limiter.apply(self.field_i[stage]) @@ -306,6 +397,62 @@ def solve_stage(self, x0, stage): if self.limiter is not None: self.limiter.apply(self.x1) + elif self.rk_formulation == RungeKuttaFormulation.linear: + + # Set combined index of stage and subcycle + cycle_stage = self.nStages*self.subcycle_idx + stage + + if stage == 0 and self.subcycle_idx == 0: + self.field_lhs = [Function(self.fs) for _ in range(self.nStages*self.ncycles)] + self.field_lhs[0].assign(self.x0) + + # All-but-last form ------------------------------------------------ + if (cycle_stage + 1 < self.ncycles*self.nStages): + # Build up RHS field to be evaluated + self.field_rhs.assign(0.0) + for i in range(stage+1): + i_cycle_stage = self.nStages*self.subcycle_idx + i + self.field_rhs.assign( + self.field_rhs + + self.butcher_matrix[stage, i]*self.field_lhs[i_cycle_stage] + ) + + # Evaluate physics and apply limiter, if necessary + for evaluate in self.evaluate_source: + evaluate(self.field_rhs, self.dt) + if self.limiter is not None: + self.limiter.apply(self.field_rhs) + + # Solve problem, placing solution in self.x1 + self.solver[0].solve() + + # Store LHS + self.field_lhs[cycle_stage+1].assign(self.x1) + + # Last stage and last subcycle ------------------------------------- + else: + # Build up RHS field to be evaluated + self.field_rhs.assign(0.0) + for i in range(self.ncycles*self.nStages): + j = i % self.nStages + self.field_rhs.assign( + self.field_rhs + + self.butcher_matrix[self.nStages-1, j]*self.field_lhs[i] + ) + + # Evaluate physics and apply limiter, if necessary + for evaluate in self.evaluate_source: + evaluate(self.field_rhs, self.original_dt) + if self.limiter is not None: + self.limiter.apply(self.field_rhs) + + # Solve problem, placing solution in self.x1 + self.solver[1].solve() + + # Final application of limiter + if self.limiter is not None: + self.limiter.apply(self.x1) + else: raise NotImplementedError( 'Runge-Kutta formulation is not implemented' @@ -319,6 +466,8 @@ def apply_cycle(self, x_out, x_in): x_in (:class:`Function`): the input field. x_out (:class:`Function`): the output field to be computed. """ + + # TODO: is this limiter application necessary? if self.limiter is not None: self.limiter.apply(x_in) diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py index 2e584516..fdfe53f1 100644 --- a/gusto/time_discretisation/time_discretisation.py +++ b/gusto/time_discretisation/time_discretisation.py @@ -423,6 +423,7 @@ def apply(self, x_out, x_in): self.x0.assign(x_in) for i in range(self.ncycles): + self.subcycle_idx = i self.apply_cycle(self.x1, self.x0) self.x0.assign(self.x1) x_out.assign(self.x1) diff --git a/gusto/timestepping/timestepper.py b/gusto/timestepping/timestepper.py index d4f63be8..bfd45613 100644 --- a/gusto/timestepping/timestepper.py +++ b/gusto/timestepping/timestepper.py @@ -2,7 +2,7 @@ from abc import ABCMeta, abstractmethod, abstractproperty from firedrake import Function, Projector, split -from firedrake.fml import drop, Term +from firedrake.fml import drop, Term, LabelledForm from pyop2.profiling import timed_stage from gusto.equations import PrognosticEquationSet from gusto.core import TimeLevelFields, StateFields @@ -133,6 +133,25 @@ def setup_transporting_velocity(self, scheme): scheme.residual = transporting_velocity.update_value(scheme.residual, uadv) + # Now also replace transporting velocity in the terms that are + # contained in labels + for idx, t in enumerate(scheme.residual.terms): + if t.has_label(transporting_velocity): + for label in t.labels.keys(): + if type(t.labels[label]) is LabelledForm: + t.labels[label] = t.labels[label].label_map( + lambda s: s.has_label(transporting_velocity), + map_if_true=lambda s: + Term(ufl.replace( + s.form, + {s.get(transporting_velocity): uadv}), + s.labels + ) + ) + + scheme.residual.terms[idx].labels[label] = \ + transporting_velocity.update_value(t.labels[label], uadv) + def log_timestep(self): """ Logs the start of a time step. diff --git a/integration-tests/transport/test_advective_then_flux.py b/integration-tests/transport/test_advective_then_flux.py index 413a1d1e..877cccbb 100644 --- a/integration-tests/transport/test_advective_then_flux.py +++ b/integration-tests/transport/test_advective_then_flux.py @@ -7,7 +7,7 @@ from gusto import * from firedrake import ( PeriodicRectangleMesh, cos, sin, SpatialCoordinate, - assemble, dx, pi, as_vector, errornorm, Function + assemble, dx, pi, as_vector, errornorm, Function, div, as_vector ) import pytest @@ -66,7 +66,7 @@ def setup_advective_then_flux(dirname, desirable_property): psi = Function(domain.spaces('H1')) psi_expr = cos(2*pi*x/domain_width)*sin(2*pi*y/domain_width) psi.interpolate(psi_expr) - u_expr = domain.perp(grad(psi_expr)) + u_expr = as_vector([-psi.dx(1), psi.dx(0)]) elif desirable_property == 'divergence_linearity': # Divergent velocity @@ -80,7 +80,7 @@ def setup_advective_then_flux(dirname, desirable_property): stepper.fields("u").project(u_expr) rho_true = Function(V_DG) - rho_true.assign(stepper.fields("rho")) + rho_true.interpolate(rho_0*(1.0 - dt*div(stepper.fields('u')))) return stepper, rho_true, dt, num_steps @@ -99,7 +99,7 @@ def test_advective_then_flux(tmpdir, desirable_property): rho = stepper.fields("rho") # Check for divergence-linearity/constancy - assert errornorm(rho, rho_true) < 2e-13, \ + assert errornorm(rho, rho_true) < 1e-11, \ "advective-then-flux form is not yielding the correct answer" # Check for conservation