Merge pull request #533 from firedrakeproject/TBendall/enhance_prescr…

…ibed_transport

Reorganise prescribed transport

gusto/timestepping/split_timestepper.py

@@ -87,15 +87,18 @@ class SplitPrescribedTransport(Timestepper):
     velocity as a prognostic variable.
     """
 
-    def __init__(self, equation, scheme, io, spatial_methods=None,
-                 physics_schemes=None,
-                 prescribed_transporting_velocity=None):
+    def __init__(self, equation, scheme, io, prescribed_transporting_velocity,
+                 spatial_methods=None, physics_schemes=None):
         """
         Args:
             equation (:class:`PrognosticEquation`): the prognostic equation
             scheme (:class:`TimeDiscretisation`): the scheme to use to timestep
                 the prognostic equation
             io (:class:`IO`): the model's object for controlling input/output.
+            prescribed_transporting_velocity: (bool): Whether a time-varying
+                transporting velocity will be defined. If True, this will
+                require the transporting velocity to be setup by calling either
+                the `setup_prescribed_expr` or `setup_prescribed_apply` methods.
             spatial_methods (iter,optional): a list of objects describing the
                 methods to use for discretising transport or diffusion terms
                 for each transported/diffused variable. Defaults to None,
@@ -131,12 +134,10 @@ def __init__(self, equation, scheme, io, spatial_methods=None,
             apply_bcs = False
             phys_scheme.setup(equation, apply_bcs, parametrisation.label)
 
-        if prescribed_transporting_velocity is not None:
-            self.velocity_projection = Projector(
-                prescribed_transporting_velocity(self.t),
-                self.fields('u'))
-        else:
-            self.velocity_projection = None
+        self.prescribed_transport_velocity = prescribed_transporting_velocity
+        self.is_velocity_setup = not self.prescribed_transport_velocity
+        self.velocity_projection = None
+        self.velocity_apply = None
 
     @property
     def transporting_velocity(self):
@@ -153,10 +154,73 @@ def setup_scheme(self):
         if self.io.output.log_courant:
             self.scheme.courant_max = self.io.courant_max
 
+    def setup_prescribed_expr(self, expr_func):
+        """
+        Sets up the prescribed transporting velocity, through a python function
+        which has time as an argument, and returns a `ufl.Expr`. This allows the
+        velocity to be updated with time.
+
+        Args:
+            expr_func (func): a python function with a single argument that
+                represents the model time, and returns a `ufl.Expr`.
+        """
+
+        if self.is_velocity_setup:
+            raise RuntimeError('Prescribed velocity already set up!')
+
+        self.velocity_projection = Projector(
+            expr_func(self.t), self.fields('u')
+        )
+
+        self.is_velocity_setup = True
+
+    def setup_prescribed_apply(self, apply_method):
+        """
+        Sets up the prescribed transporting velocity, through a python function
+        which has time as an argument. This function will perform the evaluation
+        of the transporting velocity.
+
+        Args:
+            expr_func (func): a python function with a single argument that
+                represents the model time, and performs the evaluation of the
+                transporting velocity.
+        """
+
+        if self.is_velocity_setup:
+            raise RuntimeError('Prescribed velocity already set up!')
+        self.velocity_apply = apply_method
+        self.is_velocity_setup = True
+
+    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
+        """
+
+        # Throw an error if no transporting velocity has been set up
+        if self.prescribed_transport_velocity and not self.is_velocity_setup:
+            raise RuntimeError(
+                'A time-varying prescribed velocity is required. This must be '
+                + 'set up through calling either the setup_prescribed_expr or '
+                + 'setup_prescribed_apply routines.')
+
+        # It's best to have evaluated the velocity before we start
+        if self.velocity_projection is not None:
+            self.velocity_projection.project()
+        if self.velocity_apply is not None:
+            self.velocity_apply(self.t)
+
+        super().run(t, tmax, pick_up=pick_up)
+
     def timestep(self):
 
         if self.velocity_projection is not None:
             self.velocity_projection.project()
+        if self.velocity_apply is not None:
+            self.velocity_apply(self.t)
 
         super().timestep()
 

gusto/timestepping/timestepper.py

@@ -312,22 +312,20 @@ class PrescribedTransport(Timestepper):
     """
     Implements a timeloop with a prescibed transporting velocity.
     """
-    def __init__(self, equation, scheme, io, transport_method,
-                 physics_parametrisations=None,
-                 prescribed_transporting_velocity=None):
+    def __init__(self, equation, scheme, io, prescribed_transporting_velocity,
+                 transport_method, physics_parametrisations=None):
         """
         Args:
             equation (:class:`PrognosticEquation`): the prognostic equation
             scheme (:class:`TimeDiscretisation`): the scheme to use to timestep
                 the prognostic equation.
+            io (:class:`IO`): the model's object for controlling input/output.
+            prescribed_transporting_velocity: (bool): Whether a time-varying
+                transporting velocity will be defined. If True, this will
+                require the transporting velocity to be setup by calling either
+                the `setup_prescribed_expr` or `setup_prescribed_apply` methods.
             transport_method (:class:`TransportMethod`): describes the method
                 used for discretising the transport term.
-            io (:class:`IO`): the model's object for controlling input/output.
-            physics_schemes: (list, optional): a list of :class:`Physics` and
-                :class:`TimeDiscretisation` options describing physical
-                parametrisations and timestepping schemes to use for each.
-                Timestepping schemes for physics must be explicit. Defaults to
-                None.
             physics_parametrisations: (iter, optional): an iterable of
                 :class:`PhysicsParametrisation` objects that describe physical
                 parametrisations to be included to add to the equation. They can
@@ -344,12 +342,10 @@ def __init__(self, equation, scheme, io, transport_method,
         super().__init__(equation, scheme, io, spatial_methods=transport_methods,
                          physics_parametrisations=physics_parametrisations)
 
-        if prescribed_transporting_velocity is not None:
-            self.velocity_projection = Projector(
-                prescribed_transporting_velocity(self.t),
-                self.fields('u'))
-        else:
-            self.velocity_projection = None
+        self.prescribed_transport_velocity = prescribed_transporting_velocity
+        self.is_velocity_setup = not self.prescribed_transport_velocity
+        self.velocity_projection = None
+        self.velocity_apply = None
 
     @property
     def transporting_velocity(self):
@@ -360,6 +356,43 @@ def setup_fields(self):
         self.fields = StateFields(self.x, self.equation.prescribed_fields,
                                   *self.io.output.dumplist)
 
+    def setup_prescribed_expr(self, expr_func):
+        """
+        Sets up the prescribed transporting velocity, through a python function
+        which has time as an argument, and returns a `ufl.Expr`. This allows the
+        velocity to be updated with time.
+
+        Args:
+            expr_func (func): a python function with a single argument that
+                represents the model time, and returns a `ufl.Expr`.
+        """
+
+        if self.is_velocity_setup:
+            raise RuntimeError('Prescribed velocity already set up!')
+
+        self.velocity_projection = Projector(
+            expr_func(self.t), self.fields('u')
+        )
+
+        self.is_velocity_setup = True
+
+    def setup_prescribed_apply(self, apply_method):
+        """
+        Sets up the prescribed transporting velocity, through a python function
+        which has time as an argument. This function will perform the evaluation
+        of the transporting velocity.
+
+        Args:
+            expr_func (func): a python function with a single argument that
+                represents the model time, and performs the evaluation of the
+                transporting velocity.
+        """
+
+        if self.is_velocity_setup:
+            raise RuntimeError('Prescribed velocity already set up!')
+        self.velocity_apply = apply_method
+        self.is_velocity_setup = True
+
     def run(self, t, tmax, pick_up=False):
         """
         Runs the model for the specified time, from t to tmax
@@ -368,14 +401,31 @@ def run(self, t, tmax, pick_up=False):
             tmax (float): the end time of the run
             pick_up: (bool): specify whether to pick_up from a previous run
         """
+
+        # Throw an error if no transporting velocity has been set up
+        if self.prescribed_transport_velocity and not self.is_velocity_setup:
+            raise RuntimeError(
+                'A time-varying prescribed velocity is required. This must be '
+                + 'set up through calling either the setup_prescribed_expr or '
+                + 'setup_prescribed_apply routines.')
+
         # It's best to have evaluated the velocity before we start
         if self.velocity_projection is not None:
             self.velocity_projection.project()
+        if self.velocity_apply is not None:
+            self.velocity_apply(self.t)
 
         super().run(t, tmax, pick_up=pick_up)
 
     def timestep(self):
+        """
+        Implements the time step, which possibly involves evaluation of the
+        prescribed transporting velocity.
+        """
+
         if self.velocity_projection is not None:
             self.velocity_projection.project()
+        if self.velocity_apply is not None:
+            self.velocity_apply(self.t)
 
         super().timestep()

integration-tests/equations/test_advection_diffusion.py

@@ -36,7 +36,10 @@ def run_advection_diffusion(tmpdir):
     io = IO(domain, output)
 
     # Time stepper
-    stepper = PrescribedTransport(equation, SSPRK3(domain), io, spatial_methods)
+    time_varying_velocity = False
+    stepper = PrescribedTransport(
+        equation, SSPRK3(domain), io, time_varying_velocity, spatial_methods
+    )
 
     # ------------------------------------------------------------------------ #
     # Initial conditions

integration-tests/equations/test_coupled_transport.py

@@ -51,7 +51,10 @@ def test_coupled_transport_scalar(tmpdir, geometry, equation_form1, equation_for
     transport_scheme = SSPRK3(domain)
     transport_method = [DGUpwind(eqn, 'f1'), DGUpwind(eqn, 'f2')]
 
-    timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method)
+    time_varying_velocity = False
+    timestepper = PrescribedTransport(
+        eqn, transport_scheme, setup.io, time_varying_velocity, transport_method
+    )
 
     # Initial conditions
     timestepper.fields("f1").interpolate(setup.f_init)
@@ -101,7 +104,10 @@ def test_conservative_coupled_transport(tmpdir, m_X_space, tracer_setup):
 
     transport_method = [DGUpwind(eqn, 'f1'), DGUpwind(eqn, 'f2')]
 
-    timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method)
+    time_varying_velocity = False
+    timestepper = PrescribedTransport(
+        eqn, transport_scheme, setup.io, time_varying_velocity, transport_method
+    )
 
     # Initial conditions
     timestepper.fields("f1").interpolate(setup.f_init)

integration-tests/equations/test_forced_advection.py

@@ -64,8 +64,11 @@ def run_forced_advection(tmpdir):
     io = IO(domain, output, diagnostic_fields=diagnostic_fields)
 
     # Time Stepper
-    stepper = PrescribedTransport(meqn, RK4(domain), io, transport_method,
-                                  physics_parametrisations=physics_parametrisations)
+    time_varying_velocity = False
+    stepper = PrescribedTransport(
+        meqn, RK4(domain), io, time_varying_velocity, transport_method,
+        physics_parametrisations=physics_parametrisations
+    )
 
     # ------------------------------------------------------------------------ #
     # Initial conditions

integration-tests/model/test_IMEX.py

@@ -34,7 +34,10 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup):
 
     transport_method = DGUpwind(eqn, "f")
 
-    timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method)
+    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)

integration-tests/model/test_conservative_transport_with_physics.py

@@ -58,9 +58,10 @@ def run_conservative_transport_with_physics(dirname):
     physics_schemes = [(SourceSink(eqn, 'ash', basic_expression), SSPRK3(domain))]
 
     # Time stepper
-    stepper = SplitPrescribedTransport(eqn, SSPRK3(domain, increment_form=False),
-                                       io, transport_method,
-                                       physics_schemes=physics_schemes)
+    time_varying_velocity = False
+    stepper = SplitPrescribedTransport(
+        eqn, SSPRK3(domain, increment_form=False), io, time_varying_velocity,
+        transport_method, physics_schemes=physics_schemes)
 
     # ------------------------------------------------------------------------ #
     # Initial conditions

integration-tests/model/test_nc_outputting.py

@@ -96,7 +96,10 @@ def test_nc_outputting(tmpdir, geometry, domain_and_mesh_details):
         diagnostic_fields = [ZonalComponent('u'), MeridionalComponent('u'), RadialComponent('u')]
 
     io = IO(domain, output, diagnostic_fields=diagnostic_fields)
-    stepper = PrescribedTransport(eqn, transport_scheme, io, transport_method)
+    time_varying_velocity = False
+    stepper = PrescribedTransport(
+        eqn, transport_scheme, io, time_varying_velocity, transport_method
+    )
 
     # ------------------------------------------------------------------------ #
     # Initialise fields

integration-tests/model/test_prescribed_transport.py

@@ -28,26 +28,33 @@ def test_prescribed_transport_setup(tmpdir, tracer_setup, timestep_method):
     # Make equation
     eqn = AdvectionEquation(domain, V, "f")
 
-    # Initialise fields
-    def u_evaluation(t):
-        return as_vector([2.0*cos(2*pi*t/setup.tmax),
-                          sin(2*pi*t/setup.tmax)*sin(pi*z)])
-
     transport_scheme = SSPRK3(domain)
+    time_varying_velocity = True
 
     if timestep_method == 'prescribed':
         transport_method = DGUpwind(eqn, 'f')
-        timestepper = PrescribedTransport(eqn, transport_scheme, setup.io,
-                                          transport_method,
-                                          prescribed_transporting_velocity=u_evaluation)
+        timestepper = PrescribedTransport(
+            eqn, transport_scheme, setup.io, time_varying_velocity,
+            transport_method
+        )
+
     elif timestep_method == 'split_prescribed':
         transport_method = [DGUpwind(eqn, 'f')]
-        timestepper = SplitPrescribedTransport(eqn, transport_scheme, setup.io,
-                                               transport_method,
-                                               prescribed_transporting_velocity=u_evaluation)
+        timestepper = SplitPrescribedTransport(
+            eqn, transport_scheme, setup.io, time_varying_velocity,
+            transport_method
+        )
+
     else:
         raise NotImplementedError
 
+    # Initialise fields
+    def u_evaluation(t):
+        return as_vector([2.0*cos(2*pi*t/setup.tmax),
+                          sin(2*pi*t/setup.tmax)*sin(pi*z)])
+
+    timestepper.setup_prescribed_expr(u_evaluation)
+
     # Initial conditions
     timestepper.fields("f").interpolate(setup.f_init)
     timestepper.fields("u").project(u_evaluation(Constant(0.0)))

integration-tests/model/test_sdc.py

@@ -74,7 +74,10 @@ def test_sdc(tmpdir, scheme, tracer_setup):
 
     transport_method = DGUpwind(eqn, 'f')
 
-    timestepper = PrescribedTransport(eqn, scheme, setup.io, transport_method)
+    time_varying_velocity = False
+    timestepper = PrescribedTransport(
+        eqn, scheme, setup.io, time_varying_velocity, transport_method
+    )
 
     # Initial conditions
     timestepper.fields("f").interpolate(setup.f_init)

integration-tests/model/test_time_discretisation.py

@@ -58,7 +58,10 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup):
 
     transport_method = DGUpwind(eqn, 'f')
 
-    timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method)
+    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)