diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index 798aed40..e6cae36c 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -52,6 +52,7 @@ def williamson_5( # ------------------------------------------------------------------------ # element_order = 1 + # ------------------------------------------------------------------------ # # Set up model objects # ------------------------------------------------------------------------ # diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py index ce875894..2c49304a 100644 --- a/gusto/timestepping/semi_implicit_quasi_newton.py +++ b/gusto/timestepping/semi_implicit_quasi_newton.py @@ -3,8 +3,10 @@ and GungHo dynamical cores. """ -from firedrake import (Function, Constant, TrialFunctions, DirichletBC, - LinearVariationalProblem, LinearVariationalSolver) +from firedrake import ( + Function, Constant, TrialFunctions, DirichletBC, div, Interpolator, + LinearVariationalProblem, LinearVariationalSolver +) from firedrake.fml import drop, replace_subject from pyop2.profiling import timed_stage from gusto.core import TimeLevelFields, StateFields @@ -36,8 +38,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, slow_physics_schemes=None, fast_physics_schemes=None, alpha=Constant(0.5), off_centred_u=False, num_outer=2, num_inner=2, accelerator=False, - reference_update_freq=None): - + predictor=None, reference_update_freq=None): """ Args: equation_set (:class:`PrognosticEquationSet`): the prognostic @@ -88,6 +89,15 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, accelerator (bool, optional): Whether to zero non-wind implicit forcings for transport terms in order to speed up solver convergence. Defaults to False. + predictor (str, optional): a single string corresponding to the name + of a variable to transport using the divergence predictor. This + pre-multiplies that variable by (1 - beta*dt*div(u)) before the + transport step, and calculates its transport increment from the + transport of this variable. This can improve the stability of + the time stepper at large time steps, when not using an + advective-then-flux formulation. This is only suitable for the + use on the conservative variable (e.g. depth or density). + Defaults to None, in which case no predictor is used. reference_update_freq (float, optional): frequency with which to update the reference profile with the n-th time level state fields. This variable corresponds to time in seconds, and @@ -100,6 +110,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, self.num_outer = num_outer self.num_inner = num_inner self.alpha = alpha + self.predictor = predictor self.accelerator = accelerator self.reference_update_freq = reference_update_freq self.to_update_ref_profile = False @@ -201,6 +212,14 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, self.forcing = Forcing(equation_set, self.alpha) self.bcs = equation_set.bcs + if self.predictor is not None: + V_DG = equation_set.domain.spaces('DG') + self.predictor_field_in = Function(V_DG) + div_factor = Constant(1.0) - (Constant(1.0) - self.alpha)*self.dt*div(self.x.n('u')) + self.predictor_interpolator = Interpolator( + self.x.star(predictor)*div_factor, self.predictor_field_in + ) + def _apply_bcs(self): """ Set the zero boundary conditions in the velocity. @@ -263,6 +282,33 @@ def copy_active_tracers(self, x_in, x_out): for name in self.tracers_to_copy: x_out(name).assign(x_in(name)) + def transport_field(self, name, scheme, xstar, xp): + """ + Performs the transport of a field in xstar, placing the result in xp. + + Args: + name (str): the name of the field to be transported. + scheme (:class:`TimeDiscretisation`): the time discretisation used + for the transport. + xstar (:class:`Fields`): the collection of state fields to be + transported. + xp (:class:`Fields`): the collection of state fields resulting from + the transport. + """ + + if name == self.predictor: + # Pre-multiply this variable by (1 - dt*beta*div(u)) + V = xstar(name).function_space() + field_out = Function(V) + self.predictor_interpolator.interpolate() + scheme.apply(field_out, self.predictor_field_in) + + # xp is xstar plus the increment from the transported predictor + xp(name).assign(xstar(name) + field_out - self.predictor_field_in) + else: + # Standard transport + scheme.apply(xp(name), xstar(name)) + def update_reference_profiles(self): """ Updates the reference profiles and if required also updates them in the @@ -324,7 +370,7 @@ def timestep(self): for name, scheme in self.active_transport: logger.info(f'Semi-implicit Quasi Newton: Transport {outer}: {name}') # transports a field from xstar and puts result in xp - scheme.apply(xp(name), xstar(name)) + self.transport_field(name, scheme, xstar, xp) # Fast physics ----------------------------------------------------- x_after_fast(self.field_name).assign(xp(self.field_name))