From 1f2ed7ce891d12dca87f415685b6a9a44ce56b49 Mon Sep 17 00:00:00 2001 From: jshipton Date: Mon, 12 Jan 2026 15:09:40 +0000 Subject: [PATCH 01/16] SIQN model class for compressible Euler equations and simple example script --- .../compressible_euler/sk_simple_model.py | 180 ++++++++++++++++++ gusto/__init__.py | 1 + gusto/model/__init__.py | 1 + gusto/model/model.py | 113 +++++++++++ 4 files changed, 295 insertions(+) create mode 100644 examples/compressible_euler/sk_simple_model.py create mode 100644 gusto/model/__init__.py create mode 100644 gusto/model/model.py diff --git a/examples/compressible_euler/sk_simple_model.py b/examples/compressible_euler/sk_simple_model.py new file mode 100644 index 000000000..f0cf0b0d3 --- /dev/null +++ b/examples/compressible_euler/sk_simple_model.py @@ -0,0 +1,180 @@ +""" +This example uses the non-linear compressible Euler equations to solve the +vertical slice gravity wave test case of Skamarock and Klemp, 1994: +``Efficiency and Accuracy of the Klemp-Wilhelmson Time-Splitting Technique'', +MWR. + +The domain is smaller than the "hydrostatic" gravity wave test, so that there +is difference between the hydrostatic and non-hydrostatic solutions. The test +can be run with and without a hydrostatic switch. + +Potential temperature is transported using SUPG, and the degree 1 elements are +used. +""" +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter + +from firedrake import ( + as_vector, SpatialCoordinate, PeriodicIntervalMesh, ExtrudedMesh, exp, sin, + Function, pi +) +from gusto import ( + Perturbation, CompressibleParameters, OutputParameters, + CompressibleEulerEquations, SIQNModel, + compressible_hydrostatic_balance +) + +skamarock_klemp_nonhydrostatic_defaults = { + 'ncolumns': 150, + 'nlayers': 10, + 'dt': 6.0, + 'tmax': 3000., + 'dumpfreq': 250, + 'dirname': 'skamarock_klemp_nonhydrostatic', +} + + +def skamarock_klemp_nonhydrostatic( + ncolumns=skamarock_klemp_nonhydrostatic_defaults['ncolumns'], + nlayers=skamarock_klemp_nonhydrostatic_defaults['nlayers'], + dt=skamarock_klemp_nonhydrostatic_defaults['dt'], + tmax=skamarock_klemp_nonhydrostatic_defaults['tmax'], + dumpfreq=skamarock_klemp_nonhydrostatic_defaults['dumpfreq'], + dirname=skamarock_klemp_nonhydrostatic_defaults['dirname'], +): + + # ------------------------------------------------------------------------ # + # Test case parameters + # ------------------------------------------------------------------------ # + + domain_width = 3.0e5 # Width of domain (m) + domain_height = 1.0e4 # Height of domain (m) + Tsurf = 300. # Temperature at surface (K) + wind_initial = 20. # Initial wind in x direction (m/s) + pert_width = 5.0e3 # Width parameter of perturbation (m) + deltaTheta = 1.0e-2 # Magnitude of theta perturbation (K) + N = 0.01 # Brunt-Vaisala frequency (1/s) + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain -- 3D volume mesh + base_mesh = PeriodicIntervalMesh(ncolumns, domain_width) + mesh = ExtrudedMesh(base_mesh, nlayers, layer_height=domain_height/nlayers) + + # Equation + parameters = CompressibleParameters(mesh) + eqns = CompressibleEulerEquations + + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, pddumpfreq=dumpfreq, + dump_vtus=False, dump_nc=True, + ) + + model = SIQNModel(mesh, dt, parameters, eqns, output, family="CG", + element_order=element_order) + + diagnostic_fields = [Perturbation('theta')] + model.setup(diagnostic_fields) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + stepper = model.stepper + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + + # spaces + domain = model.domain + Vt = domain.spaces("theta") + Vr = domain.spaces("DG") + + # Thermodynamic constants required for setting initial conditions + # and reference profiles + g = parameters.g + + x, z = SpatialCoordinate(mesh) + + # N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) + thetab = Tsurf*exp(N**2*z/g) + + theta_b = Function(Vt).interpolate(thetab) + rho_b = Function(Vr) + + # Calculate hydrostatic exner + compressible_hydrostatic_balance(model.equation, theta_b, rho_b) + + theta_pert = ( + deltaTheta * sin(pi*z/domain_height) + / (1 + (x - domain_width/2)**2 / pert_width**2) + ) + theta0.interpolate(theta_b + theta_pert) + rho0.assign(rho_b) + u0.project(as_vector([wind_initial, 0.0])) + + stepper.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=tmax) + +# ---------------------------------------------------------------------------- # +# MAIN +# ---------------------------------------------------------------------------- # + + +if __name__ == "__main__": + + parser = ArgumentParser( + description=__doc__, + formatter_class=ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + '--ncolumns', + help="The number of columns in the vertical slice mesh.", + type=int, + default=skamarock_klemp_nonhydrostatic_defaults['ncolumns'] + ) + parser.add_argument( + '--nlayers', + help="The number of layers for the mesh.", + type=int, + default=skamarock_klemp_nonhydrostatic_defaults['nlayers'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=skamarock_klemp_nonhydrostatic_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=skamarock_klemp_nonhydrostatic_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=skamarock_klemp_nonhydrostatic_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=skamarock_klemp_nonhydrostatic_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + skamarock_klemp_nonhydrostatic(**vars(args)) diff --git a/gusto/__init__.py b/gusto/__init__.py index 08b23ae03..686f57ae1 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -7,6 +7,7 @@ from gusto.diagnostics import * # noqa from gusto.equations import * # noqa from gusto.initialisation import * # noqa +from gusto.model import * # noqa from gusto.physics import * # noqa from gusto.recovery import * # noqa from gusto.rexi import * # noqa diff --git a/gusto/model/__init__.py b/gusto/model/__init__.py new file mode 100644 index 000000000..bbd7e2134 --- /dev/null +++ b/gusto/model/__init__.py @@ -0,0 +1 @@ +from gusto.model.model import * # noqa diff --git a/gusto/model/model.py b/gusto/model/model.py new file mode 100644 index 000000000..0af330c59 --- /dev/null +++ b/gusto/model/model.py @@ -0,0 +1,113 @@ +from abc import ABCMeta, abstractmethod +from gusto.core import (Domain, IO, EmbeddedDGOptions) +from gusto.spatial_methods import DGUpwind +from gusto.time_discretisation import SSPRK3, RungeKuttaFormulation +from gusto.timestepping import SemiImplicitQuasiNewton + + +class ModelBase(object, metaclass=ABCMeta): + """Base model class.""" + + def __init__(self, mesh, dt, parameters, equation, output, + family=None, element_order=None, + no_normal_flow_bc_ids=None): + """ + Args: + mesh (:class:`Mesh`): the model's mesh. + dt (float): the model timestep. + parameters (:class:`EquationParameters`): class storing the model + equation parameters. + equation (:class:`PrognosticEquationSet`): defines the model's + prognostic equation + output (:class:`OutputParameters`): provides parameters + controlling output + family (str, optional): the finite element space family used for + the velocity field. This determines the other finite element + spaces used via the de Rham complex. If not provided, an + appropriate choice will be made based on the cell of the mesh + (or the base mesh in 3D). + element_order (int): the element degree used for the DG + space. Defaults to None. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + """ + # if HDiv finite element family not provided figure out a + # sensible default based on the cell shape of the mesh + if family is None: + extruded_mesh = hasattr(mesh, "_base_mesh") + if extruded_mesh: + cell = mesh._base_mesh.ufl_cell().cellname() + else: + cell = mesh.ufl_cell().cellname() + if cell == "interval": + family = "CG" + elif cell == "quadrilateral": + family = "RTCF" + elif cell == "triangle": + family = "BDM" + else: + raise ValueError(f"The mesh provided (or its base mesh if 3D) must have cells of type interval, quadrilateral or triangle, not {cell}.") + + # create domain + self.domain = Domain(mesh, dt, family, element_order) + + # set up prognostic equations + self.equation = equation(self.domain, parameters, + no_normal_flow_bc_ids=no_normal_flow_bc_ids) + + # save output options for when IO is set up later + self.output = output + + @abstractmethod + def setup(self): + # Set up the model. Must be implemented in the child class. + pass + + def run(self, t, tmax, pick_up=False): + """ + Runs the model for the specified time, from t to tmax + + Args: + t (float): the start time of the run + tmax (float): the end time of the run + pick_up: (bool): specify whether to pick_up from a previous run + """ + self.stepper.run(t=t, tmax=tmax, pick_up=pick_up) + + +class SIQNModel(ModelBase): + + @property + def transported_fields(self): + theta_opts = EmbeddedDGOptions() + transported_fields = [ + SSPRK3(self.domain, "u"), + SSPRK3( + self.domain, "rho", + rk_formulation=RungeKuttaFormulation.linear + ), + SSPRK3( + self.domain, "theta", + options=theta_opts + ) + ] + return transported_fields + + @property + def transport_methods(self): + transport_methods = [ + DGUpwind(self.equation, "u"), + DGUpwind(self.equation, "rho", advective_then_flux=True), + DGUpwind(self.equation, "theta") + ] + return transport_methods + + def setup(self, diagnostic_fields): + + io = IO(self.domain, self.output, diagnostic_fields=diagnostic_fields) + self.stepper = SemiImplicitQuasiNewton( + self.equation, io, self.transported_fields, + spatial_methods=self.transport_methods, + tau_values={'rho': 1.0, 'theta': 1.0} + ) From fc703518b440b0945e31028c1457518e8423c263 Mon Sep 17 00:00:00 2001 From: jshipton Date: Mon, 12 Jan 2026 15:33:17 +0000 Subject: [PATCH 02/16] generalise SIQN model class to other equation sets --- gusto/model/model.py | 55 +++++++++++++++++++++++++++++--------------- 1 file changed, 37 insertions(+), 18 deletions(-) diff --git a/gusto/model/model.py b/gusto/model/model.py index 0af330c59..073cadc41 100644 --- a/gusto/model/model.py +++ b/gusto/model/model.py @@ -80,34 +80,53 @@ class SIQNModel(ModelBase): @property def transported_fields(self): - theta_opts = EmbeddedDGOptions() - transported_fields = [ - SSPRK3(self.domain, "u"), - SSPRK3( - self.domain, "rho", - rk_formulation=RungeKuttaFormulation.linear - ), - SSPRK3( - self.domain, "theta", - options=theta_opts - ) - ] + transported_fields = [] + for field_name in self.equation.field_names: + if self.equation.space_names[field_name] == 'L2': + transported_fields.append( + SSPRK3(self.domain, field_name, + rk_formulation=RungeKuttaFormulation.linear) + ) + elif self.equation.space_names[field_name] == 'theta': + transported_fields.append( + SSPRK3(self.domain, field_name, + options=EmbeddedDGOptions()) + ) + else: + transported_fields.append( + SSPRK3( + self.domain, field_name) + ) return transported_fields @property def transport_methods(self): - transport_methods = [ - DGUpwind(self.equation, "u"), - DGUpwind(self.equation, "rho", advective_then_flux=True), - DGUpwind(self.equation, "theta") - ] + transport_methods = [] + for field_name in self.equation.field_names: + if self.equation.space_names[field_name] == 'L2': + transport_methods.append( + DGUpwind(self.equation, field_name, + advective_then_flux=True) + ) + else: + transport_methods.append( + DGUpwind(self.equation, field_name) + ) return transport_methods + @property + def tau_values(self): + tau_values = {} + for field_name in self.equation.field_names: + if field_name is not "u": + tau_values[field_name] = 1.0 + return tau_values + def setup(self, diagnostic_fields): io = IO(self.domain, self.output, diagnostic_fields=diagnostic_fields) self.stepper = SemiImplicitQuasiNewton( self.equation, io, self.transported_fields, spatial_methods=self.transport_methods, - tau_values={'rho': 1.0, 'theta': 1.0} + tau_values=self.tau_values ) From 5e0c7e775d1c6e904a32f70b927052cf583240d7 Mon Sep 17 00:00:00 2001 From: jshipton Date: Tue, 13 Jan 2026 13:28:53 +0000 Subject: [PATCH 03/16] pass through subcycling options to model and change Williamson 5 to use model --- examples/shallow_water/williamson_5.py | 37 ++++++++------------------ gusto/model/model.py | 9 +++++-- 2 files changed, 18 insertions(+), 28 deletions(-) diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index 7bf700502..2133d0b09 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -11,11 +11,9 @@ SpatialCoordinate, as_vector, pi, sqrt, min_value, Function ) from gusto import ( - Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, - ShallowWaterParameters, ShallowWaterEquations, Sum, + OutputParameters, ShallowWaterParameters, ShallowWaterEquations, Sum, lonlatr_from_xyz, GeneralIcosahedralSphereMesh, ZonalComponent, - MeridionalComponent, RelativeVorticity, RungeKuttaFormulation, - SubcyclingOptions + MeridionalComponent, RelativeVorticity, SubcyclingOptions, SIQNModel ) williamson_5_defaults = { @@ -60,18 +58,17 @@ def williamson_5( # Domain mesh = GeneralIcosahedralSphereMesh(radius, ncells_per_edge, degree=2) - domain = Domain(mesh, dt, 'BDM', element_order) - x, y, z = SpatialCoordinate(mesh) - lamda, phi, _ = lonlatr_from_xyz(x, y, z) # Equation: topography + x, y, z = SpatialCoordinate(mesh) + lamda, phi, _ = lonlatr_from_xyz(x, y, z) rsq = min_value(R0**2, (lamda - lamda_c)**2 + (phi - phi_c)**2) r = sqrt(rsq) tpexpr = mountain_height * (1 - r/R0) parameters = ShallowWaterParameters(mesh, H=mean_depth, g=g, topog_expr=tpexpr) - eqns = ShallowWaterEquations(domain, parameters) + eqns = ShallowWaterEquations # I/O output = OutputParameters( @@ -80,31 +77,19 @@ def williamson_5( ) diagnostic_fields = [Sum('D', 'topography'), RelativeVorticity(), MeridionalComponent('u'), ZonalComponent('u')] - io = IO(domain, output, diagnostic_fields=diagnostic_fields) # Transport schemes subcycling_options = SubcyclingOptions(subcycle_by_courant=0.25) - transported_fields = [ - SSPRK3(domain, "u", subcycling_options=subcycling_options), - SSPRK3( - domain, "D", subcycling_options=subcycling_options, - rk_formulation=RungeKuttaFormulation.linear - ) - ] - transport_methods = [ - DGUpwind(eqns, "u"), - DGUpwind(eqns, "D", advective_then_flux=True) - ] - - # Time stepper - stepper = SemiImplicitQuasiNewton( - eqns, io, transported_fields, transport_methods, tau_values={'D': 1.0} - ) + + model = SIQNModel(mesh, dt, parameters, eqns, output, + family='BDM', element_order=1) + model.setup(subcycling_options, diagnostic_fields) # ------------------------------------------------------------------------ # # Initial conditions # ------------------------------------------------------------------------ # + stepper = model.stepper u0 = stepper.fields('u') D0 = stepper.fields('D') Omega = parameters.Omega @@ -124,7 +109,7 @@ def williamson_5( # Run # ------------------------------------------------------------------------ # - stepper.run(t=0, tmax=tmax) + model.run(t=0, tmax=tmax) # ---------------------------------------------------------------------------- # # MAIN diff --git a/gusto/model/model.py b/gusto/model/model.py index 073cadc41..2948fbc5e 100644 --- a/gusto/model/model.py +++ b/gusto/model/model.py @@ -85,17 +85,20 @@ def transported_fields(self): if self.equation.space_names[field_name] == 'L2': transported_fields.append( SSPRK3(self.domain, field_name, + subcycling_options=self.subcycling_options, rk_formulation=RungeKuttaFormulation.linear) ) elif self.equation.space_names[field_name] == 'theta': transported_fields.append( SSPRK3(self.domain, field_name, + subcycling_options=self.subcycling_options, options=EmbeddedDGOptions()) ) else: transported_fields.append( SSPRK3( - self.domain, field_name) + self.domain, field_name, + subcycling_options=self.subcycling_options) ) return transported_fields @@ -122,7 +125,9 @@ def tau_values(self): tau_values[field_name] = 1.0 return tau_values - def setup(self, diagnostic_fields): + def setup(self, subcycling_options, diagnostic_fields): + + self.subcycling_options = subcycling_options io = IO(self.domain, self.output, diagnostic_fields=diagnostic_fields) self.stepper = SemiImplicitQuasiNewton( From 0f29a129be7f85a99f9268fdb3c7fd22a32c51e4 Mon Sep 17 00:00:00 2001 From: jshipton Date: Tue, 13 Jan 2026 16:47:15 +0000 Subject: [PATCH 04/16] correct minor typo in docs --- gusto/equations/compressible_euler_equations.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/equations/compressible_euler_equations.py b/gusto/equations/compressible_euler_equations.py index d433f1764..c11bad8c1 100644 --- a/gusto/equations/compressible_euler_equations.py +++ b/gusto/equations/compressible_euler_equations.py @@ -47,7 +47,7 @@ def __init__(self, domain, parameters, sponge_options=None, Args: domain (:class:`Domain`): the model's domain object, containing the mesh and the compatible function spaces. - x (:class:`Configuration`, optional): an object containing + parameters (:class:`Configuration`, optional): an object containing the model's physical parameters. sponge_options (:class:`SpongeLayerParameters`, optional): any parameters for applying a sponge layer to the upper boundary. From d1bcecbb65d27dcabe92faa63c7d5a309f2010f5 Mon Sep 17 00:00:00 2001 From: jshipton Date: Tue, 13 Jan 2026 17:36:25 +0000 Subject: [PATCH 05/16] output seems to go better as an arg to the setup method --- gusto/model/model.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/gusto/model/model.py b/gusto/model/model.py index 2948fbc5e..da1780f57 100644 --- a/gusto/model/model.py +++ b/gusto/model/model.py @@ -8,7 +8,7 @@ class ModelBase(object, metaclass=ABCMeta): """Base model class.""" - def __init__(self, mesh, dt, parameters, equation, output, + def __init__(self, mesh, dt, parameters, equation, family=None, element_order=None, no_normal_flow_bc_ids=None): """ @@ -19,8 +19,6 @@ def __init__(self, mesh, dt, parameters, equation, output, equation parameters. equation (:class:`PrognosticEquationSet`): defines the model's prognostic equation - output (:class:`OutputParameters`): provides parameters - controlling output family (str, optional): the finite element space family used for the velocity field. This determines the other finite element spaces used via the de Rham complex. If not provided, an @@ -125,11 +123,17 @@ def tau_values(self): tau_values[field_name] = 1.0 return tau_values - def setup(self, subcycling_options, diagnostic_fields): + def setup(self, output, subcycling_options, diagnostic_fields): + + """ + Args: + output (:class:`OutputParameters`): provides parameters + controlling output + """ self.subcycling_options = subcycling_options - io = IO(self.domain, self.output, diagnostic_fields=diagnostic_fields) + io = IO(self.domain, output, diagnostic_fields=diagnostic_fields) self.stepper = SemiImplicitQuasiNewton( self.equation, io, self.transported_fields, spatial_methods=self.transport_methods, From 17b8d52496614608d6e96fbc01226e411b300d8f Mon Sep 17 00:00:00 2001 From: jshipton Date: Tue, 13 Jan 2026 17:39:03 +0000 Subject: [PATCH 06/16] change for output arg --- examples/compressible_euler/sk_simple_model.py | 4 ++-- examples/shallow_water/williamson_5.py | 5 +++-- gusto/model/model.py | 2 +- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/examples/compressible_euler/sk_simple_model.py b/examples/compressible_euler/sk_simple_model.py index f0cf0b0d3..f920d5681 100644 --- a/examples/compressible_euler/sk_simple_model.py +++ b/examples/compressible_euler/sk_simple_model.py @@ -77,11 +77,11 @@ def skamarock_klemp_nonhydrostatic( dump_vtus=False, dump_nc=True, ) - model = SIQNModel(mesh, dt, parameters, eqns, output, family="CG", + model = SIQNModel(mesh, dt, parameters, eqns, family="CG", element_order=element_order) diagnostic_fields = [Perturbation('theta')] - model.setup(diagnostic_fields) + model.setup(output=output, diagnostic_fields=diagnostic_fields) # ------------------------------------------------------------------------ # # Initial conditions diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index 2133d0b09..113e19440 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -81,9 +81,10 @@ def williamson_5( # Transport schemes subcycling_options = SubcyclingOptions(subcycle_by_courant=0.25) - model = SIQNModel(mesh, dt, parameters, eqns, output, + model = SIQNModel(mesh, dt, parameters, eqns, family='BDM', element_order=1) - model.setup(subcycling_options, diagnostic_fields) + model.setup(output, subcycling_options=subcycling_options, + diagnostic_fields=diagnostic_fields) # ------------------------------------------------------------------------ # # Initial conditions diff --git a/gusto/model/model.py b/gusto/model/model.py index da1780f57..51950069a 100644 --- a/gusto/model/model.py +++ b/gusto/model/model.py @@ -123,7 +123,7 @@ def tau_values(self): tau_values[field_name] = 1.0 return tau_values - def setup(self, output, subcycling_options, diagnostic_fields): + def setup(self, output, subcycling_options=None, diagnostic_fields=None): """ Args: From 2142864ab450c5fd9d5fb60be7df5e6bd9b87454 Mon Sep 17 00:00:00 2001 From: jshipton Date: Tue, 13 Jan 2026 18:19:48 +0000 Subject: [PATCH 07/16] pass through some more options, use model in Schaer test (changes some of the settings) and fix lint --- .../compressible_euler/schaer_mountain.py | 70 +++++++------------ examples/shallow_water/williamson_5.py | 2 +- gusto/model/model.py | 15 ++-- 3 files changed, 37 insertions(+), 50 deletions(-) diff --git a/examples/compressible_euler/schaer_mountain.py b/examples/compressible_euler/schaer_mountain.py index 73df3419e..322736982 100644 --- a/examples/compressible_euler/schaer_mountain.py +++ b/examples/compressible_euler/schaer_mountain.py @@ -15,11 +15,11 @@ SpatialCoordinate, exp, pi, cos, Function, Mesh, Constant ) from gusto import ( - Domain, CompressibleParameters, logger, - OutputParameters, IO, SSPRK3, DGUpwind, SemiImplicitQuasiNewton, + CompressibleParameters, logger, + OutputParameters, SIQNModel, compressible_hydrostatic_balance, SpongeLayerParameters, Exner, ZComponent, - Perturbation, SUPGOptions, TrapeziumRule, MaxKernel, MinKernel, - CompressibleEulerEquations, SubcyclingOptions, RungeKuttaFormulation + Perturbation, MaxKernel, MinKernel, + CompressibleEulerEquations, SubcyclingOptions ) schaer_mountain_defaults = { @@ -93,16 +93,19 @@ def schaer_mountain( new_coords = Function(Vc).interpolate(xexpr) mesh = Mesh(new_coords) mesh._base_mesh = base_mesh # Force new mesh to inherit original base mesh - domain = Domain(mesh, dt, "CG", element_order) # Equation parameters = CompressibleParameters(mesh, g=g, cp=cp) sponge_params = SpongeLayerParameters( mesh, H=domain_height, z_level=domain_height-sponge_depth, mubar=mu_dt/dt ) - eqns = CompressibleEulerEquations( - domain, parameters, sponge_options=sponge_params, u_transport_option=u_eqn_type - ) + eqns = CompressibleEulerEquations + + # Model + model = SIQNModel(mesh, dt, parameters, eqns, + u_transport_option=u_eqn_type, + sponge_parameters=sponge_params, + family="CG", element_order=element_order) # I/O output = OutputParameters( @@ -112,46 +115,25 @@ def schaer_mountain( Exner(parameters), ZComponent('u'), Perturbation('theta'), Perturbation('rho') ] - io = IO(domain, output, diagnostic_fields=diagnostic_fields) - # Transport schemes + # Setup model subcycling_opts = SubcyclingOptions(subcycle_by_courant=0.25) - theta_opts = SUPGOptions() - transported_fields = [ - TrapeziumRule(domain, "u", subcycling_options=subcycling_opts), - SSPRK3( - domain, "rho", rk_formulation=RungeKuttaFormulation.predictor, - subcycling_options=subcycling_opts - ), - SSPRK3( - domain, "theta", options=theta_opts, - subcycling_options=subcycling_opts - ) - ] - transport_methods = [ - DGUpwind(eqns, "u"), - DGUpwind(eqns, "rho", advective_then_flux=True), - DGUpwind(eqns, "theta", ibp=theta_opts.ibp) - ] - - # Time stepper - tau_values = {'rho': 1.0, 'theta': 1.0} - stepper = SemiImplicitQuasiNewton( - eqns, io, transported_fields, transport_methods, - alpha=alpha, tau_values=tau_values, spinup_steps=spinup_steps - ) + model.setup(output, subcycling_options=subcycling_opts, + diagnostic_fields=diagnostic_fields, + alpha=alpha, spinup_steps=spinup_steps) # ------------------------------------------------------------------------ # # Initial conditions # ------------------------------------------------------------------------ # + stepper = model.stepper u0 = stepper.fields("u") rho0 = stepper.fields("rho") theta0 = stepper.fields("theta") # spaces - Vt = domain.spaces("theta") - Vr = domain.spaces("DG") + Vt = model.domain.spaces("theta") + Vr = model.domain.spaces("DG") # Thermodynamic constants required for setting initial conditions # and reference profiles @@ -175,8 +157,8 @@ def schaer_mountain( bottom_boundary = Constant(exner_surf) logger.info(f'Solving hydrostatic with bottom Exner of {exner_surf}') compressible_hydrostatic_balance( - eqns, theta_b, rho_b, exner, top=False, exner_boundary=bottom_boundary, - solve_for_rho=True + model.equation, theta_b, rho_b, exner, top=False, + exner_boundary=bottom_boundary, solve_for_rho=True ) # Solve hydrostatic balance again, but now use minimum value from first @@ -185,7 +167,8 @@ def schaer_mountain( top_boundary = Constant(top_value) logger.info(f'Solving hydrostatic with top Exner of {top_value}') compressible_hydrostatic_balance( - eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary + model.equation, theta_b, rho_b, exner, top=True, + exner_boundary=top_boundary ) max_bottom_value = max_kernel.apply(exner) @@ -209,7 +192,8 @@ def schaer_mountain( ) compressible_hydrostatic_balance( - eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary + model.equation, theta_b, rho_b, exner, top=True, + exner_boundary=top_boundary ) max_bottom_value = max_kernel.apply(exner) @@ -224,8 +208,8 @@ def schaer_mountain( # Perform a final solve to obtain hydrostatically balanced rho compressible_hydrostatic_balance( - eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary, - solve_for_rho=True + model.equation, theta_b, rho_b, exner, top=True, + exner_boundary=top_boundary, solve_for_rho=True ) theta0.assign(theta_b) @@ -238,7 +222,7 @@ def schaer_mountain( # Run # ------------------------------------------------------------------------ # - stepper.run(t=0, tmax=tmax) + model.run(t=0, tmax=tmax) # ---------------------------------------------------------------------------- # # MAIN diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index 113e19440..d2a60f572 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -82,7 +82,7 @@ def williamson_5( subcycling_options = SubcyclingOptions(subcycle_by_courant=0.25) model = SIQNModel(mesh, dt, parameters, eqns, - family='BDM', element_order=1) + family='BDM', element_order=element_order) model.setup(output, subcycling_options=subcycling_options, diagnostic_fields=diagnostic_fields) diff --git a/gusto/model/model.py b/gusto/model/model.py index 51950069a..7f4842cd2 100644 --- a/gusto/model/model.py +++ b/gusto/model/model.py @@ -9,6 +9,8 @@ class ModelBase(object, metaclass=ABCMeta): """Base model class.""" def __init__(self, mesh, dt, parameters, equation, + u_transport_option="vector_invariant_form", + sponge_parameters=None, family=None, element_order=None, no_normal_flow_bc_ids=None): """ @@ -52,11 +54,10 @@ def __init__(self, mesh, dt, parameters, equation, # set up prognostic equations self.equation = equation(self.domain, parameters, + u_transport_option=u_transport_option, + sponge_options=sponge_parameters, no_normal_flow_bc_ids=no_normal_flow_bc_ids) - # save output options for when IO is set up later - self.output = output - @abstractmethod def setup(self): # Set up the model. Must be implemented in the child class. @@ -119,11 +120,12 @@ def transport_methods(self): def tau_values(self): tau_values = {} for field_name in self.equation.field_names: - if field_name is not "u": + if field_name != "u": tau_values[field_name] = 1.0 return tau_values - def setup(self, output, subcycling_options=None, diagnostic_fields=None): + def setup(self, output, subcycling_options=None, diagnostic_fields=None, + alpha=0.5, spinup_steps=0): """ Args: @@ -137,5 +139,6 @@ def setup(self, output, subcycling_options=None, diagnostic_fields=None): self.stepper = SemiImplicitQuasiNewton( self.equation, io, self.transported_fields, spatial_methods=self.transport_methods, - tau_values=self.tau_values + tau_values=self.tau_values, + alpha=alpha, spinup_steps=spinup_steps ) From 9f34319333b59edd7a8ec65ab952a73f308ed557 Mon Sep 17 00:00:00 2001 From: jshipton Date: Tue, 13 Jan 2026 18:51:52 +0000 Subject: [PATCH 08/16] add diffusion to model and use in straka bubble (changes some settings) --- examples/compressible_euler/straka_bubble.py | 58 ++++++-------------- gusto/model/model.py | 33 +++++++++-- 2 files changed, 44 insertions(+), 47 deletions(-) diff --git a/examples/compressible_euler/straka_bubble.py b/examples/compressible_euler/straka_bubble.py index 0c2a9c96a..c5a2461f0 100644 --- a/examples/compressible_euler/straka_bubble.py +++ b/examples/compressible_euler/straka_bubble.py @@ -13,9 +13,9 @@ Function, sqrt, conditional, as_vector ) from gusto import ( - Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, - TrapeziumRule, SUPGOptions, CourantNumber, Perturbation, - DiffusionParameters, InteriorPenaltyDiffusion, BackwardEuler, + OutputParameters, SIQNModel, + CourantNumber, Perturbation, + DiffusionParameters, CompressibleParameters, CompressibleEulerEquations, compressible_hydrostatic_balance ) @@ -67,16 +67,18 @@ def straka_bubble( # Domain base_mesh = PeriodicIntervalMesh(ncolumns, domain_width) mesh = ExtrudedMesh(base_mesh, nlayers, layer_height=delta) - domain = Domain(mesh, dt, "CG", element_order) # Equation parameters = CompressibleParameters(mesh) diffusion_params = DiffusionParameters(mesh, kappa=kappa, mu=mu0/delta) diffusion_options = [("u", diffusion_params), ("theta", diffusion_params)] - eqns = CompressibleEulerEquations( - domain, parameters, u_transport_option=u_eqn_type, - diffusion_options=diffusion_options - ) + eqns = CompressibleEulerEquations + + # Model + model = SIQNModel(mesh, dt, parameters, eqns, + u_transport_option=u_eqn_type, + diffusion_options=diffusion_options, + family="CG", element_order=element_order) # I/O output = OutputParameters( @@ -86,49 +88,21 @@ def straka_bubble( diagnostic_fields = [ CourantNumber(), Perturbation('theta'), Perturbation('rho') ] - io = IO(domain, output, diagnostic_fields=diagnostic_fields) - - # Transport schemes - theta_opts = SUPGOptions() - transported_fields = [ - TrapeziumRule(domain, "u"), - SSPRK3(domain, "rho"), - SSPRK3(domain, "theta", options=theta_opts) - ] - transport_methods = [ - DGUpwind(eqns, "u"), - DGUpwind(eqns, "rho"), - DGUpwind(eqns, "theta", ibp=theta_opts.ibp) - ] - # Diffusion schemes - diffusion_schemes = [ - BackwardEuler(domain, "u"), - BackwardEuler(domain, "theta") - ] - diffusion_methods = [ - InteriorPenaltyDiffusion(eqns, "u", diffusion_params), - InteriorPenaltyDiffusion(eqns, "theta", diffusion_params) - ] - - # Time stepper - stepper = SemiImplicitQuasiNewton( - eqns, io, transported_fields, - spatial_methods=transport_methods+diffusion_methods, - diffusion_schemes=diffusion_schemes - ) + model.setup(output, diagnostic_fields=diagnostic_fields) # ------------------------------------------------------------------------ # # Initial conditions # ------------------------------------------------------------------------ # + stepper = model.stepper u0 = stepper.fields("u") rho0 = stepper.fields("rho") theta0 = stepper.fields("theta") # spaces - Vt = domain.spaces("theta") - Vr = domain.spaces("DG") + Vt = model.domain.spaces("theta") + Vr = model.domain.spaces("DG") # Isentropic background state theta_b = Function(Vt).interpolate(Tsurf) @@ -137,7 +111,7 @@ def straka_bubble( # Calculate hydrostatic exner compressible_hydrostatic_balance( - eqns, theta_b, rho_b, exner0=exner, solve_for_rho=True + model.equation, theta_b, rho_b, exner0=exner, solve_for_rho=True ) x, z = SpatialCoordinate(mesh) @@ -162,7 +136,7 @@ def straka_bubble( # Run # ------------------------------------------------------------------------ # - stepper.run(t=0, tmax=tmax) + model.run(t=0, tmax=tmax) # ---------------------------------------------------------------------------- # # MAIN diff --git a/gusto/model/model.py b/gusto/model/model.py index 7f4842cd2..c8a43c5a0 100644 --- a/gusto/model/model.py +++ b/gusto/model/model.py @@ -1,7 +1,8 @@ from abc import ABCMeta, abstractmethod from gusto.core import (Domain, IO, EmbeddedDGOptions) -from gusto.spatial_methods import DGUpwind -from gusto.time_discretisation import SSPRK3, RungeKuttaFormulation +from gusto.spatial_methods import DGUpwind, InteriorPenaltyDiffusion +from gusto.time_discretisation import (SSPRK3, RungeKuttaFormulation, + BackwardEuler) from gusto.timestepping import SemiImplicitQuasiNewton @@ -10,7 +11,7 @@ class ModelBase(object, metaclass=ABCMeta): def __init__(self, mesh, dt, parameters, equation, u_transport_option="vector_invariant_form", - sponge_parameters=None, + sponge_options=None, diffusion_options=None, family=None, element_order=None, no_normal_flow_bc_ids=None): """ @@ -55,9 +56,12 @@ def __init__(self, mesh, dt, parameters, equation, # set up prognostic equations self.equation = equation(self.domain, parameters, u_transport_option=u_transport_option, - sponge_options=sponge_parameters, + sponge_options=sponge_options, + diffusion_options=diffusion_options, no_normal_flow_bc_ids=no_normal_flow_bc_ids) + self.diffusion_options = diffusion_options + @abstractmethod def setup(self): # Set up the model. Must be implemented in the child class. @@ -77,6 +81,24 @@ def run(self, t, tmax, pick_up=False): class SIQNModel(ModelBase): + @property + def diffusion_methods(self): + diffusion_methods = [] + for field, params in self.diffusion_options: + diffusion_methods.append( + InteriorPenaltyDiffusion(self.equation, field, params) + ) + return diffusion_methods + + @property + def diffusion_schemes(self): + diffusion_schemes = [] + for field, _ in self.diffusion_options: + diffusion_schemes.append( + BackwardEuler(self.domain, field) + ) + return diffusion_schemes + @property def transported_fields(self): transported_fields = [] @@ -138,7 +160,8 @@ def setup(self, output, subcycling_options=None, diagnostic_fields=None, io = IO(self.domain, output, diagnostic_fields=diagnostic_fields) self.stepper = SemiImplicitQuasiNewton( self.equation, io, self.transported_fields, - spatial_methods=self.transport_methods, + spatial_methods=self.transport_methods+self.diffusion_methods, + diffusion_schemes=self.diffusion_schemes, tau_values=self.tau_values, alpha=alpha, spinup_steps=spinup_steps ) From b4acacd369f7f4c0b072e1ef8bda8fc1982eb498 Mon Sep 17 00:00:00 2001 From: jshipton Date: Tue, 13 Jan 2026 19:47:50 +0000 Subject: [PATCH 09/16] be sneaky with kwargs to get over issue that different equation sets have different interfaces; also use to enable passing different physics options to SIQN timestepper. Do not expect Tom to like... --- gusto/model/model.py | 23 +++++++++-------------- 1 file changed, 9 insertions(+), 14 deletions(-) diff --git a/gusto/model/model.py b/gusto/model/model.py index c8a43c5a0..7181bddab 100644 --- a/gusto/model/model.py +++ b/gusto/model/model.py @@ -10,10 +10,8 @@ class ModelBase(object, metaclass=ABCMeta): """Base model class.""" def __init__(self, mesh, dt, parameters, equation, - u_transport_option="vector_invariant_form", - sponge_options=None, diffusion_options=None, family=None, element_order=None, - no_normal_flow_bc_ids=None): + no_normal_flow_bc_ids=None, **kwargs): """ Args: mesh (:class:`Mesh`): the model's mesh. @@ -54,13 +52,9 @@ def __init__(self, mesh, dt, parameters, equation, self.domain = Domain(mesh, dt, family, element_order) # set up prognostic equations - self.equation = equation(self.domain, parameters, - u_transport_option=u_transport_option, - sponge_options=sponge_options, - diffusion_options=diffusion_options, - no_normal_flow_bc_ids=no_normal_flow_bc_ids) + self.equation = equation(self.domain, parameters, **kwargs) - self.diffusion_options = diffusion_options + self.diffusion_options = kwargs.get("diffusion_options") @abstractmethod def setup(self): @@ -146,8 +140,7 @@ def tau_values(self): tau_values[field_name] = 1.0 return tau_values - def setup(self, output, subcycling_options=None, diagnostic_fields=None, - alpha=0.5, spinup_steps=0): + def setup(self, output, **kwargs): """ Args: @@ -155,13 +148,15 @@ def setup(self, output, subcycling_options=None, diagnostic_fields=None, controlling output """ - self.subcycling_options = subcycling_options - + diagnostic_fields = kwargs.pop("diagnostic_fields") io = IO(self.domain, output, diagnostic_fields=diagnostic_fields) + + self.subcycling_options = kwargs.get("subcycling_options") + self.stepper = SemiImplicitQuasiNewton( self.equation, io, self.transported_fields, spatial_methods=self.transport_methods+self.diffusion_methods, diffusion_schemes=self.diffusion_schemes, tau_values=self.tau_values, - alpha=alpha, spinup_steps=spinup_steps + **kwargs ) From c25d38eddb9439f2b886070fabbca151d0b51321 Mon Sep 17 00:00:00 2001 From: jshipton Date: Tue, 13 Jan 2026 21:00:43 +0000 Subject: [PATCH 10/16] add lowest order SIQN model - not sure what to do for diffusion... --- .../compressible_euler/dry_bryan_fritsch.py | 53 ++---- gusto/model/model.py | 173 ++++++++++++++---- 2 files changed, 148 insertions(+), 78 deletions(-) diff --git a/examples/compressible_euler/dry_bryan_fritsch.py b/examples/compressible_euler/dry_bryan_fritsch.py index bd3222ca7..7a975ec90 100644 --- a/examples/compressible_euler/dry_bryan_fritsch.py +++ b/examples/compressible_euler/dry_bryan_fritsch.py @@ -13,8 +13,8 @@ LinearVariationalProblem, LinearVariationalSolver ) from gusto import ( - Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, - RecoverySpaces, BoundaryMethod, Perturbation, CompressibleParameters, + OutputParameters, LowestOrderModel, + Perturbation, CompressibleParameters, CompressibleEulerEquations, compressible_hydrostatic_balance ) @@ -52,7 +52,6 @@ def dry_bryan_fritsch( # Our settings for this set up # ------------------------------------------------------------------------ # - element_order = 0 u_eqn_type = 'vector_advection_form' # ------------------------------------------------------------------------ # @@ -62,14 +61,15 @@ def dry_bryan_fritsch( # Domain base_mesh = IntervalMesh(ncolumns, domain_width) mesh = ExtrudedMesh(base_mesh, nlayers, layer_height=domain_height/nlayers) - domain = Domain(mesh, dt, "CG", element_order) # Equation params = CompressibleParameters(mesh) - eqns = CompressibleEulerEquations( - domain, params, u_transport_option=u_eqn_type, - no_normal_flow_bc_ids=[1, 2] - ) + eqns = CompressibleEulerEquations + + # Model + model = LowestOrderModel(mesh, dt, params, eqns, + u_transport_option=u_eqn_type, + family="CG", no_normal_flow_bc_ids=[1, 2]) # I/O output = OutputParameters( @@ -77,46 +77,21 @@ def dry_bryan_fritsch( dumplist=['rho'] ) diagnostic_fields = [Perturbation('theta')] - io = IO(domain, output, diagnostic_fields=diagnostic_fields) - - # Transport schemes -- set up options for using recovery wrapper - boundary_methods = {'DG': BoundaryMethod.taylor} - - recovery_spaces = RecoverySpaces( - domain, boundary_methods, use_vector_spaces=True - ) - u_opts = recovery_spaces.HDiv_options - rho_opts = recovery_spaces.DG_options - theta_opts = recovery_spaces.theta_options - - transported_fields = [ - SSPRK3(domain, "rho", options=rho_opts), - SSPRK3(domain, "theta", options=theta_opts), - SSPRK3(domain, "u", options=u_opts) - ] - - transport_methods = [ - DGUpwind(eqns, field) for field in ["u", "rho", "theta"] - ] - - # Time stepper - stepper = SemiImplicitQuasiNewton( - eqns, io, transported_fields, transport_methods, - tau_values={'rho': 1.0, 'theta': 1.0} - ) + model.setup(output, diagnostic_fields=diagnostic_fields) # ------------------------------------------------------------------------ # # Initial conditions # ------------------------------------------------------------------------ # + stepper = model.stepper u0 = stepper.fields("u") rho0 = stepper.fields("rho") theta0 = stepper.fields("theta") # spaces - Vt = domain.spaces("theta") - Vr = domain.spaces("DG") + Vt = model.domain.spaces("theta") + Vr = model.domain.spaces("DG") x, z = SpatialCoordinate(mesh) # Define constant theta_e and water_t @@ -127,7 +102,7 @@ def dry_bryan_fritsch( u0.project(as_vector([zero, zero])) # Calculate hydrostatic fields - compressible_hydrostatic_balance(eqns, theta_b, rho0, solve_for_rho=True) + compressible_hydrostatic_balance(model.equation, theta_b, rho0, solve_for_rho=True) # make mean fields rho_b = Function(Vr).assign(rho0) @@ -161,7 +136,7 @@ def dry_bryan_fritsch( # Run # ------------------------------------------------------------------------ # - stepper.run(t=0, tmax=tmax) + model.run(t=0, tmax=tmax) # ---------------------------------------------------------------------------- # # MAIN diff --git a/gusto/model/model.py b/gusto/model/model.py index 7181bddab..5a42358cf 100644 --- a/gusto/model/model.py +++ b/gusto/model/model.py @@ -1,5 +1,6 @@ -from abc import ABCMeta, abstractmethod +from abc import ABCMeta, abstractmethod, abstractproperty from gusto.core import (Domain, IO, EmbeddedDGOptions) +from gusto.recovery import BoundaryMethod, RecoverySpaces from gusto.spatial_methods import DGUpwind, InteriorPenaltyDiffusion from gusto.time_discretisation import (SSPRK3, RungeKuttaFormulation, BackwardEuler) @@ -54,7 +55,7 @@ def __init__(self, mesh, dt, parameters, equation, # set up prognostic equations self.equation = equation(self.domain, parameters, **kwargs) - self.diffusion_options = kwargs.get("diffusion_options") + self.diffusion_options = kwargs.get("diffusion_options", []) @abstractmethod def setup(self): @@ -75,88 +76,182 @@ def run(self, t, tmax, pick_up=False): class SIQNModel(ModelBase): + @abstractproperty + def diffusion_methods(self): + pass + + @abstractproperty + def diffusion_schemes(self): + pass + + @abstractproperty + def transported_fields(self): + pass + + @abstractproperty + def transport_methods(self): + pass + + @property + def tau_values(self): + _tau_values = {} + for field_name in self.equation.field_names: + if field_name != "u": + _tau_values[field_name] = 1.0 + return _tau_values + pass + + def setup(self, output, **kwargs): + + """ + Args: + output (:class:`OutputParameters`): provides parameters + controlling output + """ + + diagnostic_fields = kwargs.pop("diagnostic_fields") + io = IO(self.domain, output, diagnostic_fields=diagnostic_fields) + + self.subcycling_options = kwargs.get("subcycling_options") + + self.stepper = SemiImplicitQuasiNewton( + self.equation, io, self.transported_fields, + spatial_methods=self.transport_methods+self.diffusion_methods, + diffusion_schemes=self.diffusion_schemes, + tau_values=self.tau_values, + **kwargs + ) + + +class Model(SIQNModel): + + def __init__(self, mesh, dt, parameters, equation, + family=None, + no_normal_flow_bc_ids=None, **kwargs): + + super().__init__(mesh, dt, parameters, equation, + family=family, element_order=1, + no_normal_flow_bc_ids=no_normal_flow_bc_ids, + **kwargs) + @property def diffusion_methods(self): - diffusion_methods = [] + _diffusion_methods = [] for field, params in self.diffusion_options: - diffusion_methods.append( + _diffusion_methods.append( InteriorPenaltyDiffusion(self.equation, field, params) ) - return diffusion_methods + return _diffusion_methods @property def diffusion_schemes(self): - diffusion_schemes = [] + _diffusion_schemes = [] for field, _ in self.diffusion_options: - diffusion_schemes.append( + _diffusion_schemes.append( BackwardEuler(self.domain, field) ) - return diffusion_schemes + return _diffusion_schemes @property def transported_fields(self): - transported_fields = [] + _transported_fields = [] for field_name in self.equation.field_names: if self.equation.space_names[field_name] == 'L2': - transported_fields.append( + _transported_fields.append( SSPRK3(self.domain, field_name, subcycling_options=self.subcycling_options, rk_formulation=RungeKuttaFormulation.linear) ) elif self.equation.space_names[field_name] == 'theta': - transported_fields.append( + _transported_fields.append( SSPRK3(self.domain, field_name, subcycling_options=self.subcycling_options, options=EmbeddedDGOptions()) ) else: - transported_fields.append( + _transported_fields.append( SSPRK3( self.domain, field_name, subcycling_options=self.subcycling_options) ) - return transported_fields + return _transported_fields @property def transport_methods(self): - transport_methods = [] + _transport_methods = [] for field_name in self.equation.field_names: if self.equation.space_names[field_name] == 'L2': - transport_methods.append( + _transport_methods.append( DGUpwind(self.equation, field_name, advective_then_flux=True) ) else: - transport_methods.append( + _transport_methods.append( DGUpwind(self.equation, field_name) ) - return transport_methods + return _transport_methods - @property - def tau_values(self): - tau_values = {} - for field_name in self.equation.field_names: - if field_name != "u": - tau_values[field_name] = 1.0 - return tau_values - def setup(self, output, **kwargs): +class LowestOrderModel(SIQNModel): - """ - Args: - output (:class:`OutputParameters`): provides parameters - controlling output - """ + def __init__(self, mesh, dt, parameters, equation, + family=None, + no_normal_flow_bc_ids=None, **kwargs): - diagnostic_fields = kwargs.pop("diagnostic_fields") - io = IO(self.domain, output, diagnostic_fields=diagnostic_fields) + super().__init__(mesh, dt, parameters, equation, + family=family, element_order=0, + no_normal_flow_bc_ids=no_normal_flow_bc_ids, + **kwargs) - self.subcycling_options = kwargs.get("subcycling_options") + @property + def diffusion_methods(self): + _diffusion_methods = [] + for field, params in self.diffusion_options: + _diffusion_methods.append( + InteriorPenaltyDiffusion(self.equation, field, params) + ) + return _diffusion_methods - self.stepper = SemiImplicitQuasiNewton( - self.equation, io, self.transported_fields, - spatial_methods=self.transport_methods+self.diffusion_methods, - diffusion_schemes=self.diffusion_schemes, - tau_values=self.tau_values, - **kwargs + @property + def diffusion_schemes(self): + _diffusion_schemes = [] + for field, _ in self.diffusion_options: + _diffusion_schemes.append( + BackwardEuler(self.domain, field) + ) + return _diffusion_schemes + + @property + def transported_fields(self): + boundary_methods = {'DG': BoundaryMethod.taylor} + recovery_spaces = RecoverySpaces( + self.domain, boundary_methods, use_vector_spaces=True ) + _transported_fields = [] + for field_name in self.equation.field_names: + if self.equation.space_names[field_name] == 'HDiv': + _transported_fields.append( + SSPRK3(self.domain, field_name, + subcycling_options=self.subcycling_options, + options=recovery_spaces.HDiv_options) + ) + elif self.equation.space_names[field_name] == 'L2': + _transported_fields.append( + SSPRK3(self.domain, field_name, + subcycling_options=self.subcycling_options, + options=recovery_spaces.DG_options) + ) + elif self.equation.space_names[field_name] == 'theta': + _transported_fields.append( + SSPRK3(self.domain, field_name, + subcycling_options=self.subcycling_options, + options=recovery_spaces.theta_options) + ) + return _transported_fields + + @property + def transport_methods(self): + _transport_methods = [] + for field_name in self.equation.field_names: + _transport_methods.append(DGUpwind(self.equation, field_name)) + return _transport_methods From 92ca0b4a40ddb2f0444daaebb9fc59cf0b3ad73d Mon Sep 17 00:00:00 2001 From: jshipton Date: Tue, 13 Jan 2026 22:13:23 +0000 Subject: [PATCH 11/16] model name change in example scripts --- examples/compressible_euler/schaer_mountain.py | 10 +++++----- examples/compressible_euler/straka_bubble.py | 10 +++++----- examples/shallow_water/williamson_5.py | 6 +++--- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/examples/compressible_euler/schaer_mountain.py b/examples/compressible_euler/schaer_mountain.py index 322736982..300abef6f 100644 --- a/examples/compressible_euler/schaer_mountain.py +++ b/examples/compressible_euler/schaer_mountain.py @@ -16,7 +16,7 @@ ) from gusto import ( CompressibleParameters, logger, - OutputParameters, SIQNModel, + OutputParameters, Model, compressible_hydrostatic_balance, SpongeLayerParameters, Exner, ZComponent, Perturbation, MaxKernel, MinKernel, CompressibleEulerEquations, SubcyclingOptions @@ -102,10 +102,10 @@ def schaer_mountain( eqns = CompressibleEulerEquations # Model - model = SIQNModel(mesh, dt, parameters, eqns, - u_transport_option=u_eqn_type, - sponge_parameters=sponge_params, - family="CG", element_order=element_order) + model = Model(mesh, dt, parameters, eqns, + u_transport_option=u_eqn_type, + sponge_parameters=sponge_params, + family="CG", element_order=element_order) # I/O output = OutputParameters( diff --git a/examples/compressible_euler/straka_bubble.py b/examples/compressible_euler/straka_bubble.py index c5a2461f0..e7a4236c5 100644 --- a/examples/compressible_euler/straka_bubble.py +++ b/examples/compressible_euler/straka_bubble.py @@ -13,7 +13,7 @@ Function, sqrt, conditional, as_vector ) from gusto import ( - OutputParameters, SIQNModel, + OutputParameters, Model, CourantNumber, Perturbation, DiffusionParameters, CompressibleParameters, CompressibleEulerEquations, @@ -75,10 +75,10 @@ def straka_bubble( eqns = CompressibleEulerEquations # Model - model = SIQNModel(mesh, dt, parameters, eqns, - u_transport_option=u_eqn_type, - diffusion_options=diffusion_options, - family="CG", element_order=element_order) + model = Model(mesh, dt, parameters, eqns, + u_transport_option=u_eqn_type, + diffusion_options=diffusion_options, + family="CG", element_order=element_order) # I/O output = OutputParameters( diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index d2a60f572..0364d9776 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -13,7 +13,7 @@ from gusto import ( OutputParameters, ShallowWaterParameters, ShallowWaterEquations, Sum, lonlatr_from_xyz, GeneralIcosahedralSphereMesh, ZonalComponent, - MeridionalComponent, RelativeVorticity, SubcyclingOptions, SIQNModel + MeridionalComponent, RelativeVorticity, SubcyclingOptions, Model ) williamson_5_defaults = { @@ -81,8 +81,8 @@ def williamson_5( # Transport schemes subcycling_options = SubcyclingOptions(subcycle_by_courant=0.25) - model = SIQNModel(mesh, dt, parameters, eqns, - family='BDM', element_order=element_order) + model = Model(mesh, dt, parameters, eqns, + family='BDM', element_order=element_order) model.setup(output, subcycling_options=subcycling_options, diagnostic_fields=diagnostic_fields) From c861b604ed2198225795e17708cc7f932b134262 Mon Sep 17 00:00:00 2001 From: jshipton Date: Wed, 14 Jan 2026 08:30:15 +0000 Subject: [PATCH 12/16] fixes --- examples/compressible_euler/schaer_mountain.py | 3 +-- examples/compressible_euler/sk_simple_model.py | 6 +++--- examples/shallow_water/williamson_5.py | 8 +------- gusto/model/model.py | 2 +- 4 files changed, 6 insertions(+), 13 deletions(-) diff --git a/examples/compressible_euler/schaer_mountain.py b/examples/compressible_euler/schaer_mountain.py index 300abef6f..18763d9c8 100644 --- a/examples/compressible_euler/schaer_mountain.py +++ b/examples/compressible_euler/schaer_mountain.py @@ -66,7 +66,6 @@ def schaer_mountain( spinup_steps = 5 # Not necessary but helps balance initial conditions alpha = 0.51 # Necessary to absorb grid scale waves - element_order = 1 u_eqn_type = 'vector_invariant_form' # ------------------------------------------------------------------------ # @@ -105,7 +104,7 @@ def schaer_mountain( model = Model(mesh, dt, parameters, eqns, u_transport_option=u_eqn_type, sponge_parameters=sponge_params, - family="CG", element_order=element_order) + family="CG") # I/O output = OutputParameters( diff --git a/examples/compressible_euler/sk_simple_model.py b/examples/compressible_euler/sk_simple_model.py index f920d5681..d604d865e 100644 --- a/examples/compressible_euler/sk_simple_model.py +++ b/examples/compressible_euler/sk_simple_model.py @@ -19,7 +19,7 @@ ) from gusto import ( Perturbation, CompressibleParameters, OutputParameters, - CompressibleEulerEquations, SIQNModel, + CompressibleEulerEquations, Model, compressible_hydrostatic_balance ) @@ -77,8 +77,8 @@ def skamarock_klemp_nonhydrostatic( dump_vtus=False, dump_nc=True, ) - model = SIQNModel(mesh, dt, parameters, eqns, family="CG", - element_order=element_order) + model = Model(mesh, dt, parameters, eqns, family="CG", + element_order=element_order) diagnostic_fields = [Perturbation('theta')] model.setup(output=output, diagnostic_fields=diagnostic_fields) diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index 0364d9776..c572de5e1 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -46,12 +46,6 @@ def williamson_5( lamda_c = -pi/2. # longitudinal centre of mountain (rad) phi_c = pi/6. # latitudinal centre of mountain (rad) - # ------------------------------------------------------------------------ # - # Our settings for this set up - # ------------------------------------------------------------------------ # - - element_order = 1 - # ------------------------------------------------------------------------ # # Set up model objects # ------------------------------------------------------------------------ # @@ -82,7 +76,7 @@ def williamson_5( subcycling_options = SubcyclingOptions(subcycle_by_courant=0.25) model = Model(mesh, dt, parameters, eqns, - family='BDM', element_order=element_order) + family='BDM') model.setup(output, subcycling_options=subcycling_options, diagnostic_fields=diagnostic_fields) diff --git a/gusto/model/model.py b/gusto/model/model.py index 5a42358cf..dc1636bac 100644 --- a/gusto/model/model.py +++ b/gusto/model/model.py @@ -112,7 +112,7 @@ def setup(self, output, **kwargs): diagnostic_fields = kwargs.pop("diagnostic_fields") io = IO(self.domain, output, diagnostic_fields=diagnostic_fields) - self.subcycling_options = kwargs.get("subcycling_options") + self.subcycling_options = kwargs.pop("subcycling_options") self.stepper = SemiImplicitQuasiNewton( self.equation, io, self.transported_fields, From 731e762252cb87468f9da2d84bcc251386a2b21d Mon Sep 17 00:00:00 2001 From: jshipton Date: Wed, 14 Jan 2026 08:48:24 +0000 Subject: [PATCH 13/16] fix default to None --- gusto/model/model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gusto/model/model.py b/gusto/model/model.py index dc1636bac..857053cf8 100644 --- a/gusto/model/model.py +++ b/gusto/model/model.py @@ -109,10 +109,10 @@ def setup(self, output, **kwargs): controlling output """ - diagnostic_fields = kwargs.pop("diagnostic_fields") + diagnostic_fields = kwargs.pop("diagnostic_fields", None) io = IO(self.domain, output, diagnostic_fields=diagnostic_fields) - self.subcycling_options = kwargs.pop("subcycling_options") + self.subcycling_options = kwargs.pop("subcycling_options", None) self.stepper = SemiImplicitQuasiNewton( self.equation, io, self.transported_fields, From 49e35a6ab2da3a90d71f09c83f48c63dbd4ee48b Mon Sep 17 00:00:00 2001 From: jshipton Date: Wed, 14 Jan 2026 08:50:14 +0000 Subject: [PATCH 14/16] fix call to Model class --- examples/compressible_euler/straka_bubble.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/compressible_euler/straka_bubble.py b/examples/compressible_euler/straka_bubble.py index e7a4236c5..a2a016769 100644 --- a/examples/compressible_euler/straka_bubble.py +++ b/examples/compressible_euler/straka_bubble.py @@ -57,7 +57,6 @@ def straka_bubble( delta = domain_height/nlayers ncolumns = 8 * nlayers - element_order = 1 u_eqn_type = 'vector_advection_form' # ------------------------------------------------------------------------ # @@ -78,7 +77,7 @@ def straka_bubble( model = Model(mesh, dt, parameters, eqns, u_transport_option=u_eqn_type, diffusion_options=diffusion_options, - family="CG", element_order=element_order) + family="CG") # I/O output = OutputParameters( From c4e2c0dcff5eac9158a13f00019a93a0a9c105d7 Mon Sep 17 00:00:00 2001 From: jshipton Date: Wed, 14 Jan 2026 08:52:46 +0000 Subject: [PATCH 15/16] use correction keyword for sponge options --- examples/compressible_euler/schaer_mountain.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/compressible_euler/schaer_mountain.py b/examples/compressible_euler/schaer_mountain.py index 18763d9c8..e6c5c0f12 100644 --- a/examples/compressible_euler/schaer_mountain.py +++ b/examples/compressible_euler/schaer_mountain.py @@ -103,7 +103,7 @@ def schaer_mountain( # Model model = Model(mesh, dt, parameters, eqns, u_transport_option=u_eqn_type, - sponge_parameters=sponge_params, + sponge_options=sponge_params, family="CG") # I/O From 25d98704481e5da71db1c0a7a718c7468a5fa1f1 Mon Sep 17 00:00:00 2001 From: jshipton Date: Wed, 14 Jan 2026 12:10:55 +0000 Subject: [PATCH 16/16] pass limiters through when settin gup model and adapt unsaturated bubble to use model --- .../compressible_euler/unsaturated_bubble.py | 97 ++++++++----------- gusto/model/model.py | 19 ++-- 2 files changed, 52 insertions(+), 64 deletions(-) diff --git a/examples/compressible_euler/unsaturated_bubble.py b/examples/compressible_euler/unsaturated_bubble.py index 47ba48de3..7d443fedb 100644 --- a/examples/compressible_euler/unsaturated_bubble.py +++ b/examples/compressible_euler/unsaturated_bubble.py @@ -23,8 +23,8 @@ ) from firedrake.slope_limiter.vertex_based_limiter import VertexBasedLimiter from gusto import ( - Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, - Perturbation, RecoverySpaces, BoundaryMethod, Recoverer, Fallout, + OutputParameters, LowestOrderModel, DGUpwind, SSPRK3, + Perturbation, BoundaryMethod, Recoverer, Fallout, Coalescence, SaturationAdjustment, EvaporationOfRain, thermodynamics, CompressibleParameters, CompressibleEulerEquations, unsaturated_hydrostatic_balance, WaterVapour, CloudWater, Rain, @@ -71,7 +71,6 @@ def unsaturated_bubble( # Our settings for this set up # ------------------------------------------------------------------------ # - element_order = 0 u_eqn_type = 'vector_advection_form' # ------------------------------------------------------------------------ # @@ -81,14 +80,16 @@ def unsaturated_bubble( # Domain base_mesh = PeriodicIntervalMesh(ncolumns, domain_width) mesh = ExtrudedMesh(base_mesh, nlayers, layer_height=domain_height/nlayers) - domain = Domain(mesh, dt, "CG", element_order) # Equation - params = CompressibleParameters(mesh) + parameters = CompressibleParameters(mesh) tracers = [WaterVapour(), CloudWater(), Rain()] - eqns = CompressibleEulerEquations( - domain, params, active_tracers=tracers, u_transport_option=u_eqn_type - ) + eqns = CompressibleEulerEquations + + model = LowestOrderModel(mesh, dt, parameters, eqns, + u_transport_option=u_eqn_type, + active_tracers=tracers, + family="CG") # I/O output = OutputParameters( @@ -96,61 +97,41 @@ def unsaturated_bubble( dumplist=['cloud_water', 'rain'] ) diagnostic_fields = [ - RelativeHumidity(eqns), Perturbation('theta'), Perturbation('rho'), + RelativeHumidity(model.equation), Perturbation('theta'), + Perturbation('rho'), Perturbation('water_vapour'), Perturbation('RelativeHumidity') ] - io = IO(domain, output, diagnostic_fields=diagnostic_fields) - - # Transport schemes -- specify options for using recovery wrapper - boundary_methods = {'DG': BoundaryMethod.taylor} - - recovery_spaces = RecoverySpaces(domain, boundary_method=boundary_methods, use_vector_spaces=True) - - u_opts = recovery_spaces.HDiv_options - rho_opts = recovery_spaces.DG_options - theta_opts = recovery_spaces.theta_options - VDG1 = domain.spaces("DG1_equispaced") + VDG1 = model.domain.spaces("DG1_equispaced") limiter = VertexBasedLimiter(VDG1) - - transported_fields = [ - SSPRK3(domain, "u", options=u_opts), - SSPRK3(domain, "rho", options=rho_opts), - SSPRK3(domain, "theta", options=theta_opts), - SSPRK3(domain, "water_vapour", options=theta_opts, limiter=limiter), - SSPRK3(domain, "cloud_water", options=theta_opts, limiter=limiter), - SSPRK3(domain, "rain", options=theta_opts, limiter=limiter) - ] - - transport_methods = [ - DGUpwind(eqns, field) for field in - ["u", "rho", "theta", "water_vapour", "cloud_water", "rain"] - ] + limiters = {"water_vapour": limiter, + "cloud_water": limiter, + "rain": limiter} # Physics schemes - Vt = domain.spaces('theta') - rainfall_method = DGUpwind(eqns, 'rain', outflow=True) + Vt = model.domain.spaces('theta') + rainfall_method = DGUpwind(model.equation, 'rain', outflow=True) zero_limiter = MixedFSLimiter( - eqns, + model.equation, {'water_vapour': ZeroLimiter(Vt), 'cloud_water': ZeroLimiter(Vt)} ) physics_schemes = [ - (Fallout(eqns, 'rain', domain, rainfall_method), SSPRK3(domain)), - (Coalescence(eqns), ForwardEuler(domain)), - (EvaporationOfRain(eqns), ForwardEuler(domain)), - (SaturationAdjustment(eqns), ForwardEuler(domain, limiter=zero_limiter)) + (Fallout(model.equation, 'rain', model.domain, rainfall_method), SSPRK3(model.domain)), + (Coalescence(model.equation), ForwardEuler(model.domain)), + (EvaporationOfRain(model.equation), ForwardEuler(model.domain)), + (SaturationAdjustment(model.equation), ForwardEuler(model.domain, limiter=zero_limiter)) ] - # Time stepper - stepper = SemiImplicitQuasiNewton( - eqns, io, transported_fields, transport_methods, - final_physics_schemes=physics_schemes, tau_values={'rho': 1.0, 'theta': 1.0} - ) + # Setup model + model.setup(output, limiters=limiters, + final_physics_schemes=physics_schemes, + diagnostic_fields=diagnostic_fields) # ------------------------------------------------------------------------ # # Initial conditions # ------------------------------------------------------------------------ # + stepper = model.stepper u0 = stepper.fields("u") rho0 = stepper.fields("rho") theta0 = stepper.fields("theta") @@ -159,7 +140,7 @@ def unsaturated_bubble( water_r0 = stepper.fields("rain") # spaces - Vr = domain.spaces("DG") + Vr = model.domain.spaces("DG") x, z = SpatialCoordinate(mesh) quadrature_degree = (4, 4) dxp = dx(degree=(quadrature_degree)) @@ -167,14 +148,14 @@ def unsaturated_bubble( physics_boundary_method = BoundaryMethod.extruded # Define constant theta_e and water_t - exner_surf = (float(psurf) / float(eqns.parameters.p_0)) ** float(eqns.parameters.kappa) - theta_surf = thermodynamics.theta(eqns.parameters, Tsurf, psurf) + exner_surf = (float(psurf) / float(parameters.p_0)) ** float(parameters.kappa) + theta_surf = thermodynamics.theta(parameters, Tsurf, psurf) theta_d = Function(Vt).interpolate(theta_surf * exp(S*z)) rel_hum = Function(Vt).assign(rel_hum_background) # Calculate hydrostatic fields unsaturated_hydrostatic_balance( - eqns, stepper.fields, theta_d, rel_hum, + model.equation, stepper.fields, theta_d, rel_hum, exner_boundary=Constant(exner_surf) ) @@ -208,16 +189,16 @@ def unsaturated_bubble( water_v_eval = Function(Vt) delta = 1.0 - R_d = eqns.parameters.R_d - R_v = eqns.parameters.R_v + R_d = parameters.R_d + R_v = parameters.R_v epsilon = R_d / R_v # make expressions to evaluate residual - exner_expr = thermodynamics.exner_pressure(eqns.parameters, rho_averaged, theta0) - p_expr = thermodynamics.p(eqns.parameters, exner_expr) - T_expr = thermodynamics.T(eqns.parameters, theta0, exner_expr, water_v0) - water_v_expr = thermodynamics.r_v(eqns.parameters, rel_hum, T_expr, p_expr) - rel_hum_expr = thermodynamics.RH(eqns.parameters, water_v0, T_expr, p_expr) + exner_expr = thermodynamics.exner_pressure(parameters, rho_averaged, theta0) + p_expr = thermodynamics.p(parameters, exner_expr) + T_expr = thermodynamics.T(parameters, theta0, exner_expr, water_v0) + water_v_expr = thermodynamics.r_v(parameters, rel_hum, T_expr, p_expr) + rel_hum_expr = thermodynamics.RH(parameters, water_v0, T_expr, p_expr) rel_hum_eval = Function(Vt) # set-up rho problem to keep exner constant @@ -275,7 +256,7 @@ def unsaturated_bubble( # Run # ------------------------------------------------------------------------ # - stepper.run(t=0, tmax=tmax) + model.run(t=0, tmax=tmax) # ---------------------------------------------------------------------------- # # MAIN diff --git a/gusto/model/model.py b/gusto/model/model.py index 857053cf8..b63b13a60 100644 --- a/gusto/model/model.py +++ b/gusto/model/model.py @@ -113,6 +113,7 @@ def setup(self, output, **kwargs): io = IO(self.domain, output, diagnostic_fields=diagnostic_fields) self.subcycling_options = kwargs.pop("subcycling_options", None) + self.limiters = kwargs.pop("limiters", {}) self.stepper = SemiImplicitQuasiNewton( self.equation, io, self.transported_fields, @@ -160,19 +161,22 @@ def transported_fields(self): _transported_fields.append( SSPRK3(self.domain, field_name, subcycling_options=self.subcycling_options, - rk_formulation=RungeKuttaFormulation.linear) + rk_formulation=RungeKuttaFormulation.linear, + limiter=self.limiters.get(field_name)) ) elif self.equation.space_names[field_name] == 'theta': _transported_fields.append( SSPRK3(self.domain, field_name, subcycling_options=self.subcycling_options, - options=EmbeddedDGOptions()) + options=EmbeddedDGOptions(), + limiter=self.limiters.get(field_name)) ) else: _transported_fields.append( SSPRK3( self.domain, field_name, - subcycling_options=self.subcycling_options) + subcycling_options=self.subcycling_options, + limiter=self.limiters.get(field_name)) ) return _transported_fields @@ -233,19 +237,22 @@ def transported_fields(self): _transported_fields.append( SSPRK3(self.domain, field_name, subcycling_options=self.subcycling_options, - options=recovery_spaces.HDiv_options) + options=recovery_spaces.HDiv_options, + limiter=self.limiters.get(field_name)) ) elif self.equation.space_names[field_name] == 'L2': _transported_fields.append( SSPRK3(self.domain, field_name, subcycling_options=self.subcycling_options, - options=recovery_spaces.DG_options) + options=recovery_spaces.DG_options, + limiter=self.limiters.get(field_name)) ) elif self.equation.space_names[field_name] == 'theta': _transported_fields.append( SSPRK3(self.domain, field_name, subcycling_options=self.subcycling_options, - options=recovery_spaces.theta_options) + options=recovery_spaces.theta_options, + limiter=self.limiters.get(field_name)) ) return _transported_fields