From af8f6040e3eab621b949e69e754fe34edcb76577 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 23 Oct 2024 11:21:44 +0100 Subject: [PATCH 01/15] Split HV added. Works for cartesian --- gusto/core/labels.py | 2 + gusto/equations/common_forms.py | 73 ++++++++- gusto/spatial_methods/spatial_methods.py | 2 +- gusto/spatial_methods/transport_methods.py | 170 ++++++++++++++++++++- 4 files changed, 240 insertions(+), 7 deletions(-) diff --git a/gusto/core/labels.py b/gusto/core/labels.py index 5f928d6d2..f882046f8 100644 --- a/gusto/core/labels.py +++ b/gusto/core/labels.py @@ -92,6 +92,8 @@ def __call__(self, target, value=None): # ---------------------------------------------------------------------------- # implicit = Label("implicit") explicit = Label("explicit") +horizontal = Label("horizontal") +vertical = Label("vertical") 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]) diff --git a/gusto/equations/common_forms.py b/gusto/equations/common_forms.py index 91ce01368..64a9594f9 100644 --- a/gusto/equations/common_forms.py +++ b/gusto/equations/common_forms.py @@ -7,7 +7,7 @@ from firedrake.fml import subject, drop from gusto.core.configuration import TransportEquationType from gusto.core.labels import (transport, transporting_velocity, diffusion, - prognostic, linearisation) + prognostic, linearisation, horizontal, vertical) __all__ = ["advection_form", "advection_form_1d", "continuity_form", "continuity_form_1d", "vector_invariant_form", @@ -15,7 +15,7 @@ "diffusion_form", "diffusion_form_1d", "linear_advection_form", "linear_continuity_form", "linear_continuity_form_1d", - "split_continuity_form", "tracer_conservative_form"] + "split_continuity_form", "tracer_conservative_form", "split_hv_advective_form"] def advection_form(test, q, ubar): @@ -366,3 +366,72 @@ def tracer_conservative_form(test, q, rho, ubar): form = transporting_velocity(L, ubar) return transport(form, TransportEquationType.tracer_conservative) + + +def split_hv_advective_form(equation, field_name): + u""" + Splits advective term into horizontal and vertical terms. + This describes splitting u.∇(q) terms into u.(∇_h)q and w dq/dz, + for transporting velocity u and transported q. + Args: + equation (:class:`PrognosticEquation`): the model's equation. + Returns: + :class:`PrognosticEquation`: the model's equation. + """ + k = equation.domain.k # vertical unit vector + for t in equation.residual: + print(t.get(prognostic)) + if (t.get(transport) == TransportEquationType.advective and t.get(prognostic) == field_name): + # Get fields and test functions + subj = t.get(subject) + idx = equation.field_names.index(field_name) + W = equation.function_space + test = TestFunctions(W)[idx] + q = split(subj)[idx] + # u is either a prognostic or prescribed field + if (hasattr(equation, "field_names") + and 'u' in equation.field_names): + u_idx = equation.field_names.index('u') + uadv = split(equation.X)[u_idx] + elif 'u' in equation.prescribed_fields._field_names: + uadv = equation.prescribed_fields('u') + else: + raise ValueError('Cannot get velocity field') + + # Create new advective and divergence terms + u_vertical = k*inner(uadv, k) + u_horizontal = uadv - u_vertical + + L = inner(test, dot(u_vertical, grad(q)))*dx + vertical_adv_term = prognostic(vertical(transport(transporting_velocity(L, uadv), TransportEquationType.advective)), field_name) + L = inner(test, dot(u_horizontal, grad(q)))*dx + horizontal_adv_term = prognostic( horizontal(transport(transporting_velocity(L, uadv), TransportEquationType.advective)), field_name) + #horizontal_adv_term = prognostic(advection_form(test, q, u_horizontal), field_name) + + + # # Add linearisations of new terms if required + if (t.has_label(linearisation)): + u_trial = TrialFunctions(W)[u_idx] + u_trial_vert = k*inner(u_trial, k) + u_trial_horiz = u_trial - u_trial_vert + qbar = split(equation.X_ref)[idx] + # Add linearisation to adv_term + linear_hori_term = horizontal(transport(transporting_velocity(test*dot(u_trial_horiz, grad(qbar))*dx, u_trial), TransportEquationType.advective)) + adv_horiz_term = linearisation(horizontal_adv_term, linear_hori_term) + # Add linearisation to div_term + linear_vert_term = horizontal(transport(transporting_velocity(test*dot(u_trial_vert, grad(qbar))*dx, u_trial), TransportEquationType.advective)) + adv_vert_term = linearisation(vertical_adv_term, linear_vert_term) + else: + adv_vert_term = vertical_adv_term + adv_horiz_term = horizontal_adv_term + # Drop old term + equation.residual = equation.residual.label_map( + lambda t: t.get(transport) == TransportEquationType.advective and t.get(prognostic) == field_name, + map_if_true=drop) + + + # Add new terms onto residual + equation.residual += subject(adv_horiz_term + adv_vert_term, subj) + + + return equation diff --git a/gusto/spatial_methods/spatial_methods.py b/gusto/spatial_methods/spatial_methods.py index 6ad214dcd..ba1cb7a00 100644 --- a/gusto/spatial_methods/spatial_methods.py +++ b/gusto/spatial_methods/spatial_methods.py @@ -44,7 +44,7 @@ def __init__(self, equation, variable, term_label): map_if_true=keep, map_if_false=drop) num_terms = len(self.original_form.terms) - assert num_terms == 1, f'Unable to find {term_label.label} term ' \ + assert num_terms >= 1, f'Unable to find {term_label.label} term ' \ + f'for {variable}. {num_terms} found' def replace_form(self, equation): diff --git a/gusto/spatial_methods/transport_methods.py b/gusto/spatial_methods/transport_methods.py index dc2f9d6e2..0df66115f 100644 --- a/gusto/spatial_methods/transport_methods.py +++ b/gusto/spatial_methods/transport_methods.py @@ -9,11 +9,11 @@ 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) + mass_weighted, horizontal, vertical) from gusto.core.logging import logger from gusto.spatial_methods.spatial_methods import SpatialMethod -__all__ = ["DefaultTransport", "DGUpwind"] +__all__ = ["DefaultTransport", "DGUpwind", "Split_DGUpwind"] # ---------------------------------------------------------------------------- # @@ -105,8 +105,43 @@ def replace_form(self, equation): map_if_true=lambda t: new_term) else: - raise RuntimeError('Found multiple transport terms for ' - + f'{self.variable}. {len(original_form.terms)} found') + horizontal_form = original_form = equation.residual.label_map( + lambda t: t.has_label(transport) and t.has_label(horizontal) and t.get(prognostic) == self.variable, + map_if_true=keep, map_if_false=drop + ) + vertical_form = original_form = equation.residual.label_map( + lambda t: t.has_label(transport) and t.has_label(vertical) and t.get(prognostic) == self.variable, + map_if_true=keep, map_if_false=drop + ) + + # Replace form + horizontal_term = horizontal_form.terms[0] + vertical_term = vertical_form.terms[0] + + # Update transporting velocity + new_horizontal_transporting_velocity = self.form_h.terms[0].get(transporting_velocity) + new_vertical_transporting_velocity = self.form_v.terms[0].get(transporting_velocity) + horizontal_term = transporting_velocity.update_value(horizontal_term, new_horizontal_transporting_velocity) + vertical_term = transporting_velocity.update_value(vertical_term, new_vertical_transporting_velocity) + + # Create new term + new_horizontal_term = Term(self.form_h.form, horizontal_term.labels) + new_vertical_term = Term(self.form_v.form, vertical_term.labels) + + # Check if this is a conservative transport + if horizontal_term.has_label(mass_weighted) or vertical_term.has_label(mass_weighted): + raise RuntimeError('Mass weighted transport terms not yet supported for multiple terms') + + # Replace original term with new term + equation.residual = equation.residual.label_map( + lambda t: t.has_label(transport) and t.has_label(horizontal) and t.get(prognostic) == self.variable, + map_if_true=lambda t: new_horizontal_term) + + equation.residual = equation.residual.label_map( + lambda t: t.has_label(transport) and t.has_label(vertical) and t.get(prognostic) == self.variable, + map_if_true=lambda t: new_vertical_term) + + # ---------------------------------------------------------------------------- # @@ -238,9 +273,136 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, self.form = form +class Split_DGUpwind(TransportMethod): + """ + The Discontinuous Galerkin Upwind transport scheme. + Discretises the gradient of a field weakly, taking the upwind value of the + transported variable at facets. + """ + def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, + vector_manifold_correction=False, outflow=False): + """ + Args: + equation (:class:`PrognosticEquation`): the equation, which includes + a transport term. + variable (str): name of the variable to set the transport scheme for + ibp (:class:`IntegrateByParts`, optional): an enumerator for how + many times to integrate by parts. Defaults to `ONCE`. + vector_manifold_correction (bool, optional): whether to include a + 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. + """ + + super().__init__(equation, variable) + self.ibp = ibp + self.vector_manifold_correction = vector_manifold_correction + self.outflow = outflow + + # -------------------------------------------------------------------- # + # Determine appropriate form to use + # -------------------------------------------------------------------- # + # first check for 1d mesh and scalar velocity space + 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) + elif self.transport_equation_type == TransportEquationType.conservative: + form = upwind_continuity_form_1d(self.domain, self.test, + self.field, + ibp=ibp, outflow=outflow) + + else: + if self.transport_equation_type == TransportEquationType.advective: + + form_h, form_v = split_upwind_advection_form(self.domain, self.test, + self.field, + ibp=ibp, outflow=outflow) + + else: + raise NotImplementedError('Upwind transport scheme has not been ' + + 'implemented for this transport equation type') + self.form_v = form_v + self.form_h = form_h + + + # ---------------------------------------------------------------------------- # # Forms for DG Upwind transport # ---------------------------------------------------------------------------- # +def split_upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): + u""" + The form corresponding to the DG upwind advective transport operator. + This discretises (u.∇)q, for transporting velocity u and transported + variable q. An upwind discretisation is used for the facet terms when the + form is integrated by parts. + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ibp (:class:`IntegrateByParts`, optional): an enumerator representing + the number of times to integrate by parts. Defaults to + `IntegrateByParts.ONCE`. + outflow (bool, optional): whether to include outflow at the domain + boundaries, through exterior facet terms. Defaults to False. + Raises: + ValueError: Can only use outflow option when the integration by parts + option is not "never". + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + if outflow and ibp == IntegrateByParts.NEVER: + raise ValueError("outflow is True and ibp is None are incompatible options") + Vu = domain.spaces("HDiv") + k=domain.k + dS_ = (dS_v + dS_h) if Vu.extruded else dS + ubar = Function(Vu) + ubar_v = k*inner(ubar, k) + ubar_h = ubar - ubar_v + + if ibp == IntegrateByParts.ONCE: + L_h = -inner(div(outer(test, ubar_h)), q)*dx + L_v = -inner(div(outer(test, ubar_v)), q)*dx + else: + L_h = inner(outer(test, ubar_h), grad(q))*dx + L_v = inner(outer(test, ubar_v), grad(q))*dx + + if ibp != IntegrateByParts.NEVER: + n = FacetNormal(domain.mesh) + un_h = 0.5*(dot(ubar_h, n) + abs(dot(ubar_h, n))) + + L_h += dot(jump(test), (un_h('+')*q('+') - un_h('-')*q('-')))*dS_ + + un_v = 0.5*(dot(ubar_v, n) + abs(dot(ubar_v, n))) + + L_v += dot(jump(test), (un_v('+')*q('+') - un_v('-')*q('-')))*dS_ + + if ibp == IntegrateByParts.TWICE: + L_h -= (inner(test('+'), dot(ubar_h('+'), n('+'))*q('+')) + + inner(test('-'), dot(ubar_h('-'), n('-'))*q('-')))*dS_ + + L_v -= (inner(test('+'), dot(ubar_v('+'), n('+'))*q('+')) + + inner(test('-'), dot(ubar_v('-'), n('-'))*q('-')))*dS_ + + if outflow: + n = FacetNormal(domain.mesh) + un_h = 0.5*(dot(ubar_h, n) + abs(dot(ubar_h, n))) + L_h += test*un_h*q*(ds_v + ds_t + ds_b) + + un_v = 0.5*(dot(ubar_v, n) + abs(dot(ubar_v, n))) + L_v += test*un_v*q*(ds_v + ds_t + ds_b) + + form_h = transporting_velocity(L_h, ubar) + form_v = transporting_velocity(L_v, ubar) + labelled_form_h = ibp_label(transport(form_h, TransportEquationType.advective), ibp) + labelled_form_v = ibp_label(transport(form_v, TransportEquationType.advective), ibp) + return labelled_form_h, labelled_form_v + + def upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): u""" The form corresponding to the DG upwind advective transport operator. From 13e2eb90aa60cd3b63c44304c75182fdcf169029 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Tue, 5 Nov 2024 16:45:09 +0000 Subject: [PATCH 02/15] Testing spherical nature --- .../dcmip_3_1_gravity_wave.py | 57 ++++++++++++++----- 1 file changed, 43 insertions(+), 14 deletions(-) diff --git a/examples/compressible_euler/dcmip_3_1_gravity_wave.py b/examples/compressible_euler/dcmip_3_1_gravity_wave.py index 2fccfbe3f..80d6e6b53 100644 --- a/examples/compressible_euler/dcmip_3_1_gravity_wave.py +++ b/examples/compressible_euler/dcmip_3_1_gravity_wave.py @@ -17,7 +17,10 @@ TrapeziumRule, SUPGOptions, lonlatr_from_xyz, CompressibleParameters, CompressibleEulerEquations, CompressibleSolver, ZonalComponent, compressible_hydrostatic_balance, Perturbation, GeneralCubedSphereMesh, + Timestepper, IMEX_SSP3, split_continuity_form, explicit, implicit, + time_derivative, transport, IMEXRungeKutta ) +import numpy as np dcmip_3_1_gravity_wave_defaults = { 'ncells_per_edge': 8, @@ -85,6 +88,13 @@ def dcmip_3_1_gravity_wave( eqns = CompressibleEulerEquations( domain, parameters, u_transport_option=u_eqn_type ) + print("Opt Cores:", eqns.X.function_space().dim()/50000.) + eqns = split_continuity_form(eqns) + eqns.label_terms(lambda t: not any(t.has_label(time_derivative, transport)), implicit) + #eqns.label_terms(lambda t: not any(t.has_label(time_derivative, horizontal)), implicit) + eqns.label_terms(lambda t: t.has_label(transport), explicit) + # eqns.label_terms(lambda t: t.has_label(transport) and t.has_label(vertical), implicit) + # eqns.label_terms(lambda t: t.has_label(transport) and not any(t.has_label(horizontal, vertical)), explicit) # I/O output = OutputParameters( @@ -96,24 +106,43 @@ def dcmip_3_1_gravity_wave( io = IO(domain, output, diagnostic_fields=diagnostic_fields) # Transport schemes - transported_fields = [ - TrapeziumRule(domain, "u"), - SSPRK3(domain, "rho", fixed_subcycles=2), - SSPRK3(domain, "theta", options=SUPGOptions(), fixed_subcycles=2) - ] + # transported_fields = [ + # TrapeziumRule(domain, "u"), + # SSPRK3(domain, "rho", fixed_subcycles=2), + # SSPRK3(domain, "theta", options=SUPGOptions(), fixed_subcycles=2) + # ] transport_methods = [ DGUpwind(eqns, field) for field in ["u", "rho", "theta"] ] - # Linear solver - linear_solver = CompressibleSolver(eqns) - - # Time stepper - stepper = SemiImplicitQuasiNewton( - eqns, io, transported_fields, transport_methods, - linear_solver=linear_solver - ) - + nl_solver_parameters = { + "snes_converged_reason": None, + "snes_lag_preconditioner_persists":None, + "snes_lag_preconditioner":-2, + "mat_type": "matfree", + "ksp_type": "gmres", + "ksp_converged_reason": None, + "ksp_atol": 1e-5, + "ksp_rtol": 1e-5, + "ksp_max_it": 400, + "pc_type": "python", + "pc_python_type": "firedrake.AssembledPC", + "assembled_pc_star_sub_sub_pc_type": "lu", + "assembled_pc_type": "python", + "assembled_pc_python_type": "firedrake.ASMStarPC", + "assembled_pc_star_construct_dim": 0, + "assembled_pc_star_sub_sub_pc_factor_mat_ordering_type": "rcm", + "assembled_pc_star_sub_sub_pc_factor_reuse_ordering": None, + "assembled_pc_star_sub_sub_pc_factor_reuse_fill": None, + "assembled_pc_star_sub_sub_pc_factor_fill": 1.2} + + + # IMEX time stepper + + butcher_imp = np.array([[0.0, 0.0], [0.0, 0.5], [0.0, 1.]]) + butcher_exp = np.array([[0.0, 0.0], [0.5, 0.0], [0.0, 1.]]) + scheme = IMEXRungeKutta(domain, butcher_imp, butcher_exp, solver_parameters=nl_solver_parameters) + stepper = Timestepper(eqns, scheme, io, transport_methods) # ------------------------------------------------------------------------ # # Initial conditions # ------------------------------------------------------------------------ # From 20c3846c405fea3a4be8f79021a37adf53c9246b Mon Sep 17 00:00:00 2001 From: atb1995 Date: Fri, 8 Nov 2024 09:43:18 +0000 Subject: [PATCH 03/15] Intermediate commit --- examples/compressible_euler/dcmip_3_1_gravity_wave.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/compressible_euler/dcmip_3_1_gravity_wave.py b/examples/compressible_euler/dcmip_3_1_gravity_wave.py index c9c93a7cb..91801840c 100644 --- a/examples/compressible_euler/dcmip_3_1_gravity_wave.py +++ b/examples/compressible_euler/dcmip_3_1_gravity_wave.py @@ -92,8 +92,8 @@ def dcmip_3_1_gravity_wave( print("Opt Cores:", eqns.X.function_space().dim()/50000.) eqns = split_continuity_form(eqns) eqns = split_hv_advective_form(eqns, "theta") - # eqns = split_hv_advective_form(eqns, "rho") - #eqns = split_hv_advective_form(eqns, "u") + eqns = split_hv_advective_form(eqns, "rho") + eqns = split_hv_advective_form(eqns, "u") # eqns.label_terms(lambda t: not any(t.has_label(time_derivative, transport)), implicit) # #eqns.label_terms(lambda t: not any(t.has_label(time_derivative, horizontal)), implicit) # eqns.label_terms(lambda t: t.has_label(transport), explicit) @@ -103,7 +103,7 @@ def dcmip_3_1_gravity_wave( eqns.label_terms(lambda t: not any(t.has_label(time_derivative, transport)), implicit) #eqns.label_terms(lambda t: not any(t.has_label(time_derivative, horizontal)), implicit) eqns.label_terms(lambda t: t.has_label(transport) and t.has_label(horizontal), explicit) - eqns.label_terms(lambda t: t.has_label(transport) and t.has_label(vertical), implicit) + eqns.label_terms(lambda t: t.has_label(transport) and t.has_label(vertical), explicit) eqns.label_terms(lambda t: t.has_label(transport) and not (t.has_label(horizontal) and t.has_label(vertical)), explicit) # I/O output = OutputParameters( @@ -117,11 +117,11 @@ def dcmip_3_1_gravity_wave( # Transport schemes # transported_fields = [ # TrapeziumRule(domain, "u"), - # SSPRK3(domain, "rho", fixed_subcycles=2), + # SSPRK3(domain, "rho", fixed_subcycles=2), # SSPRK3(domain, "theta", options=SUPGOptions(), fixed_subcycles=2) # ] transport_methods = [ - DGUpwind(eqns, "u"), DGUpwind(eqns, "rho") + Split_DGUpwind(eqns, "u"), Split_DGUpwind(eqns, "rho"), Split_DGUpwind(eqns, "theta") ] nl_solver_parameters = { From 8354245b811fbc96d9633dfa54959ac1c9fc2571 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 5 Mar 2025 15:24:12 +0000 Subject: [PATCH 04/15] Tidy of code, exceptions added, test added --- .../dcmip_3_1_gravity_wave.py | 15 +++-- gusto/equations/common_forms.py | 26 ++++---- gusto/spatial_methods/transport_methods.py | 61 ++++++++---------- .../transport/test_split_dg_transport.py | 64 +++++++++++++++++++ 4 files changed, 115 insertions(+), 51 deletions(-) create mode 100644 integration-tests/transport/test_split_dg_transport.py diff --git a/examples/compressible_euler/dcmip_3_1_gravity_wave.py b/examples/compressible_euler/dcmip_3_1_gravity_wave.py index 6bf9e2b64..6862dd94f 100644 --- a/examples/compressible_euler/dcmip_3_1_gravity_wave.py +++ b/examples/compressible_euler/dcmip_3_1_gravity_wave.py @@ -19,7 +19,6 @@ compressible_hydrostatic_balance, Perturbation, GeneralCubedSphereMesh, SubcyclingOptions ) -import numpy as np dcmip_3_1_gravity_wave_defaults = { 'ncells_per_edge': 8, @@ -69,7 +68,7 @@ def dcmip_3_1_gravity_wave( # ------------------------------------------------------------------------ # element_order = 1 - u_eqn_type = 'vector_advection_form' + u_eqn_type = 'vector_invariant_form' # ------------------------------------------------------------------------ # # Set up model objects @@ -87,6 +86,7 @@ def dcmip_3_1_gravity_wave( eqns = CompressibleEulerEquations( domain, parameters, u_transport_option=u_eqn_type ) + # I/O output = OutputParameters( dirname=dirname, dumpfreq=dumpfreq, dump_vtus=False, dump_nc=True @@ -107,11 +107,18 @@ def dcmip_3_1_gravity_wave( ), ] transport_methods = [ - DGUpwind(eqns, "u"), DGUpwind(eqns, "rho"), DGUpwind(eqns, "theta") + DGUpwind(eqns, field) for field in ["u", "rho", "theta"] ] + # Linear solver + linear_solver = CompressibleSolver(eqns) + + # Time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods, + linear_solver=linear_solver + ) - stepper = SemiImplicitQuasiNewton(eqns, io, spatial_methods=transport_methods, transported_fields=transported_fields) # ------------------------------------------------------------------------ # # Initial conditions # ------------------------------------------------------------------------ # diff --git a/gusto/equations/common_forms.py b/gusto/equations/common_forms.py index 6b90d1c49..8ea031fde 100644 --- a/gusto/equations/common_forms.py +++ b/gusto/equations/common_forms.py @@ -363,17 +363,21 @@ def split_hv_advective_form(equation, field_name): if (t.get(transport) == TransportEquationType.advective and t.get(prognostic) == field_name): # Get fields and test functions subj = t.get(subject) - idx = equation.field_names.index(field_name) - W = equation.function_space - test = TestFunctions(W)[idx] - q = split(subj)[idx] + # u is either a prognostic or prescribed field if (hasattr(equation, "field_names") and 'u' in equation.field_names): + idx = equation.field_names.index(field_name) + W = equation.function_space + test = TestFunctions(W)[idx] + q = split(subj)[idx] u_idx = equation.field_names.index('u') uadv = split(equation.X)[u_idx] elif 'u' in equation.prescribed_fields._field_names: uadv = equation.prescribed_fields('u') + q = subj + W = equation.function_space + test = TestFunction(W) else: raise ValueError('Cannot get velocity field') @@ -381,14 +385,10 @@ def split_hv_advective_form(equation, field_name): u_vertical = k*inner(uadv, k) u_horizontal = uadv - u_vertical - L = inner(test, dot(u_vertical, grad(q)))*dx - vertical_adv_term = prognostic(vertical(transport(transporting_velocity(L, uadv), TransportEquationType.advective)), field_name) - L = inner(test, dot(u_horizontal, grad(q)))*dx - horizontal_adv_term = prognostic( horizontal(transport(transporting_velocity(L, uadv), TransportEquationType.advective)), field_name) - #horizontal_adv_term = prognostic(advection_form(test, q, u_horizontal), field_name) - + vertical_adv_term = prognostic(vertical(transport(transporting_velocity(inner(test, dot(u_vertical, grad(q)))*dx, uadv), TransportEquationType.advective)), field_name) + horizontal_adv_term = prognostic(horizontal(transport(transporting_velocity(inner(test, dot(u_horizontal, grad(q)))*dx, uadv), TransportEquationType.advective)), field_name) - # # Add linearisations of new terms if required + # Add linearisations of new terms if required if (t.has_label(linearisation)): u_trial = TrialFunctions(W)[u_idx] u_trial_vert = k*inner(u_trial, k) @@ -398,7 +398,7 @@ def split_hv_advective_form(equation, field_name): linear_hori_term = horizontal(transport(transporting_velocity(test*dot(u_trial_horiz, grad(qbar))*dx, u_trial), TransportEquationType.advective)) adv_horiz_term = linearisation(horizontal_adv_term, linear_hori_term) # Add linearisation to div_term - linear_vert_term = horizontal(transport(transporting_velocity(test*dot(u_trial_vert, grad(qbar))*dx, u_trial), TransportEquationType.advective)) + linear_vert_term = horizontal(transport(transporting_velocity(test*dot(u_trial_vert, grad(qbar))*dx, u_trial), TransportEquationType.advective)) adv_vert_term = linearisation(vertical_adv_term, linear_vert_term) else: adv_vert_term = vertical_adv_term @@ -408,9 +408,7 @@ def split_hv_advective_form(equation, field_name): lambda t: t.get(transport) == TransportEquationType.advective and t.get(prognostic) == field_name, map_if_true=drop) - # Add new terms onto residual equation.residual += subject(adv_horiz_term, subj) + subject(adv_vert_term, subj) - return equation diff --git a/gusto/spatial_methods/transport_methods.py b/gusto/spatial_methods/transport_methods.py index fc07fcbd6..8a960bbbd 100644 --- a/gusto/spatial_methods/transport_methods.py +++ b/gusto/spatial_methods/transport_methods.py @@ -112,13 +112,13 @@ def replace_form(self, equation): else: horizontal_form = equation.residual.label_map( - lambda t: t.has_label(transport) and t.has_label(horizontal) and t.get(prognostic) == self.variable, + lambda t: t.has_label(transport) and t.has_label(horizontal) and t.get(prognostic) == self.variable, map_if_true=keep, map_if_false=drop - ) + ) vertical_form = equation.residual.label_map( - lambda t: t.has_label(transport) and t.has_label(vertical) and t.get(prognostic) == self.variable, + lambda t: t.has_label(transport) and t.has_label(vertical) and t.get(prognostic) == self.variable, map_if_true=keep, map_if_false=drop - ) + ) # Replace form horizontal_term = horizontal_form.terms[0] @@ -136,18 +136,16 @@ def replace_form(self, equation): # Check if this is a conservative transport if horizontal_term.has_label(mass_weighted) or vertical_term.has_label(mass_weighted): - raise RuntimeError('Mass weighted transport terms not yet supported for multiple terms') - - # Replace original term with new term - # equation.residual = equation.residual.label_map( - # lambda t: t.has_label(transport) and t.has_label(horizontal) and t.get(prognostic) == self.variable, - # map_if_true=lambda t: new_horizontal_term) + raise RuntimeError('Mass weighted transport terms not yet supported for multiple terms') + # Replace original terms with new terms equation.residual = equation.residual.label_map( - lambda t: t.has_label(transport) and t.has_label(vertical) and t.get(prognostic) == self.variable, - map_if_true=lambda t: new_vertical_term) - + lambda t: t.has_label(transport) and t.has_label(horizontal) and t.get(prognostic) == self.variable, + map_if_true=lambda _: new_horizontal_term) + equation.residual = equation.residual.label_map( + lambda t: t.has_label(transport) and t.has_label(vertical) and t.get(prognostic) == self.variable, + map_if_true=lambda _: new_vertical_term) # ---------------------------------------------------------------------------- # @@ -314,7 +312,8 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, class Split_DGUpwind(TransportMethod): """ - The Discontinuous Galerkin Upwind transport scheme. + The Discontinuous Galerkin Upwind transport scheme applied separately in the + horizontal and vertical directions. Discretises the gradient of a field weakly, taking the upwind value of the transported variable at facets. """ @@ -342,39 +341,35 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, # Determine appropriate form to use # -------------------------------------------------------------------- # # first check for 1d mesh and scalar velocity space + if equation.domain.on_sphere: + raise NotImplementedError('Split hv Upwind transport scheme has not been ' + + 'implemented for spherical geometry') 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) - elif self.transport_equation_type == TransportEquationType.conservative: - form = upwind_continuity_form_1d(self.domain, self.test, - self.field, - ibp=ibp, outflow=outflow) + raise ValueError('You cannot do horizontal and vertical splitting in 1D') else: if self.transport_equation_type == TransportEquationType.advective: form_h, form_v = split_upwind_advection_form(self.domain, self.test, - self.field, - ibp=ibp, outflow=outflow) + self.field, + ibp=ibp, outflow=outflow) else: - raise NotImplementedError('Upwind transport scheme has not been ' + raise NotImplementedError('Split hv Upwind transport scheme has not been ' + 'implemented for this transport equation type') self.form_v = form_v self.form_h = form_h - # ---------------------------------------------------------------------------- # # Forms for DG Upwind transport # ---------------------------------------------------------------------------- # def split_upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): u""" - The form corresponding to the DG upwind advective transport operator. - This discretises (u.∇)q, for transporting velocity u and transported + The forms corresponding to the DG upwind advective transport operator in + the horizontal and vertical directions. + This discretises u_h.(∇_h)q and w dq/dz, for transporting velocity u and transported variable q. An upwind discretisation is used for the facet terms when the form is integrated by parts. Args: @@ -397,7 +392,7 @@ def split_upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outf if outflow and ibp == IntegrateByParts.NEVER: raise ValueError("outflow is True and ibp is None are incompatible options") Vu = domain.spaces("HDiv") - k=domain.k + k = domain.k dS_ = (dS_v + dS_h) if Vu.extruded else dS ubar = Function(Vu) ubar_v = k*inner(ubar, k) @@ -421,11 +416,11 @@ def split_upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outf L_v += dot(jump(test), (un_v('+')*q('+') - un_v('-')*q('-')))*dS_ if ibp == IntegrateByParts.TWICE: - L_h -= (inner(test('+'), dot(ubar_h('+'), n('+'))*q('+')) - + inner(test('-'), dot(ubar_h('-'), n('-'))*q('-')))*dS_ + L_h -= (inner(test('+'), dot(ubar_h('+'), n('+')) * q('+')) + + inner(test('-'), dot(ubar_h('-'), n('-')) * q('-'))) * dS_ - L_v -= (inner(test('+'), dot(ubar_v('+'), n('+'))*q('+')) - + inner(test('-'), dot(ubar_v('-'), n('-'))*q('-')))*dS_ + L_v -= (inner(test('+'), dot(ubar_v('+'), n('+')) * q('+')) + + inner(test('-'), dot(ubar_v('-'), n('-')) * q('-'))) * dS_ if outflow: n = FacetNormal(domain.mesh) diff --git a/integration-tests/transport/test_split_dg_transport.py b/integration-tests/transport/test_split_dg_transport.py new file mode 100644 index 000000000..8293eb64a --- /dev/null +++ b/integration-tests/transport/test_split_dg_transport.py @@ -0,0 +1,64 @@ +""" +Tests the Split horizontal and vertical DG upwind transport scheme for +advective form transport equation. This tests that the +field is transported to the correct place. +""" + +from firedrake import norm, VectorFunctionSpace, as_vector +from gusto import * + + +def run(timestepper, tmax, f_end): + timestepper.run(0, tmax) + return norm(timestepper.fields("f") - f_end) / norm(f_end) + + +def test_split_dg_transport_scalar(tmpdir, tracer_setup): + setup = tracer_setup(tmpdir, "slice") + domain = setup.domain + V = domain.spaces("DG") + + eqn = AdvectionEquation(domain, V, "f") + eqn = split_hv_advective_form(eqn, "f") + + transport_method = Split_DGUpwind(eqn, "f") + transport_scheme = SSPRK3(domain) + + time_varying_velocity = False + timestepper = PrescribedTransport( + eqn, transport_scheme, setup.io, time_varying_velocity, transport_method + ) + + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) + + error = run(timestepper, setup.tmax, setup.f_end) + assert error < setup.tol, \ + 'The transport error is greater than the permitted tolerance' + + +def test_split_dg_transport_vector(tmpdir, tracer_setup): + setup = tracer_setup(tmpdir, "slice") + domain = setup.domain + gdim = domain.mesh.geometric_dimension() + f_init = as_vector([setup.f_init]*gdim) + V = VectorFunctionSpace(domain.mesh, "DG", 1) + eqn = AdvectionEquation(domain, V, "f") + eqn = split_hv_advective_form(eqn, "f") + + transport_scheme = SSPRK3(domain) + transport_method = Split_DGUpwind(eqn, "f") + + time_varying_velocity = False + timestepper = PrescribedTransport( + eqn, transport_scheme, setup.io, time_varying_velocity, transport_method + ) + + # Initial conditions + timestepper.fields("f").interpolate(f_init) + timestepper.fields("u").project(setup.uexpr) + f_end = as_vector([setup.f_end]*gdim) + error = run(timestepper, setup.tmax, f_end) + assert error < setup.tol, \ + 'The transport error is greater than the permitted tolerance' From 62739207d65a2dc22ec084ee92e3dfbdb44094d3 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 27 Mar 2025 11:58:18 +0000 Subject: [PATCH 05/15] Changes in response to TB review --- gusto/core/labels.py | 4 +- gusto/equations/common_forms.py | 47 ++++++++-- gusto/spatial_methods/spatial_methods.py | 24 ++--- gusto/spatial_methods/transport_methods.py | 90 ++++++++++--------- .../transport/test_split_dg_transport.py | 4 +- 5 files changed, 107 insertions(+), 62 deletions(-) diff --git a/gusto/core/labels.py b/gusto/core/labels.py index c325e0aa0..995f03b70 100644 --- a/gusto/core/labels.py +++ b/gusto/core/labels.py @@ -92,8 +92,8 @@ def __call__(self, target, value=None): # ---------------------------------------------------------------------------- # implicit = Label("implicit") explicit = Label("explicit") -horizontal = Label("horizontal") -vertical = Label("vertical") +horizontal_transport = Label("horizontal_transport") +vertical_transport = Label("vertical_transport") source_label = Label("source_label") 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) diff --git a/gusto/equations/common_forms.py b/gusto/equations/common_forms.py index 8ea031fde..7234cc984 100644 --- a/gusto/equations/common_forms.py +++ b/gusto/equations/common_forms.py @@ -7,7 +7,8 @@ from firedrake.fml import subject, drop from gusto.core.configuration import TransportEquationType from gusto.core.labels import (transport, transporting_velocity, diffusion, - prognostic, linearisation, horizontal, vertical) + prognostic, linearisation, horizontal_transport, + vertical_transport) __all__ = ["advection_form", "advection_form_1d", "continuity_form", "continuity_form_1d", "vector_invariant_form", @@ -385,8 +386,28 @@ def split_hv_advective_form(equation, field_name): u_vertical = k*inner(uadv, k) u_horizontal = uadv - u_vertical - vertical_adv_term = prognostic(vertical(transport(transporting_velocity(inner(test, dot(u_vertical, grad(q)))*dx, uadv), TransportEquationType.advective)), field_name) - horizontal_adv_term = prognostic(horizontal(transport(transporting_velocity(inner(test, dot(u_horizontal, grad(q)))*dx, uadv), TransportEquationType.advective)), field_name) + vertical_adv_term = prognostic( + vertical_transport( + transport( + transporting_velocity( + inner(test, dot(u_vertical, grad(q))) * dx, uadv + ), + TransportEquationType.advective + ) + ), + field_name + ) + horizontal_adv_term = prognostic( + horizontal_transport( + transport( + transporting_velocity( + inner(test, dot(u_horizontal, grad(q))) * dx, uadv + ), + TransportEquationType.advective + ) + ), + field_name + ) # Add linearisations of new terms if required if (t.has_label(linearisation)): @@ -395,10 +416,26 @@ def split_hv_advective_form(equation, field_name): u_trial_horiz = u_trial - u_trial_vert qbar = split(equation.X_ref)[idx] # Add linearisation to adv_term - linear_hori_term = horizontal(transport(transporting_velocity(test*dot(u_trial_horiz, grad(qbar))*dx, u_trial), TransportEquationType.advective)) + linear_hori_term = horizontal_transport( + transport( + transporting_velocity( + test * dot(u_trial_horiz, grad(qbar)) * dx, + u_trial + ), + TransportEquationType.advective + ) + ) adv_horiz_term = linearisation(horizontal_adv_term, linear_hori_term) # Add linearisation to div_term - linear_vert_term = horizontal(transport(transporting_velocity(test*dot(u_trial_vert, grad(qbar))*dx, u_trial), TransportEquationType.advective)) + linear_vert_term = vertical_transport( + transport( + transporting_velocity( + test * dot(u_trial_vert, grad(qbar)) * dx, + u_trial + ), + TransportEquationType.advective + ) + ) adv_vert_term = linearisation(vertical_adv_term, linear_vert_term) else: adv_vert_term = vertical_adv_term diff --git a/gusto/spatial_methods/spatial_methods.py b/gusto/spatial_methods/spatial_methods.py index ba1cb7a00..009681052 100644 --- a/gusto/spatial_methods/spatial_methods.py +++ b/gusto/spatial_methods/spatial_methods.py @@ -15,19 +15,19 @@ class SpatialMethod(object): The base object for describing a spatial discretisation of some term. """ - def __init__(self, equation, variable, term_label): + def __init__(self, equation, variable, term_labels): """ Args: equation (:class:`PrognosticEquation`): the equation, which includes the original type of this term. variable (str): name of the variable to set the method for - term_label (:class:`Label`): the label specifying which type of term - to be discretised. + term_labels (list of :class:`Label`, optional): the label specifying + which type of term to be discretised. """ self.equation = equation self.variable = variable self.domain = self.equation.domain - self.term_label = term_label + self.term_labels = term_labels if hasattr(equation, "field_names"): # Equation with multiple prognostic variables @@ -38,14 +38,16 @@ def __init__(self, equation, variable, term_label): self.field = equation.X self.test = equation.test - # Find the original term to be used - self.original_form = equation.residual.label_map( - lambda t: t.has_label(term_label) and t.get(prognostic) == variable, - map_if_true=keep, map_if_false=drop) + # Loop through terms + for term_label in term_labels: + # Find the original term to be used + self.original_form = equation.residual.label_map( + lambda t: t.has_label(term_label) and t.get(prognostic) == variable, + map_if_true=keep, map_if_false=drop) - num_terms = len(self.original_form.terms) - assert num_terms >= 1, f'Unable to find {term_label.label} term ' \ - + f'for {variable}. {num_terms} found' + num_terms = len(self.original_form.terms) + assert num_terms == 1, f'Unable to find {term_label.label} term ' \ + + f'for {variable}. {num_terms} found' def replace_form(self, equation): """ diff --git a/gusto/spatial_methods/transport_methods.py b/gusto/spatial_methods/transport_methods.py index f92f8811f..566d8af70 100644 --- a/gusto/spatial_methods/transport_methods.py +++ b/gusto/spatial_methods/transport_methods.py @@ -10,12 +10,12 @@ from gusto.core.configuration import IntegrateByParts, TransportEquationType from gusto.core.labels import ( prognostic, transport, transporting_velocity, ibp_label, mass_weighted, - all_but_last, horizontal, vertical, explicit + all_but_last, horizontal_transport, vertical_transport, explicit ) from gusto.core.logging import logger from gusto.spatial_methods.spatial_methods import SpatialMethod -__all__ = ["DefaultTransport", "DGUpwind", "Split_DGUpwind"] +__all__ = ["DefaultTransport", "DGUpwind", "SplitDGUpwind"] # ---------------------------------------------------------------------------- # @@ -26,16 +26,18 @@ class TransportMethod(SpatialMethod): The base object for describing a transport scheme. """ - def __init__(self, equation, variable): + def __init__(self, equation, variable, term_labels=[transport]): """ Args: equation (:class:`PrognosticEquation`): the equation, which includes a transport term. variable (str): name of the variable to set the transport scheme for + term_labels (list of :class:`Label`, optional): the label specifying + which type of term to be discretised. Defaults to [transport]. """ # Inherited init method extracts original term to be replaced - super().__init__(equation, variable, transport) + super().__init__(equation, variable, term_labels) # If this is term has a mass_weighted label, then we need to # use the tracer_conservative version of the transport method. @@ -112,41 +114,44 @@ def replace_form(self, equation): else: horizontal_form = equation.residual.label_map( - lambda t: t.has_label(transport) and t.has_label(horizontal) and t.get(prognostic) == self.variable, + lambda t: t.has_label(transport) and t.has_label(horizontal_transport) and t.get(prognostic) == self.variable, map_if_true=keep, map_if_false=drop ) vertical_form = equation.residual.label_map( - lambda t: t.has_label(transport) and t.has_label(vertical) and t.get(prognostic) == self.variable, + lambda t: t.has_label(transport) and t.has_label(vertical_transport) and t.get(prognostic) == self.variable, map_if_true=keep, map_if_false=drop ) - - # Replace form - horizontal_term = horizontal_form.terms[0] - vertical_term = vertical_form.terms[0] - - # Update transporting velocity - new_horizontal_transporting_velocity = self.form_h.terms[0].get(transporting_velocity) - new_vertical_transporting_velocity = self.form_v.terms[0].get(transporting_velocity) - horizontal_term = transporting_velocity.update_value(horizontal_term, new_horizontal_transporting_velocity) - vertical_term = transporting_velocity.update_value(vertical_term, new_vertical_transporting_velocity) - - # Create new term - new_horizontal_term = Term(self.form_h.form, horizontal_term.labels) - new_vertical_term = Term(self.form_v.form, vertical_term.labels) - - # Check if this is a conservative transport - if horizontal_term.has_label(mass_weighted) or vertical_term.has_label(mass_weighted): - raise RuntimeError('Mass weighted transport terms not yet supported for multiple terms') - - # Replace original terms with new terms - equation.residual = equation.residual.label_map( - lambda t: t.has_label(transport) and t.has_label(horizontal) and t.get(prognostic) == self.variable, - map_if_true=lambda _: new_horizontal_term) - - equation.residual = equation.residual.label_map( - lambda t: t.has_label(transport) and t.has_label(vertical) and t.get(prognostic) == self.variable, - map_if_true=lambda _: new_vertical_term) - + if len(horizontal_form.terms) == 1 and len(vertical_form.terms) == 1: + + # Replace forms + horizontal_term = horizontal_form.terms[0] + vertical_term = vertical_form.terms[0] + + # Update transporting velocity + new_horizontal_transporting_velocity = self.form_h.terms[0].get(transporting_velocity) + new_vertical_transporting_velocity = self.form_v.terms[0].get(transporting_velocity) + horizontal_term = transporting_velocity.update_value(horizontal_term, new_horizontal_transporting_velocity) + vertical_term = transporting_velocity.update_value(vertical_term, new_vertical_transporting_velocity) + + # Create new terms + new_horizontal_term = Term(self.form_h.form, horizontal_term.labels) + new_vertical_term = Term(self.form_v.form, vertical_term.labels) + + # Check if this is a conservative transport + if horizontal_term.has_label(mass_weighted) or vertical_term.has_label(mass_weighted): + raise NotImplementedError('Mass weighted transport terms not yet supported for multiple terms') + + # Replace original terms with new terms + equation.residual = equation.residual.label_map( + lambda t: t.has_label(transport) and t.has_label(horizontal_transport) and t.get(prognostic) == self.variable, + map_if_true=lambda _: new_horizontal_term) + + equation.residual = equation.residual.label_map( + lambda t: t.has_label(transport) and t.has_label(vertical_transport) and t.get(prognostic) == self.variable, + map_if_true=lambda _: new_vertical_term) + else: + raise RuntimeError('Found multiple transport terms for the same ' + 'variable in the equation where there should only be one') # ---------------------------------------------------------------------------- # # TransportMethod for using underlying default transport form @@ -310,7 +315,7 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, self.form = form -class Split_DGUpwind(TransportMethod): +class SplitDGUpwind(TransportMethod): """ The Discontinuous Galerkin Upwind transport scheme applied separately in the horizontal and vertical directions. @@ -332,7 +337,7 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, boundaries, through exterior facet terms. Defaults to False. """ - super().__init__(equation, variable) + super().__init__(equation, variable, [horizontal_transport, vertical_transport]) self.ibp = ibp self.vector_manifold_correction = vector_manifold_correction self.outflow = outflow @@ -393,17 +398,18 @@ def split_upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outf raise ValueError("outflow is True and ibp is None are incompatible options") Vu = domain.spaces("HDiv") k = domain.k - dS_ = (dS_v + dS_h) if Vu.extruded else dS + quad = domain.max_quad_degree + dS_ = (dS_v(degree=quad) + dS_h(degree=quad)) if Vu.extruded else dS ubar = Function(Vu) ubar_v = k*inner(ubar, k) ubar_h = ubar - ubar_v if ibp == IntegrateByParts.ONCE: - L_h = -inner(div(outer(test, ubar_h)), q)*dx - L_v = -inner(div(outer(test, ubar_v)), q)*dx + L_h = -inner(div(outer(test, ubar_h)), q)*dx(degree=quad) + L_v = -inner(div(outer(test, ubar_v)), q)*dx(degree=quad) else: - L_h = inner(outer(test, ubar_h), grad(q))*dx - L_v = inner(outer(test, ubar_v), grad(q))*dx + L_h = inner(outer(test, ubar_h), grad(q))*dx(degree=quad) + L_v = inner(outer(test, ubar_v), grad(q))*dx(degree=quad) if ibp != IntegrateByParts.NEVER: n = FacetNormal(domain.mesh) @@ -432,7 +438,7 @@ def split_upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outf form_h = transporting_velocity(L_h, ubar) form_v = transporting_velocity(L_v, ubar) - labelled_form_h = ibp_label(transport(explicit(form_h), TransportEquationType.advective), ibp) + labelled_form_h = ibp_label(transport(form_h, TransportEquationType.advective), ibp) labelled_form_v = ibp_label(transport(form_v, TransportEquationType.advective), ibp) return labelled_form_h, labelled_form_v diff --git a/integration-tests/transport/test_split_dg_transport.py b/integration-tests/transport/test_split_dg_transport.py index 8293eb64a..3099a10d8 100644 --- a/integration-tests/transport/test_split_dg_transport.py +++ b/integration-tests/transport/test_split_dg_transport.py @@ -21,7 +21,7 @@ def test_split_dg_transport_scalar(tmpdir, tracer_setup): eqn = AdvectionEquation(domain, V, "f") eqn = split_hv_advective_form(eqn, "f") - transport_method = Split_DGUpwind(eqn, "f") + transport_method = SplitDGUpwind(eqn, "f") transport_scheme = SSPRK3(domain) time_varying_velocity = False @@ -48,7 +48,7 @@ def test_split_dg_transport_vector(tmpdir, tracer_setup): eqn = split_hv_advective_form(eqn, "f") transport_scheme = SSPRK3(domain) - transport_method = Split_DGUpwind(eqn, "f") + transport_method = SplitDGUpwind(eqn, "f") time_varying_velocity = False timestepper = PrescribedTransport( From b22afa8d9495074cc9b0e2a05c4d5569aeb903c9 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 27 Mar 2025 11:59:51 +0000 Subject: [PATCH 06/15] Lint fix --- gusto/spatial_methods/transport_methods.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/gusto/spatial_methods/transport_methods.py b/gusto/spatial_methods/transport_methods.py index 566d8af70..11245fff5 100644 --- a/gusto/spatial_methods/transport_methods.py +++ b/gusto/spatial_methods/transport_methods.py @@ -10,7 +10,7 @@ from gusto.core.configuration import IntegrateByParts, TransportEquationType from gusto.core.labels import ( prognostic, transport, transporting_velocity, ibp_label, mass_weighted, - all_but_last, horizontal_transport, vertical_transport, explicit + all_but_last, horizontal_transport, vertical_transport ) from gusto.core.logging import logger from gusto.spatial_methods.spatial_methods import SpatialMethod @@ -151,7 +151,8 @@ def replace_form(self, equation): map_if_true=lambda _: new_vertical_term) else: raise RuntimeError('Found multiple transport terms for the same ' - 'variable in the equation where there should only be one') + 'variable in the equation where there should only be one') + # ---------------------------------------------------------------------------- # # TransportMethod for using underlying default transport form From 4d3dcdc4f0362ded7073d2152017aa6c5db186bc Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 27 Mar 2025 12:35:09 +0000 Subject: [PATCH 07/15] Tidy of common forms and fix to spatial methods --- gusto/equations/common_forms.py | 79 +++++++++++++++--------- gusto/spatial_methods/spatial_methods.py | 21 +++---- 2 files changed, 59 insertions(+), 41 deletions(-) diff --git a/gusto/equations/common_forms.py b/gusto/equations/common_forms.py index 7234cc984..e59423894 100644 --- a/gusto/equations/common_forms.py +++ b/gusto/equations/common_forms.py @@ -348,11 +348,53 @@ def tracer_conservative_form(test, q, rho, ubar): return transport(form, TransportEquationType.tracer_conservative) +def split_advection_form(test, q, ubar, ubar_full): + u""" + The form corresponding to the advective transport operator in either horzontal + or vertical directions (dependent on ubar). + + This describes either u_h.(∇)q or w dq/dz, for transporting velocity u and transported q. + + Args: + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ubar (:class:`ufl.Expr`): the transporting velocity in a subset of dimensions. + ubar_full (:class:`ufl.Expr`): the transporting velocity in all dimensions. + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + L = inner(test, dot(ubar, grad(q)))*dx + form = transporting_velocity(L, ubar_full) + + return transport(form, TransportEquationType.advective) + +def split_linear_advection_form(test, qbar, ubar, ubar_full): + """ + The form corresponding to the linearised advective transport operator in + either horzontal or vertical directions (dependent on ubar). + + Args: + test (:class:`TestFunction`): the test function. + qbar (:class:`ufl.Expr`): the variable to be transported. + ubar (:class:`ufl.Expr`): the transporting velocity in a subset of dimensions. + ubar_full (:class:`ufl.Expr`): the transporting velocity in all dimensions. + + Returns: + :class:`LabelledForm`: a labelled transport form. + """ + + L = test*dot(ubar, grad(qbar))*dx + form = transporting_velocity(L, ubar_full) + + return transport(form, TransportEquationType.advective) + def split_hv_advective_form(equation, field_name): u""" Splits advective term into horizontal and vertical terms. - This describes splitting u.∇(q) terms into u.(∇_h)q and w dq/dz, + This describes splitting u.∇(q) terms into u_h.(∇)q and w dq/dz, for transporting velocity u and transported q. Args: equation (:class:`PrognosticEquation`): the model's equation. @@ -385,26 +427,15 @@ def split_hv_advective_form(equation, field_name): # Create new advective and divergence terms u_vertical = k*inner(uadv, k) u_horizontal = uadv - u_vertical - vertical_adv_term = prognostic( vertical_transport( - transport( - transporting_velocity( - inner(test, dot(u_vertical, grad(q))) * dx, uadv - ), - TransportEquationType.advective - ) + split_advection_form(test, q, u_vertical, uadv) ), field_name ) horizontal_adv_term = prognostic( horizontal_transport( - transport( - transporting_velocity( - inner(test, dot(u_horizontal, grad(q))) * dx, uadv - ), - TransportEquationType.advective - ) + split_advection_form(test, q, u_horizontal, uadv) ), field_name ) @@ -415,26 +446,14 @@ def split_hv_advective_form(equation, field_name): u_trial_vert = k*inner(u_trial, k) u_trial_horiz = u_trial - u_trial_vert qbar = split(equation.X_ref)[idx] - # Add linearisation to adv_term + # Add linearisations linear_hori_term = horizontal_transport( - transport( - transporting_velocity( - test * dot(u_trial_horiz, grad(qbar)) * dx, - u_trial - ), - TransportEquationType.advective - ) + split_linear_advection_form(test, qbar, u_trial_horiz, u_trial) ) adv_horiz_term = linearisation(horizontal_adv_term, linear_hori_term) - # Add linearisation to div_term + linear_vert_term = vertical_transport( - transport( - transporting_velocity( - test * dot(u_trial_vert, grad(qbar)) * dx, - u_trial - ), - TransportEquationType.advective - ) + split_linear_advection_form(test, qbar, u_trial_vert, u_trial) ) adv_vert_term = linearisation(vertical_adv_term, linear_vert_term) else: diff --git a/gusto/spatial_methods/spatial_methods.py b/gusto/spatial_methods/spatial_methods.py index 009681052..875cf8d42 100644 --- a/gusto/spatial_methods/spatial_methods.py +++ b/gusto/spatial_methods/spatial_methods.py @@ -21,13 +21,14 @@ def __init__(self, equation, variable, term_labels): equation (:class:`PrognosticEquation`): the equation, which includes the original type of this term. variable (str): name of the variable to set the method for - term_labels (list of :class:`Label`, optional): the label specifying + term_labels (list of :class:`Label`): list of the labels specifying which type of term to be discretised. """ self.equation = equation self.variable = variable self.domain = self.equation.domain - self.term_labels = term_labels + # Most schemes have only one term label + self.term_label = term_label[0] if hasattr(equation, "field_names"): # Equation with multiple prognostic variables @@ -38,16 +39,14 @@ def __init__(self, equation, variable, term_labels): self.field = equation.X self.test = equation.test - # Loop through terms - for term_label in term_labels: - # Find the original term to be used - self.original_form = equation.residual.label_map( - lambda t: t.has_label(term_label) and t.get(prognostic) == variable, - map_if_true=keep, map_if_false=drop) + # Find the original term to be used (for first term label) + self.original_form = equation.residual.label_map( + lambda t: t.has_label(self.term_label) and t.get(prognostic) == variable, + map_if_true=keep, map_if_false=drop) - num_terms = len(self.original_form.terms) - assert num_terms == 1, f'Unable to find {term_label.label} term ' \ - + f'for {variable}. {num_terms} found' + num_terms_per_label = len(self.original_form.terms)/len(term_labels) + assert num_terms_per_label == 1, f'Unable to find {term_label.label} term ' \ + + f'for {variable}. {num_terms} found' def replace_form(self, equation): """ From df562adf1ef3372127a93393a7dc627074757f42 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 27 Mar 2025 12:45:24 +0000 Subject: [PATCH 08/15] More fixes to code.. --- gusto/equations/common_forms.py | 2 ++ gusto/spatial_methods/spatial_methods.py | 13 +++++++++---- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/gusto/equations/common_forms.py b/gusto/equations/common_forms.py index e59423894..a96c84b8d 100644 --- a/gusto/equations/common_forms.py +++ b/gusto/equations/common_forms.py @@ -348,6 +348,7 @@ def tracer_conservative_form(test, q, rho, ubar): return transport(form, TransportEquationType.tracer_conservative) + def split_advection_form(test, q, ubar, ubar_full): u""" The form corresponding to the advective transport operator in either horzontal @@ -370,6 +371,7 @@ def split_advection_form(test, q, ubar, ubar_full): return transport(form, TransportEquationType.advective) + def split_linear_advection_form(test, qbar, ubar, ubar_full): """ The form corresponding to the linearised advective transport operator in diff --git a/gusto/spatial_methods/spatial_methods.py b/gusto/spatial_methods/spatial_methods.py index 875cf8d42..d73387010 100644 --- a/gusto/spatial_methods/spatial_methods.py +++ b/gusto/spatial_methods/spatial_methods.py @@ -28,7 +28,7 @@ def __init__(self, equation, variable, term_labels): self.variable = variable self.domain = self.equation.domain # Most schemes have only one term label - self.term_label = term_label[0] + self.term_label = term_labels[0] if hasattr(equation, "field_names"): # Equation with multiple prognostic variables @@ -44,9 +44,14 @@ def __init__(self, equation, variable, term_labels): lambda t: t.has_label(self.term_label) and t.get(prognostic) == variable, map_if_true=keep, map_if_false=drop) - num_terms_per_label = len(self.original_form.terms)/len(term_labels) - assert num_terms_per_label == 1, f'Unable to find {term_label.label} term ' \ - + f'for {variable}. {num_terms} found' + num_terms_per_label = len(self.original_form.terms) // len(term_labels) + assert len(self.original_form.terms) % len(term_labels) == 0, ( + "The terms do not divide evenly into labels." + ) + assert num_terms_per_label == 1, ( + f"Unable to find terms {[term.label for term in term_labels]} for " + f"{variable}. {num_terms_per_label} terms per expected term found" + ) def replace_form(self, equation): """ From f58d7b725aed77bea4c0d963d11f5bdd76afef89 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 27 Mar 2025 13:17:19 +0000 Subject: [PATCH 09/15] Spherical test + fixes added --- gusto/spatial_methods/diffusion_methods.py | 2 +- gusto/spatial_methods/transport_methods.py | 4 ---- .../transport/test_split_dg_transport.py | 12 ++++++------ 3 files changed, 7 insertions(+), 11 deletions(-) diff --git a/gusto/spatial_methods/diffusion_methods.py b/gusto/spatial_methods/diffusion_methods.py index bc33da4df..7d2ab21db 100644 --- a/gusto/spatial_methods/diffusion_methods.py +++ b/gusto/spatial_methods/diffusion_methods.py @@ -22,7 +22,7 @@ def __init__(self, equation, variable): """ # Inherited init method extracts original term to be replaced - super().__init__(equation, variable, diffusion) + super().__init__(equation, variable, [diffusion]) def interior_penalty_diffusion_form(domain, test, q, parameters): diff --git a/gusto/spatial_methods/transport_methods.py b/gusto/spatial_methods/transport_methods.py index 11245fff5..ce7cabbb5 100644 --- a/gusto/spatial_methods/transport_methods.py +++ b/gusto/spatial_methods/transport_methods.py @@ -347,13 +347,9 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, # Determine appropriate form to use # -------------------------------------------------------------------- # # first check for 1d mesh and scalar velocity space - if equation.domain.on_sphere: - raise NotImplementedError('Split hv Upwind transport scheme has not been ' - + 'implemented for spherical geometry') if equation.domain.mesh.topological_dimension() == 1 and len(equation.domain.spaces("HDiv").shape) == 0: assert not vector_manifold_correction raise ValueError('You cannot do horizontal and vertical splitting in 1D') - else: if self.transport_equation_type == TransportEquationType.advective: diff --git a/integration-tests/transport/test_split_dg_transport.py b/integration-tests/transport/test_split_dg_transport.py index 3099a10d8..0e686bfa3 100644 --- a/integration-tests/transport/test_split_dg_transport.py +++ b/integration-tests/transport/test_split_dg_transport.py @@ -12,9 +12,9 @@ def run(timestepper, tmax, f_end): timestepper.run(0, tmax) return norm(timestepper.fields("f") - f_end) / norm(f_end) - -def test_split_dg_transport_scalar(tmpdir, tracer_setup): - setup = tracer_setup(tmpdir, "slice") +@pytest.mark.parametrize("geometry", ["slice", "sphere"]) +def test_split_dg_transport_scalar(tmpdir, geometry, tracer_setup): + setup = tracer_setup(tmpdir, geometry) domain = setup.domain V = domain.spaces("DG") @@ -37,9 +37,9 @@ def test_split_dg_transport_scalar(tmpdir, tracer_setup): assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' - -def test_split_dg_transport_vector(tmpdir, tracer_setup): - setup = tracer_setup(tmpdir, "slice") +@pytest.mark.parametrize("geometry", ["slice", "sphere"]) +def test_split_dg_transport_vector(tmpdir, geometry, tracer_setup): + setup = tracer_setup(tmpdir, geometry) domain = setup.domain gdim = domain.mesh.geometric_dimension() f_init = as_vector([setup.f_init]*gdim) From 9f10f4dca6ec640e1c0cba3599630d665d58e67f Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 27 Mar 2025 13:17:56 +0000 Subject: [PATCH 10/15] Lint fix --- integration-tests/transport/test_split_dg_transport.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/integration-tests/transport/test_split_dg_transport.py b/integration-tests/transport/test_split_dg_transport.py index 0e686bfa3..37baac61a 100644 --- a/integration-tests/transport/test_split_dg_transport.py +++ b/integration-tests/transport/test_split_dg_transport.py @@ -12,6 +12,7 @@ def run(timestepper, tmax, f_end): timestepper.run(0, tmax) return norm(timestepper.fields("f") - f_end) / norm(f_end) + @pytest.mark.parametrize("geometry", ["slice", "sphere"]) def test_split_dg_transport_scalar(tmpdir, geometry, tracer_setup): setup = tracer_setup(tmpdir, geometry) @@ -37,6 +38,7 @@ def test_split_dg_transport_scalar(tmpdir, geometry, tracer_setup): assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' + @pytest.mark.parametrize("geometry", ["slice", "sphere"]) def test_split_dg_transport_vector(tmpdir, geometry, tracer_setup): setup = tracer_setup(tmpdir, geometry) From e045287afad4503a10a76896c0970da704479dfe Mon Sep 17 00:00:00 2001 From: atb1995 Date: Mon, 7 Apr 2025 15:14:37 +0100 Subject: [PATCH 11/15] pytest fix --- integration-tests/transport/test_split_dg_transport.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/integration-tests/transport/test_split_dg_transport.py b/integration-tests/transport/test_split_dg_transport.py index 37baac61a..d6864dbed 100644 --- a/integration-tests/transport/test_split_dg_transport.py +++ b/integration-tests/transport/test_split_dg_transport.py @@ -6,7 +6,7 @@ from firedrake import norm, VectorFunctionSpace, as_vector from gusto import * - +import pytest def run(timestepper, tmax, f_end): timestepper.run(0, tmax) From c362d337afd8438d734597f0f3e714e2d6731631 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Mon, 7 Apr 2025 15:21:45 +0100 Subject: [PATCH 12/15] lint fix --- integration-tests/transport/test_split_dg_transport.py | 1 + 1 file changed, 1 insertion(+) diff --git a/integration-tests/transport/test_split_dg_transport.py b/integration-tests/transport/test_split_dg_transport.py index d6864dbed..b1cc43d2c 100644 --- a/integration-tests/transport/test_split_dg_transport.py +++ b/integration-tests/transport/test_split_dg_transport.py @@ -8,6 +8,7 @@ from gusto import * import pytest + def run(timestepper, tmax, f_end): timestepper.run(0, tmax) return norm(timestepper.fields("f") - f_end) / norm(f_end) From 9f15d8d2da2ea08ea08e9baace7281c1574f9e7f Mon Sep 17 00:00:00 2001 From: atb1995 Date: Mon, 7 Apr 2025 15:59:08 +0100 Subject: [PATCH 13/15] Fix error for check --- gusto/spatial_methods/spatial_methods.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/gusto/spatial_methods/spatial_methods.py b/gusto/spatial_methods/spatial_methods.py index d73387010..be216227a 100644 --- a/gusto/spatial_methods/spatial_methods.py +++ b/gusto/spatial_methods/spatial_methods.py @@ -4,7 +4,7 @@ """ from firedrake import split -from firedrake.fml import Term, keep, drop +from firedrake.fml import Term, keep, drop, all_terms from gusto.core.labels import prognostic __all__ = ['SpatialMethod'] @@ -39,10 +39,17 @@ def __init__(self, equation, variable, term_labels): self.field = equation.X self.test = equation.test - # Find the original term to be used (for first term label) + # Find the original terms to be used self.original_form = equation.residual.label_map( - lambda t: t.has_label(self.term_label) and t.get(prognostic) == variable, - map_if_true=keep, map_if_false=drop) + all_terms, + map_if_true=drop + ) + for term in term_labels: + self.original_form += equation.residual.label_map( + lambda t: t.has_label(term) and t.get(prognostic) == variable, + map_if_true=keep, + map_if_false=drop + ) num_terms_per_label = len(self.original_form.terms) // len(term_labels) assert len(self.original_form.terms) % len(term_labels) == 0, ( From 39e74f1e9ff345242d9ec46ca4ac5656a667f791 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Mon, 7 Apr 2025 17:02:52 +0100 Subject: [PATCH 14/15] Source term fix --- gusto/time_discretisation/sdc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/time_discretisation/sdc.py b/gusto/time_discretisation/sdc.py index 3e634f17a..7c419cb21 100644 --- a/gusto/time_discretisation/sdc.py +++ b/gusto/time_discretisation/sdc.py @@ -338,7 +338,7 @@ def res(self, m): r_source_k = self.residual.label_map( lambda t: t.has_label(source_label), - map_if_true=replace_subject(self.source_Ukp1[i+1], old_idx=self.idx), + map_if_true=replace_subject(self.source_Uk[i+1], old_idx=self.idx), map_if_false=drop) r_source_k = r_source_k.label_map( all_terms, From 0e669f25e5790a830935c2f580d336937d6e1d63 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Tue, 8 Apr 2025 14:22:29 +0100 Subject: [PATCH 15/15] Changes in response to review --- gusto/spatial_methods/diffusion_methods.py | 2 +- gusto/spatial_methods/spatial_methods.py | 54 +++++++++++-------- gusto/spatial_methods/transport_methods.py | 13 +++-- .../transport/test_split_dg_transport.py | 13 ++--- 4 files changed, 46 insertions(+), 36 deletions(-) diff --git a/gusto/spatial_methods/diffusion_methods.py b/gusto/spatial_methods/diffusion_methods.py index 7d2ab21db..bc33da4df 100644 --- a/gusto/spatial_methods/diffusion_methods.py +++ b/gusto/spatial_methods/diffusion_methods.py @@ -22,7 +22,7 @@ def __init__(self, equation, variable): """ # Inherited init method extracts original term to be replaced - super().__init__(equation, variable, [diffusion]) + super().__init__(equation, variable, diffusion) def interior_penalty_diffusion_form(domain, test, q, parameters): diff --git a/gusto/spatial_methods/spatial_methods.py b/gusto/spatial_methods/spatial_methods.py index be216227a..2fff98728 100644 --- a/gusto/spatial_methods/spatial_methods.py +++ b/gusto/spatial_methods/spatial_methods.py @@ -15,20 +15,19 @@ class SpatialMethod(object): The base object for describing a spatial discretisation of some term. """ - def __init__(self, equation, variable, term_labels): + def __init__(self, equation, variable, *term_labels): """ Args: equation (:class:`PrognosticEquation`): the equation, which includes the original type of this term. variable (str): name of the variable to set the method for - term_labels (list of :class:`Label`): list of the labels specifying - which type of term to be discretised. + term_labels (:class:`Label`): One or more labels specifying which type(s) + of terms should be discretized. """ self.equation = equation self.variable = variable self.domain = self.equation.domain - # Most schemes have only one term label - self.term_label = term_labels[0] + self.term_labels = list(term_labels) if hasattr(equation, "field_names"): # Equation with multiple prognostic variables @@ -39,26 +38,37 @@ def __init__(self, equation, variable, term_labels): self.field = equation.X self.test = equation.test - # Find the original terms to be used - self.original_form = equation.residual.label_map( - all_terms, - map_if_true=drop - ) - for term in term_labels: - self.original_form += equation.residual.label_map( - lambda t: t.has_label(term) and t.get(prognostic) == variable, + if (len(self.term_labels) == 1): + # Most cases only have one term to be replaced + self.term_label = self.term_labels[0] + self.original_form = equation.residual.label_map( + lambda t: t.has_label(self.term_label) and t.get(prognostic) == variable, map_if_true=keep, map_if_false=drop ) - - num_terms_per_label = len(self.original_form.terms) // len(term_labels) - assert len(self.original_form.terms) % len(term_labels) == 0, ( - "The terms do not divide evenly into labels." - ) - assert num_terms_per_label == 1, ( - f"Unable to find terms {[term.label for term in term_labels]} for " - f"{variable}. {num_terms_per_label} terms per expected term found" - ) + # Check that the original form has the correct number of terms + num_terms = len(self.original_form.terms) + assert num_terms == 1, f'Unable to find {self.term_label.label} term ' \ + + f'for {variable}. {num_terms} found' + else: + # Multiple terms to be replaced. Find the original terms to be used + self.term_label = self.term_labels[0] + self.original_form = equation.residual.label_map( + all_terms, + map_if_true=drop + ) + for term in self.term_labels: + original_form = equation.residual.label_map( + lambda t: t.has_label(term) and t.get(prognostic) == variable, + map_if_true=keep, + map_if_false=drop + ) + # Check that the original form has the correct number of terms + num_terms = len(original_form.terms) + assert num_terms == 1, f'Unable to find {term.label} term ' \ + + f'for {variable}. {num_terms} found' + # Add the terms form to original forms + self.original_form += original_form def replace_form(self, equation): """ diff --git a/gusto/spatial_methods/transport_methods.py b/gusto/spatial_methods/transport_methods.py index ce7cabbb5..edfd79a5c 100644 --- a/gusto/spatial_methods/transport_methods.py +++ b/gusto/spatial_methods/transport_methods.py @@ -26,18 +26,21 @@ class TransportMethod(SpatialMethod): The base object for describing a transport scheme. """ - def __init__(self, equation, variable, term_labels=[transport]): + def __init__(self, equation, variable, *term_labels): """ Args: equation (:class:`PrognosticEquation`): the equation, which includes a transport term. variable (str): name of the variable to set the transport scheme for - term_labels (list of :class:`Label`, optional): the label specifying - which type of term to be discretised. Defaults to [transport]. + term_labels (:class:`Label`): One or more labels specifying which type(s) + of terms should be discretized. """ # Inherited init method extracts original term to be replaced - super().__init__(equation, variable, term_labels) + if not term_labels: + super().__init__(equation, variable, transport) + else: + super().__init__(equation, variable, *term_labels) # If this is term has a mass_weighted label, then we need to # use the tracer_conservative version of the transport method. @@ -338,7 +341,7 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, boundaries, through exterior facet terms. Defaults to False. """ - super().__init__(equation, variable, [horizontal_transport, vertical_transport]) + super().__init__(equation, variable, horizontal_transport, vertical_transport) self.ibp = ibp self.vector_manifold_correction = vector_manifold_correction self.outflow = outflow diff --git a/integration-tests/transport/test_split_dg_transport.py b/integration-tests/transport/test_split_dg_transport.py index b1cc43d2c..5fe1030df 100644 --- a/integration-tests/transport/test_split_dg_transport.py +++ b/integration-tests/transport/test_split_dg_transport.py @@ -1,12 +1,11 @@ """ -Tests the Split horizontal and vertical DG upwind transport scheme for +Tests the split horizontal and vertical DG upwind transport scheme for advective form transport equation. This tests that the field is transported to the correct place. """ from firedrake import norm, VectorFunctionSpace, as_vector from gusto import * -import pytest def run(timestepper, tmax, f_end): @@ -14,9 +13,8 @@ def run(timestepper, tmax, f_end): return norm(timestepper.fields("f") - f_end) / norm(f_end) -@pytest.mark.parametrize("geometry", ["slice", "sphere"]) -def test_split_dg_transport_scalar(tmpdir, geometry, tracer_setup): - setup = tracer_setup(tmpdir, geometry) +def test_split_dg_transport_scalar(tmpdir, tracer_setup): + setup = tracer_setup(tmpdir, "slice") domain = setup.domain V = domain.spaces("DG") @@ -40,9 +38,8 @@ def test_split_dg_transport_scalar(tmpdir, geometry, tracer_setup): 'The transport error is greater than the permitted tolerance' -@pytest.mark.parametrize("geometry", ["slice", "sphere"]) -def test_split_dg_transport_vector(tmpdir, geometry, tracer_setup): - setup = tracer_setup(tmpdir, geometry) +def test_split_dg_transport_vector(tmpdir, tracer_setup): + setup = tracer_setup(tmpdir, "slice") domain = setup.domain gdim = domain.mesh.geometric_dimension() f_init = as_vector([setup.f_init]*gdim)