D witt/tr bdf2 model

gusto/diagnostics/diagnostics.py

@@ -185,7 +185,7 @@ def setup(self, domain, state_fields, space=None):
             if not hasattr(domain.spaces, space.name):
                 domain.spaces.add_space(space.name, space)
 
-            self.field = state_fields(self.name, space=space, dump=self.to_dump, pick_up=False)
+            self.field = state_fields(self.name, space=space, dump=self.to_dump, pick_up=True)
 
             if self.method != 'solve':
                 assert self.expr is not None, \

gusto/model/model.py

@@ -1,10 +1,11 @@
 from abc import ABCMeta, abstractmethod, abstractproperty
+from firedrake import sqrt
 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)
-from gusto.timestepping import SemiImplicitQuasiNewton
+from gusto.timestepping import SemiImplicitQuasiNewton, TRBDF2QuasiNewton
 
 
 class ModelBase(object, metaclass=ABCMeta):
@@ -97,6 +98,218 @@ def run(self, t, tmax, pick_up=False):
         self.stepper.run(t=t, tmax=tmax, pick_up=pick_up)
 
 
+class TRBDF2ModelBase(ModelBase):
+    """
+    Base for model classes using SIQN. Child classes should define the
+    standard transport and diffusion schemes and methods.
+    """
+
+    @abstractproperty
+    def diffusion_methods(self):
+        """
+        Needs to return a list of instances of the diffusion discretisation
+        class for each of the prognostic fields in the model equation set that
+        has diffusion options defined.
+        """
+        pass
+
+    @abstractproperty
+    def diffusion_schemes(self):
+        """
+        Need to return a list of instances of the timestepping scheme
+        class to be used for each of the prognostic fields in the model
+        equation set that has diffusion options defined.
+        """
+        pass
+
+    @abstractproperty
+    def transported_fields(self):
+        """
+        Needs to return a list of instances of the timestepping scheme class
+        to be used for each of the prognostic fields in the model equation set.
+        """
+        pass
+
+    @abstractproperty
+    def transport_methods(self):
+        """
+        Needs to return a list of instances of the transport method class
+        for each of the prognostic fields in the model equation set.
+        """
+        pass
+
+    @property
+    def tr_tau_values(self):
+        _tau = {}
+        for field_name in self.equation.field_names:
+            if field_name != "u":
+                _tau[field_name] = 2 * self.gamma
+        return _tau
+
+    @property
+    def bdf_tau_values(self):
+        _tau = {}
+        for field_name in self.equation.field_names:
+            if field_name != "u":
+                _tau[field_name] = 1.0
+        return _tau
+
+    def setup(self, output, **kwargs):
+        """
+        Args:
+            output (:class:`OutputParameters`): provides parameters
+                controlling output
+
+        Kwargs:
+            diagnostic_fields (list, optional): an iterable of `DiagnosticField`
+                objects. Defaults to None.
+            subcycling_options (:class:`SubcyclingOptions`, optional): an
+                object containing options for subcycling the time
+                discretisation. Defaults to None.
+
+            Remaining kwargs are passed straight through to the stepper - see
+            there for documentation.
+        """
+        self.gamma = 1 - sqrt(2) / 2
+        diagnostic_fields = kwargs.pop("diagnostic_fields", None)
+        io = IO(self.domain, output, diagnostic_fields=diagnostic_fields)
+
+        self.subcycling_options = kwargs.pop("subcycling_options", None)
+        self.limiters = kwargs.pop("limiters", {})
+        self.stepper = TRBDF2QuasiNewton(
+            self.equation, io, self.transported_fields,
+            spatial_methods=self.transport_methods+self.diffusion_methods,
+            diffusion_schemes=self.diffusion_schemes,
+            tau_values_tr=self.tr_tau_values,
+            tau_values_bdf=self.bdf_tau_values,
+            **kwargs
+        )
+
+
+class TRBDF2Model(TRBDF2ModelBase):
+    """
+    TRBDF2 model class encapsulating the best settings for next-to-lowest order
+    methods.
+    """
+    def __init__(self, mesh, dt, parameters, equations,
+                 family=None, **kwargs):
+        """
+        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
+            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). Defaults to None.
+
+        Kwargs:
+            Kwargs are passed straight through to the equation class; see
+            those classes for full documentation. Some common kwargs are
+            listed below.
+
+            u_transport_option (str, optional): specifies the transport term
+                used for the velocity equation. Supported options are:
+                'vector_invariant_form', 'vector_advection_form', and
+                'circulation_form'. Defaults to 'vector_invariant_form'.
+            diffusion_options (iterable, optional): iterable of
+                ``(field_name, diffusion_parameters)`` pairs where
+                diffusion_parameters is a :class:`DiffusionParameters`
+                object specifying the diffusion parameters to be applied
+                to the field field_name. 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.
+            active_tracers (list, optional): a list of `ActiveTracer` objects
+                that encode the metadata for any active tracers to be included
+                in the equations. Defaults to None.
+        """
+        super().__init__(mesh, dt, parameters, equations, family=family,
+                         element_order=1, **kwargs)
+
+    @property
+    def diffusion_methods(self):
+        """
+        Returns a list of instances of the InteriorPenaltyDiffusion
+        class for each of the prognostic fields in the model equation set that
+        has diffusion options defined.
+        """
+        _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):
+        """
+        Returns a list of instances of the Backward Euler timestepping scheme
+        class for each of the prognostic fields in the model equation set that
+        has diffusion options defined.
+        """
+        _diffusion_schemes = []
+        for field, _ in self.diffusion_options:
+            _diffusion_schemes.append(
+                BackwardEuler(self.domain, field)
+            )
+        return _diffusion_schemes
+
+    @property
+    def transported_fields(self):
+        """
+        Returns a list of instances of the SSPRK3 timestepping scheme class
+        for each of the prognostic fields in the model equation set.
+        """
+        _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,
+                           subcycling_options=self.subcycling_options,
+                           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(),
+                           limiter=self.limiters.get(field_name))
+                )
+            else:
+                _transported_fields.append(
+                    SSPRK3(
+                        self.domain, field_name,
+                        subcycling_options=self.subcycling_options,
+                        limiter=self.limiters.get(field_name))
+                )
+        return _transported_fields
+
+    @property
+    def transport_methods(self):
+        """
+        Returns a list of instances of the DGUpwind transport method class
+        for each of the prognostic fields in the model equation set.
+        """
+        _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
+
+
 class SIQNModelBase(ModelBase):
     """
     Base for model classes using SIQN. Child classes should define the

gusto/time_discretisation/explicit_runge_kutta.py

@@ -126,7 +126,7 @@ def __init__(self, domain, butcher_matrix, field_name=None,
         self.nStages = int(np.shape(self.butcher_matrix)[0])
         self.rk_formulation = rk_formulation
 
-    def setup(self, equation, apply_bcs=True, *active_labels):
+    def setup(self, equation, apply_bcs=True, *active_labels, dt_scale=None,):
         """
         Set up the time discretisation based on the equation.
 
@@ -135,7 +135,7 @@ def setup(self, equation, apply_bcs=True, *active_labels):
             *active_labels (:class:`Label`): labels indicating which terms of
                 the equation to include.
         """
-        super().setup(equation, apply_bcs, *active_labels)
+        super().setup(equation, apply_bcs, *active_labels, dt_scale=dt_scale)
 
         if self.rk_formulation == RungeKuttaFormulation.predictor:
             self.field_i = [Function(self.fs) for _ in range(self.nStages+1)]
@@ -448,7 +448,6 @@ def apply_cycle(self, x_out, x_in):
             self.limiter.apply(x_in)
 
         self.x1.assign(x_in)
-
         for i in range(self.nStages):
             self.solve_stage(x_in, i)
         x_out.assign(self.x1)
@@ -501,7 +500,8 @@ def __init__(
                          rk_formulation=rk_formulation,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options,
-                         augmentation=augmentation)
+                         augmentation=augmentation
+                         )
 
 
 class SSPRK2(ExplicitRungeKutta):

gusto/time_discretisation/implicit_runge_kutta.py

@@ -90,7 +90,7 @@ def __init__(self, domain, butcher_matrix, field_name=None,
         self.nStages = int(np.shape(self.butcher_matrix)[1])
         self.rk_formulation = rk_formulation
 
-    def setup(self, equation, apply_bcs=True, *active_labels):
+    def setup(self, equation, apply_bcs=True, *active_labels, dt_scale=None):
         """
         Set up the time discretisation based on the equation.
 
@@ -100,7 +100,7 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                 the equation to include.
         """
 
-        super().setup(equation, apply_bcs, *active_labels)
+        super().setup(equation, apply_bcs, *active_labels, dt_scale=dt_scale)
 
         self.k = [Function(self.fs) for i in range(self.nStages)]
 

gusto/time_discretisation/time_discretisation.py

@@ -1,4 +1,4 @@
-u"""
+"""
 Objects for discretising time derivatives.
 
 Time discretisation objects discretise ∂y/∂t = F(y), for variable y, time t and
@@ -9,9 +9,9 @@
 from functools import cached_property
 import math
 
 from firedrake import (Function, TestFunction, TestFunctions, DirichletBC,
                        NonlinearVariationalProblem, NonlinearVariationalSolver,
-                       FunctionSpace)
+                       FunctionSpace, dot)
 from firedrake.fml import (replace_subject, replace_test_function, Term,
                            all_terms, drop)
 from firedrake.formmanipulation import split_form
@@ -91,17 +91,17 @@
         self.domain = domain
         self.field_name = field_name
         self.equation = None
-        R = FunctionSpace(domain.mesh, "R", 0)
-        self.dt = Function(R, val=0.0)
+        self.R = FunctionSpace(domain.mesh, "R", 0)
+        self.dt = Function(self.R, val=0.0)
         self.dt.assign(domain.dt)
-        self.original_dt = Function(R, val=0.0)
+        self.original_dt = Function(self.R, val=0.0)
         self.original_dt.assign(self.dt)
         self.options = options
         self.limiter = limiter
         self.courant_max = None
         self.augmentation = augmentation
         self.subcycling_options = subcycling_options
-
+        
         if self.subcycling_options is not None:
             self.subcycling_options.check_options()
 
@@ -144,7 +144,7 @@
         else:
             self.solver_parameters = solver_parameters
 
-    def setup(self, equation, apply_bcs=True, *active_labels):
+    def setup(self, equation, apply_bcs=True, *active_labels, dt_scale=None,):
         """
         Set up the time discretisation based on the equation.
 
@@ -158,6 +158,11 @@
         self.equation = equation
         self.residual = equation.residual
 
+        if dt_scale is not None:
+            dt_scale_func = Function(self.R).assign(dt_scale)
+            self.dt.assign(self.dt * dt_scale_func)
+            self.original_dt.assign(self.original_dt * dt_scale_func)
+
         if self.field_name is not None and hasattr(equation, "field_names"):
             if isinstance(self.field_name, list):
                 # Multiple fields are being solved for simultaneously.
@@ -494,9 +499,10 @@
                          subcycling_options=subcycling_options,
                          solver_parameters=solver_parameters,
                          limiter=limiter, options=options,
-                         augmentation=augmentation)
+                         augmentation=augmentation
+                         )
 
-    def setup(self, equation, apply_bcs=True, *active_labels):
+    def setup(self, equation, apply_bcs=True,  *active_labels, dt_scale=None):
         """
         Set up the time discretisation based on the equation.
 
@@ -507,7 +513,7 @@
             *active_labels (:class:`Label`): labels indicating which terms of
                 the equation to include.
         """
-        super().setup(equation, apply_bcs, *active_labels)
+        super().setup(equation, apply_bcs, *active_labels, dt_scale=dt_scale)
 
         # get default solver options if none passed in
         self.solver_parameters.update(mass_parameters(