From 60f58b6c42d812f565d257151480e458dd0cff98 Mon Sep 17 00:00:00 2001 From: Joe Wallwork Date: Mon, 18 Dec 2023 07:50:09 +0000 Subject: [PATCH] Temporarily revert to sphinx 4.5.0 --- _modules/index.html | 19 +- _modules/thetis/assembledschur.html | 17 +- _modules/thetis/callback.html | 137 +- _modules/thetis/configuration.html | 71 +- _modules/thetis/coordsys.html | 209 +- _modules/thetis/coupled_timeintegrator.html | 78 +- .../thetis/coupled_timeintegrator_2d.html | 62 +- _modules/thetis/diagnostics.html | 136 +- _modules/thetis/equation.html | 42 +- _modules/thetis/exner_eq.html | 22 +- _modules/thetis/exporter.html | 96 +- _modules/thetis/implicitexplicit.html | 38 +- _modules/thetis/interpolation.html | 111 +- _modules/thetis/inversion_tools.html | 486 +- _modules/thetis/limiter.html | 41 +- _modules/thetis/log.html | 14 +- _modules/thetis/momentum_eq.html | 38 +- _modules/thetis/optimisation.html | 144 +- _modules/thetis/options.html | 178 +- _modules/thetis/rungekutta.html | 110 +- _modules/thetis/sediment_eq_2d.html | 29 +- _modules/thetis/sediment_model.html | 39 +- _modules/thetis/shallowwater_eq.html | 98 +- _modules/thetis/solver.html | 225 +- _modules/thetis/solver2d.html | 376 +- _modules/thetis/stability_functions.html | 64 +- _modules/thetis/timeintegrator.html | 193 +- _modules/thetis/timezone.html | 22 +- _modules/thetis/tracer_eq.html | 26 +- _modules/thetis/tracer_eq_2d.html | 304 +- _modules/thetis/turbines.html | 348 +- _modules/thetis/turbulence.html | 88 +- _modules/thetis/utility.html | 285 +- _modules/thetis/utility3d.html | 108 +- _sources/model_options_2d.rst.txt | 4 +- _sources/model_options_3d.rst.txt | 2 +- _static/basic.css | 73 +- _static/doctools.js | 420 +- _static/documentation_options.js | 5 +- _static/jquery.js | 4 +- _static/language_data.js | 102 +- _static/searchtools.js | 831 ++- animationgrid.html | 24 +- contact.html | 26 +- demos/demo_2d_channel.py.html | 22 +- demos/demo_2d_channel_bnd.py.html | 22 +- demos/demo_2d_multiple_tracers.py.html | 22 +- demos/demo_2d_north_sea.py.html | 68 +- demos/demo_2d_tracer.py.html | 33 +- demos/demo_3d_channel.py.html | 22 +- documentation.html | 30 +- download.html | 28 +- eos_options.html | 24 +- field_documentation.html | 24 +- field_list.html | 20 +- funding.html | 27 +- genindex.html | 115 +- index.html | 26 +- model_formulation_2d.html | 43 +- model_formulation_3d.html | 57 +- model_options_2d.html | 26 +- model_options_3d.html | 24 +- objects.inv | Bin 9829 -> 9604 bytes outputs_and_visu.html | 30 +- publications.html | 36 +- py-modindex.html | 23 +- search.html | 18 +- searchindex.js | 2 +- sediment_formulation_2d.html | 34 +- sediment_model_options.html | 22 +- team.html | 4 +- teamgrid.html | 26 +- thetis.html | 4598 ++++++++--------- tracer_formulation_2d.html | 40 +- turbulence_options.html | 24 +- 75 files changed, 6097 insertions(+), 5138 deletions(-) diff --git a/_modules/index.html b/_modules/index.html index 2fc38f1..516823d 100644 --- a/_modules/index.html +++ b/_modules/index.html @@ -1,15 +1,17 @@ + - + Overview: module code — Thetis 0+untagged.2009.gd7af522 documentation - - - - - + + + + + + @@ -57,7 +56,7 @@

Source code for thetis.assembledschur

 
 
 
[docs]class AssembledSchurPC(PCBase):
-    """
+    """
     Preconditioner for the Schur complement
 
     The preconditioner matrix is assembled by explicitly matrix multiplying
@@ -83,7 +82,7 @@ 
Source code for thetis.assembledschur
 
         test, trial = a.arguments()
         W = test.function_space()
-        V, Q = W.split()
+        V, Q = W.subfunctions
         v = TestFunction(V)
         u = TrialFunction(V)
         mass = dot(v, u)*dx
@@ -151,8 +150,8 @@ 
Source code for thetis.assembledschur
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/callback.html b/_modules/thetis/callback.html
index 9ecc343..fbfc9a9 100644
--- a/_modules/thetis/callback.html
+++ b/_modules/thetis/callback.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.callback — Thetis 0+untagged.1888.gc4e776a documentation
+    thetis.callback — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -65,7 +65,7 @@ 
Source code for thetis.callback
 
 
 
[docs]class CallbackManager(defaultdict):
-    """
+    """
     Stores callbacks in different categories and provides methods for
     evaluating them.
 
@@ -91,7 +91,7 @@ 
Source code for thetis.callback
         super(CallbackManager, self).__init__(OrderedDict)
 
 
[docs]    def add(self, callback, mode):
-        """
+        """
         Add a callback under the given mode
 
         :arg callback: a :class:`.DiagnosticCallback` object
@@ -101,7 +101,7 @@ 
Source code for thetis.callback
         self[mode][key] = callback

 
 
[docs]    def evaluate(self, mode, index=None):
-        """
+        """
         Evaluate all callbacks registered under the given mode
 
         :arg str mode: evaluate all callbacks under this mode
@@ -113,23 +113,27 @@ 
Source code for thetis.callback
 
 
 
[docs]class DiagnosticHDF5(object):
-    """
+    """
     A HDF5 file for storing diagnostic time series arrays.
     """
     @PETSc.Log.EventDecorator("thetis.DiagnosticHDF5.__init__")
     def __init__(self, filename, varnames, array_dim=1, attrs=None,
-                 comm=COMM_WORLD, new_file=True, dtype='d',
-                 include_time=True):
-        """
+                 var_attrs=None, comm=COMM_WORLD, new_file=True,
+                 dtype='d', include_time=True):
+        """
         :arg str filename: Full filename of the HDF5 file.
         :arg varnames: List of variable names that the diagnostic callback
             provides
         :kwarg array_dim: Dimension of the output array.
             Can be a tuple for multi-dimensional output. Use "1" for scalars.
-        :kwarg dict attrs: Additional attributes to be saved in the hdf5 file.
+        :kwarg dict attrs: Global attributes to be saved in the hdf5 file.
+        :kwarg dict var_attrs: nested dict of variable specific attributes,
+             e.g. {'time': {'units': 'seconds since 1950-01-01'}}
         :kwarg comm: MPI communicator
         :kwarg bool new_file: Define whether to create a new hdf5 file or
             append to an existing one (if any)
+        :kwarg dtype: array datatype
+        :kwarg include_time: whether to include time array in the file
         """
         self.comm = comm
         self.filename = filename
@@ -141,8 +145,10 @@ 
Source code for thetis.callback
             # create empty file with correct datasets
             with h5py.File(filename, 'w') as hdf5file:
                 if include_time:
-                    hdf5file.create_dataset('time', (0, 1),
-                                            maxshape=(None, 1), dtype=dtype)
+                    ds = hdf5file.create_dataset(
+                        'time', (0, 1), maxshape=(None, 1), dtype=dtype)
+                    if var_attrs is not None and 'time' in var_attrs:
+                        ds.attrs.update(var_attrs['time'])
                 dim_list = array_dim
                 if isinstance(dim_list, tuple):
                     dim_list = list(dim_list)
@@ -151,20 +157,22 @@ 
Source code for thetis.callback
                 shape = tuple([0] + dim_list)
                 max_shape = tuple([None] + dim_list)
                 for var in self.varnames:
-                    hdf5file.create_dataset(var, shape,
-                                            maxshape=max_shape, dtype=dtype)
+                    ds = hdf5file.create_dataset(
+                        var, shape, maxshape=max_shape, dtype=dtype)
+                    if var_attrs is not None and var in var_attrs:
+                        ds.attrs.update(var_attrs[var])
                 if attrs is not None:
                     hdf5file.attrs.update(attrs)
 
     def _expand_array(self, hdf5file, varname):
-        """Expands array varname by 1 entry"""
+        """Expands array varname by 1 entry"""
         arr = hdf5file[varname]
         new_shape = list(arr.shape)
         new_shape[0] += 1
         arr.resize(tuple(new_shape))
 
     def _expand(self, hdf5file):
-        """Expands data arrays by 1 entry"""
+        """Expands data arrays by 1 entry"""
         for var in self.varnames:
             self._expand_array(hdf5file, var)
         if self.include_time:
@@ -175,7 +183,7 @@ 
Source code for thetis.callback
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.DiagnosticHDF5.export")
     def export(self, variables, time=None, index=None):
-        """
+        """
         Appends a new entry of (time, variables) to the file.
 
         The HDF5 is updated immediately.
@@ -205,7 +213,7 @@ 
Source code for thetis.callback
 
 
 
[docs]class DiagnosticCallback(ABC):
-    """
+    """
     A base class for all Callback classes
     """
 
@@ -217,13 +225,13 @@ 
Source code for thetis.callback
                  hdf5_dtype='d',
                  start_time=None,
                  end_time=None):
-        """
+        """
         :arg solver_obj: Thetis solver object
         :kwarg str outputdir: Custom directory where hdf5 files will be stored.
             By default solver's output directory is used.
         :kwarg array_dim: Dimension of the output array.
             Can be a tuple for multi-dimensional output. Use "1" for scalars.
-        :kwarg dict attrs: Additional attributes to be saved in the hdf5 file.
+        :kwarg dict attrs: Global attributes to be saved in the hdf5 file.
         :kwarg bool export_to_hdf5: If True, diagnostics will be stored in hdf5
             format
         :kwarg bool append_to_log: If True, callback output messages will be
@@ -234,10 +242,13 @@ 
Source code for thetis.callback
         :kwarg start_time: Optional start time for callback evaluation
         :kwarg end_time: Optional end time for callback evaluation
         """
+        if attrs is None:
+            attrs = {}
         self.solver_obj = solver_obj
         self.outputdir = outputdir or self.solver_obj.options.output_directory
         self.array_dim = array_dim
         self.attrs = attrs
+        self.var_attrs = {}
         self.append_to_hdf5 = export_to_hdf5
         self.append_to_log = append_to_log
         self.hdf5_dtype = hdf5_dtype
@@ -247,8 +258,13 @@ 
Source code for thetis.callback
         self.start_time = start_time or -numpy.inf
         self.end_time = end_time or numpy.inf
 
+        init_date = self.solver_obj.options.simulation_initial_date
+        if init_date is not None and include_time:
+            time_units = 'seconds since ' + init_date.isoformat()
+            self.var_attrs['time'] = {'units': time_units}
+
 
[docs]    def set_write_mode(self, mode):
-        """
+        """
         Define whether to create a new hdf5 file or append to an existing one
 
         :arg str mode: Either 'create' (default) or 'append'
@@ -257,7 +273,7 @@ 
Source code for thetis.callback
         self._create_new_file = mode == 'create'

 
     def _create_hdf5_file(self):
-        """
+        """
         Creates an empty hdf5 file with correct datasets.
         """
         if self.append_to_hdf5:
@@ -269,23 +285,24 @@ 
Source code for thetis.callback
                                                array_dim=self.array_dim,
                                                new_file=self._create_new_file,
                                                attrs=self.attrs,
+                                               var_attrs=self.var_attrs,
                                                comm=comm, dtype=self.hdf5_dtype,
                                                include_time=self.include_time)
         self._hdf5_initialized = True
 
     @abstractproperty
     def name(self):
-        """The name of the diagnostic"""
+        """The name of the diagnostic"""
         pass
 
     @abstractproperty
     def variable_names(self):
-        """Names of all scalar values"""
+        """Names of all scalar values"""
         pass
 
     @abstractmethod
     def __call__(self):
-        """
+        """
         Evaluate the diagnostic value.
 
         .. note::
@@ -295,7 +312,7 @@ 
Source code for thetis.callback
 
 
[docs]    @abstractmethod
     def message_str(self, *args):
-        """
+        """
         A string representation.
 
         :arg args: If provided, these will be the return value from
@@ -304,7 +321,7 @@ 
Source code for thetis.callback
         return "{} diagnostic".format(self.name)

 
 
[docs]    def push_to_log(self, time, args):
-        """
+        """
         Push callback status message to log
 
         :arg time: time stamp of entry
@@ -313,7 +330,7 @@ 
Source code for thetis.callback
         print_output(self.message_str(*args))

 
 
[docs]    def push_to_hdf5(self, time, args, index=None):
-        """
+        """
         Append values to HDF5 file.
 
         :arg time: time stamp of entry
@@ -324,7 +341,7 @@ 
Source code for thetis.callback
         self.hdf_exporter.export(args, time=time, index=index)

 
 
[docs]    def evaluate(self, index=None):
-        """
+        """
         Evaluates callback and pushes values to log and hdf file (if enabled)
         """
         time = self.solver_obj.simulation_time
@@ -338,11 +355,11 @@ 
Source code for thetis.callback
 
 
 
[docs]class ScalarConservationCallback(DiagnosticCallback):
-    """Base class for callbacks that check conservation of a scalar quantity"""
+    """Base class for callbacks that check conservation of a scalar quantity"""
     variable_names = ['integral', 'relative_difference']
 
     def __init__(self, scalar_callback, solver_obj, **kwargs):
-        """
+        """
         Creates scalar conservation check callback object
 
         :arg scalar_callback: Python function that takes the solver object as
@@ -368,11 +385,11 @@ 
Source code for thetis.callback
 
 
 
[docs]class VolumeConservation3DCallback(ScalarConservationCallback):
-    """Checks conservation of 3D volume (volume of 3D mesh)"""
+    """Checks conservation of 3D volume (volume of 3D mesh)"""
     name = 'volume3d'
 
     def __init__(self, solver_obj, **kwargs):
-        """
+        """
         :arg solver_obj: Thetis solver object
         :arg kwargs: any additional keyword arguments, see
             :class:`.DiagnosticCallback`.
@@ -383,11 +400,11 @@ 
Source code for thetis.callback
 
 
 
[docs]class VolumeConservation2DCallback(ScalarConservationCallback):
-    """Checks conservation of 2D volume (integral of water elevation field)"""
+    """Checks conservation of 2D volume (integral of water elevation field)"""
     name = 'volume2d'
 
     def __init__(self, solver_obj, **kwargs):
-        """
+        """
         :arg solver_obj: Thetis solver object
         :arg kwargs: any additional keyword arguments, see
             :class:`.DiagnosticCallback`.
@@ -400,7 +417,7 @@ 
Source code for thetis.callback
 
 
 
[docs]class TracerMassConservation2DCallback(ScalarConservationCallback):
-    """
+    """
     Checks conservation of depth-averaged tracer mass
 
     Depth-averaged tracer mass is defined as the integral of 2D tracer
@@ -409,7 +426,7 @@ 
Source code for thetis.callback
     name = 'tracer mass'
 
     def __init__(self, tracer_name, solver_obj, **kwargs):
-        """
+        """
         :arg tracer_name: Name of the tracer. Use canonical field names as in
             :class:`.FieldDict`.
         :arg solver_obj: Thetis solver object
@@ -425,13 +442,13 @@ 
Source code for thetis.callback
 
 
 
[docs]class ConservativeTracerMassConservation2DCallback(ScalarConservationCallback):
-    """
+    """
     Checks conservation of conservative tracer mass which is depth integrated.
     """
     name = 'tracer mass'
 
     def __init__(self, tracer_name, solver_obj, **kwargs):
-        """
+        """
         :arg tracer_name: Name of the tracer. Use canonical field names as in
             :class:`.FieldDict`.
         :arg solver_obj: Thetis solver object
@@ -448,11 +465,11 @@ 
Source code for thetis.callback
 
 
 
[docs]class TracerMassConservationCallback(ScalarConservationCallback):
-    """Checks conservation of total tracer mass"""
+    """Checks conservation of total tracer mass"""
     name = 'tracer mass'
 
     def __init__(self, tracer_name, solver_obj, **kwargs):
-        """
+        """
         :arg tracer_name: Name of the tracer. Use canonical field names as in
             :class:`.FieldDict`.
         :arg solver_obj: Thetis solver object
@@ -467,11 +484,11 @@ 
Source code for thetis.callback
 
 
 
[docs]class MinMaxConservationCallback(DiagnosticCallback):
-    """Base class for callbacks that check conservation of a minimum/maximum"""
+    """Base class for callbacks that check conservation of a minimum/maximum"""
     variable_names = ['min_value', 'max_value', 'undershoot', 'overshoot']
 
     def __init__(self, minmax_callback, solver_obj, **kwargs):
-        """
+        """
         :arg minmax_callback: Python function that takes the solver object as
             an argument and returns a (min, max) value tuple
         :arg solver_obj: Thetis solver object
@@ -496,11 +513,11 @@ 
Source code for thetis.callback
 
 
 
[docs]class TracerOvershootCallBack(MinMaxConservationCallback):
-    """Checks overshoots of the given tracer field."""
+    """Checks overshoots of the given tracer field."""
     name = 'tracer overshoot'
 
     def __init__(self, tracer_name, solver_obj, **kwargs):
-        """
+        """
         :arg tracer_name: Name of the tracer. Use canonical field names as in
             :class:`.FieldDict`.
         :arg solver_obj: Thetis solver object
@@ -519,7 +536,7 @@ 
Source code for thetis.callback
 
 
 
[docs]class DetectorsCallback(DiagnosticCallback):
-    """
+    """
     Callback that evaluates the specified fields at the specified locations
     """
     def __init__(self, solver_obj,
@@ -528,7 +545,7 @@ 
Source code for thetis.callback
                  name,
                  detector_names=None,
                  **kwargs):
-        """
+        """
         :arg solver_obj: Thetis solver object
         :arg detector_locations: List of x, y locations in which fields are to
             be interpolated.
@@ -574,7 +591,7 @@ 
Source code for thetis.callback
         return self.detector_names
 
     def _values_per_field(self, values):
-        """
+        """
         Given all values evaulated in a detector location, return the values per field"""
         i = 0
         result = []
@@ -593,7 +610,7 @@ 
Source code for thetis.callback
         return self.solver_obj.fields[field_name](self.detector_locations)
 
     def __call__(self):
-        """
+        """
         Evaluate all current fields in all detector locations
 
         Returns a ndetectors x ndims array, where ndims is the sum of the
@@ -608,7 +625,7 @@ 
Source code for thetis.callback
 
 
 
[docs]class AccumulatorCallback(DiagnosticCallback):
-    """
+    """
     Callback that evaluates a (scalar) functional involving integrals in both
     time and space.
 
@@ -619,7 +636,7 @@ 
Source code for thetis.callback
     variable_names = ['spatial integral at current timestep']
 
     def __init__(self, scalar_callback, solver_obj, **kwargs):
-        """
+        """
         :arg scalar_callback: Python function that returns a list of values of an objective functional.
         :arg solver_obj: Thetis solver object
         :arg kwargs: any additional keyword arguments, see DiagnosticCallback
@@ -649,7 +666,7 @@ 
Source code for thetis.callback
 
 
 
[docs]class TimeSeriesCallback2D(DiagnosticCallback):
-    """
+    """
     Extract a time series of a 2D field at a given (x,y) location
 
     Currently implemented only for the 3D model.
@@ -664,7 +681,7 @@ 
Source code for thetis.callback
                  append_to_log=True,
                  tolerance=1e-3, eval_func=None,
                  start_time=None, end_time=None):
-        """
+        """
         :arg solver_obj: Thetis :class:`FlowSolver` object
         :arg fieldnames: List of fields to extract
         :arg x, y: location coordinates in model coordinate system.
@@ -763,7 +780,7 @@ 
Source code for thetis.callback
 
 
 
[docs]class TimeSeriesCallback3D(DiagnosticCallback):
-    """
+    """
     Extract a time series of a 3D field at a given (x,y,z) location
 
     Currently implemented only for the 3D model.
@@ -776,7 +793,7 @@ 
Source code for thetis.callback
                  location_name,
                  outputdir=None, export_to_hdf5=True, append_to_log=True,
                  start_time=None, end_time=None):
-        """
+        """
         :arg solver_obj: Thetis :class:`FlowSolver` object
         :arg fieldnames: List of fields to extract
         :arg x, y, z: location coordinates in model coordinate system.
@@ -862,7 +879,7 @@ 
Source code for thetis.callback
 
 
 
[docs]class VerticalProfileCallback(DiagnosticCallback):
-    """
+    """
     Extract a vertical profile of a 3D field at a given (x,y) location
 
     Only for the 3D model.
@@ -876,7 +893,7 @@ 
Source code for thetis.callback
                  npoints=48,
                  outputdir=None, export_to_hdf5=True,
                  append_to_log=True):
-        """
+        """
         :arg solver_obj: Thetis :class:`FlowSolver` object
         :arg fieldnames: List of fields to extract
         :arg x, y: location coordinates in model coordinate system.
@@ -970,7 +987,7 @@ 
Source code for thetis.callback
 
 
 
[docs]class TransectCallback(DiagnosticCallback):
-    """
+    """
     Extract a vertical transect of a 3D field at a given (x,y) locations.
 
     Only for the 3D model.
@@ -985,7 +1002,7 @@ 
Source code for thetis.callback
                  z_min=None, z_max=None,
                  outputdir=None, export_to_hdf5=True,
                  append_to_log=True):
-        """
+        """
         :arg solver_obj: Thetis :class:`FlowSolver` object
         :arg fieldnames: List of fields to extract
         :arg x, y: location coordinates in model coordinate system.
@@ -1128,8 +1145,8 @@ 
Source code for thetis.callback
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/configuration.html b/_modules/thetis/configuration.html
index ad4bbfe..b65b729 100644
--- a/_modules/thetis/configuration.html
+++ b/_modules/thetis/configuration.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.configuration — Thetis 0+untagged.1922.g02d08b3.dirty documentation
+    thetis.configuration — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -55,9 +55,9 @@ 
Source code for thetis.configuration
 Utility function and extensions to traitlets used for specifying Thetis options
 """
 import textwrap
+import traitlets
 from textwrap import dedent
 from traitlets.config.configurable import Configurable
-from traitlets import *
 from firedrake import Constant, Function
 import datetime
 import ufl
@@ -70,7 +70,7 @@ 
Source code for thetis.configuration
 
 
 
[docs]def rst_all_options(cls, nspace=0, prefix=None):
-    """Recursively generate rST for a provided Configurable class.
+    """Recursively generate rST for a provided Configurable class.
 
     :arg cls: The Configurable class.
     :arg nspace: Indentation level.
@@ -125,7 +125,7 @@ 
Source code for thetis.configuration
     return "\n".join(lines)

 
 
-
[docs]class PositiveInteger(Integer):
+
[docs]class PositiveInteger(traitlets.Integer):
 
[docs]    def info(self):
         return u'a positive integer'

 
@@ -135,7 +135,7 @@ 
Source code for thetis.configuration
         return proposal

 
 
-
[docs]class PositiveFloat(Float):
+
[docs]class PositiveFloat(traitlets.Float):
 
[docs]    def info(self):
         return u'a positive float'

 
@@ -145,7 +145,7 @@ 
Source code for thetis.configuration
         return proposal

 
 
-
[docs]class NonNegativeInteger(Integer):
+
[docs]class NonNegativeInteger(traitlets.Integer):
 
[docs]    def info(self):
         return u'a non-negative integer'

 
@@ -155,7 +155,7 @@ 
Source code for thetis.configuration
         return proposal

 
 
-
[docs]class NonNegativeFloat(Float):
+
[docs]class NonNegativeFloat(traitlets.Float):
 
[docs]    def info(self):
         return u'a non-negative float'

 
@@ -165,8 +165,8 @@ 
Source code for thetis.configuration
         return proposal

 
 
-
[docs]class BoundedInteger(Integer):
-    def __init__(self, default_value=Undefined, bounds=None, **kwargs):
+
[docs]class BoundedInteger(traitlets.Integer):
+    def __init__(self, default_value=traitlets.Undefined, bounds=None, **kwargs):
         super(BoundedInteger, self).__init__(default_value, **kwargs)
         self.minval = bounds[0]
         self.maxval = bounds[1]
@@ -181,8 +181,8 @@ 
Source code for thetis.configuration
         return proposal

 
 
-
[docs]class BoundedFloat(Float):
-    def __init__(self, default_value=Undefined, bounds=None, **kwargs):
+
[docs]class BoundedFloat(traitlets.Float):
+    def __init__(self, default_value=traitlets.Undefined, bounds=None, **kwargs):
         self.minval = bounds[0]
         self.maxval = bounds[1]
         super(BoundedFloat, self).__init__(default_value, **kwargs)
@@ -197,7 +197,7 @@ 
Source code for thetis.configuration
         return proposal

 
 
-
[docs]class FiredrakeConstantTraitlet(TraitType):
+
[docs]class FiredrakeConstantTraitlet(traitlets.TraitType):
     default_value = None
     info_text = 'a Firedrake Constant'
 
@@ -210,7 +210,7 @@ 
Source code for thetis.configuration
         return 'Constant({:})'.format(float(self.default_value))

 
 
-
[docs]class FiredrakeCoefficient(TraitType):
+
[docs]class FiredrakeCoefficient(traitlets.TraitType):
     default_value = None
     info_text = 'a Firedrake Constant or Function'
 
@@ -225,7 +225,7 @@ 
Source code for thetis.configuration
         return 'Function'

 
 
-
[docs]class FiredrakeScalarExpression(TraitType):
+
[docs]class FiredrakeScalarExpression(traitlets.TraitType):
     default_value = None
     info_text = 'a scalar UFL expression'
 
@@ -243,7 +243,7 @@ 
Source code for thetis.configuration
         return 'UFL scalar expression'

 
 
-
[docs]class FiredrakeVectorExpression(TraitType):
+
[docs]class FiredrakeVectorExpression(traitlets.TraitType):
     default_value = None
     info_text = 'a vector UFL expression'
 
@@ -261,7 +261,7 @@ 
Source code for thetis.configuration
         return 'UFL vector expression'

 
 
-
[docs]class DatetimeTraitlet(TraitType):
+
[docs]class DatetimeTraitlet(traitlets.TraitType):
     default_value = None
     info_text = 'a timezone-aware datetime object'
 
@@ -271,8 +271,8 @@ 
Source code for thetis.configuration
         self.error(obj, value)

 
 
-
[docs]class PETScSolverParameters(Dict):
-    """PETSc solver options dictionary"""
+
[docs]class PETScSolverParameters(traitlets.Dict):
+    """PETSc solver options dictionary"""
     info_text = 'a PETSc solver options dictionary'
 
 
[docs]    def validate(self, obj, value):
@@ -281,8 +281,8 @@ 
Source code for thetis.configuration
         self.error(obj, value)

 
 
-
[docs]class PairedEnum(Enum):
-    """A enum whose value must be in a given sequence.
+
[docs]class PairedEnum(traitlets.Enum):
+    """A enum whose value must be in a given sequence.
 
     This enum controls a slaved option, with default values provided here.
 
@@ -292,7 +292,7 @@ 
Source code for thetis.configuration
     :arg paired_name: trait name this enum is paired with.
     :arg default_value: default value.
     """
-    def __init__(self, values, paired_name, default_value=Undefined, **kwargs):
+    def __init__(self, values, paired_name, default_value=traitlets.Undefined, **kwargs):
         self.paired_defaults = dict(values)
         self.paired_name = paired_name
         values, _ = zip(*values)
@@ -304,16 +304,16 @@ 
Source code for thetis.configuration
 
 
 
[docs]class OptionsBase:
-    """Abstract base class for all options classes"""
+    """Abstract base class for all options classes"""
     __metaclass__ = ABCMeta
 
     @abstractproperty
     def name(self):
-        """Human readable name of the configurable object"""
+        """Human readable name of the configurable object"""
         pass
 
 
[docs]    def update(self, options):
-        """
+        """
         Assign options from another container
 
         :arg options: Either a dictionary of options or another
@@ -322,13 +322,13 @@ 
Source code for thetis.configuration
         if isinstance(options, dict):
             params_dict = options
         else:
-            assert isinstance(options, HasTraits), 'options must be a dict or HasTraits object'
+            assert isinstance(options, traitlets.HasTraits), 'options must be a dict or HasTraits object'
             params_dict = options._trait_values
         for key in params_dict:
             self.__setattr__(key, params_dict[key])

 
     def __str__(self):
-        """Returs a summary of all defined parameters and their values in a string"""
+        """Returs a summary of all defined parameters and their values in a string"""
         output = '{:} parameters\n'.format(self.name)
         params_dict = self._trait_values
         for k in sorted(params_dict.keys()):
@@ -339,14 +339,14 @@ 
Source code for thetis.configuration
 # HasTraits and Configurable (and all their subclasses) have MetaHasTraits as their metaclass
 # to subclass from HasTraits and another class with ABCMeta as its metaclass, we need a combined
 # meta class that sub(meta)classes from ABCMeta and MetaHasTraits
-
[docs]class ABCMetaHasTraits(ABCMeta, MetaHasTraits):
-    """Combined metaclass of ABCMeta and MetaHasTraits"""
+
[docs]class ABCMetaHasTraits(ABCMeta, traitlets.MetaHasTraits):
+    """Combined metaclass of ABCMeta and MetaHasTraits"""
     pass

 
 
-
[docs]class FrozenHasTraits(OptionsBase, HasTraits):
+
[docs]class FrozenHasTraits(OptionsBase, traitlets.HasTraits):
     __metaclass__ = ABCMetaHasTraits
-    """
+    """
     A HasTraits class that only allows adding new attributes in the class
     definition or when  self._isfrozen is False.
     """
@@ -364,7 +364,7 @@ 
Source code for thetis.configuration
 
 
 
[docs]class FrozenConfigurable(OptionsBase, Configurable):
-    """
+    """
     A Configurable class that only allows adding new attributes in the class
     definition or when  self._isfrozen is False.
     """
@@ -384,7 +384,7 @@ 
Source code for thetis.configuration
 
 
 
[docs]def attach_paired_options(name, name_trait, value_trait):
-    """Attach paired options to a Configurable object.
+    """Attach paired options to a Configurable object.
 
     :arg name: the name of the enum trait
     :arg name_trait: the enum trait (a PairedEnum)
@@ -400,8 +400,8 @@ 
Source code for thetis.configuration
         if hasattr(name_trait, 'default_value') and name_trait.default_value is not None:
             return name_trait.paired_defaults[name_trait.default_value]()
 
-    obs_handler = observe(name, type="change")
-    def_handler = default(name_trait.paired_name)
+    obs_handler = traitlets.observe(name, type="change")
+    def_handler = traitlets.default(name_trait.paired_name)
 
     def update_class(cls):
         "Programmatically update the class"
@@ -415,6 +415,7 @@ 
Source code for thetis.configuration
         value_trait.class_init(cls, name_trait.paired_name)
         obs_handler.class_init(cls, "_%s_observer" % name)
         def_handler.class_init(cls, "_%s_default" % name_trait.paired_name)
+        cls.setup_class(cls.__dict__)
 
         return cls
     return update_class

@@ -428,7 +429,7 @@ 
Source code for thetis.configuration
       

     

     
   
diff --git a/_modules/thetis/coordsys.html b/_modules/thetis/coordsys.html
index f3e630c..57bc15d 100644
--- a/_modules/thetis/coordsys.html
+++ b/_modules/thetis/coordsys.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.coordsys — Thetis 0+untagged.1891.g4528d31 documentation
+    thetis.coordsys — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -58,46 +58,131 @@ 
Source code for thetis.coordsys
 import firedrake as fd
 import pyproj
 import numpy
+from abc import ABC, abstractmethod
 
-
-UTM_ZONE10 = pyproj.Proj(
-    proj='utm',
-    zone=10,
-    datum='WGS84',
-    units='m',
-    errcheck=True)
 LL_WGS84 = pyproj.Proj(proj='latlong', datum='WGS84', errcheck=True)
 
 
-
[docs]def convert_coords(source_sys, target_sys, x, y):
-    """
-    Converts coordinates from source_sys to target_sys
+
[docs]class CoordinateSystem(ABC):
+    """
+    Base class for horizontal coordinate systems
+
+    Provides methods for coordinate transformations etc.
+    """
+
[docs]    @abstractmethod
+    def to_lonlat(self, x, y):
+        """Convert coordinates to latitude and longitude"""
+        pass

+
+
[docs]    @abstractmethod
+    def get_vector_rotator(self, x, y):
+        """
+        Returns a vector rotator object.
+
+        The rotator converst vector-valued data to/from longitude, latitude
+        coordinates.
+        """
+        pass

 
-    This function extends pyproj.transform method by handling NaNs correctly.
 
-    :arg source_sys: pyproj coordinate system where (x, y) are defined in
-    :arg target_sys: target pyproj coordinate system
-    :arg x: x coordinate
-    :arg y: y coordinate
-    :type x: float or numpy.array_like
-    :type y: float or numpy.array_like
+
[docs]def proj_transform(x, y, trans):
+    """
+    Transform coordinates from source to target system.
+
+    :arg x,y: coordinates, float or numpy.array_like
+    :kwarg trans: pyproj Transformer object
     """
-    if isinstance(x, numpy.ndarray):
-        # proj may give wrong results if nans in the arrays
-        lon = numpy.full_like(x, numpy.nan)
-        lat = numpy.full_like(y, numpy.nan)
-        goodIx = numpy.logical_and(numpy.isfinite(x), numpy.isfinite(y))
-        lon[goodIx], lat[goodIx] = pyproj.transform(
-            source_sys, target_sys, x[goodIx], y[goodIx])
+    x_is_array = isinstance(x, numpy.ndarray)
+    y_is_array = isinstance(y, numpy.ndarray)
+    numpy_inputs = x_is_array or y_is_array
+    if numpy_inputs:
+        assert x_is_array and y_is_array, 'both x and y must be numpy arrays'
+        assert x.shape == y.shape, 'x and y must have same shape'
+        # transform only non-nan entries as proj behavior can be erratic
+        a = numpy.full_like(x, numpy.nan)
+        b = numpy.full_like(y, numpy.nan)
+        good_ix = numpy.logical_and(numpy.isfinite(x), numpy.isfinite(y))
+        a[good_ix], b[good_ix] = trans.transform(x[good_ix], y[good_ix])
     else:
-        lon, lat = pyproj.transform(source_sys, target_sys, x, y)
-    return lon, lat

+        a, b = trans.transform(x, y)
+    return a, b

+
+
+
[docs]class UTMCoordinateSystem(CoordinateSystem):
+    """
+    Represents Universal Transverse Mercator coordinate systems
+    """
+    def __init__(self, utm_zone):
+        self.proj_obj = pyproj.Proj(proj='utm', zone=utm_zone, datum='WGS84',
+                                    units='m', errcheck=True)
+        self.transformer_lonlat = pyproj.Transformer.from_crs(
+            self.proj_obj.srs, LL_WGS84.srs)
+        self.transformer_xy = pyproj.Transformer.from_crs(
+            LL_WGS84.srs, self.proj_obj.srs)
+
+
[docs]    def to_lonlat(self, x, y, positive_lon=False):
+        """
+        Convert (x, y) coordinates to (latitude, longitude)
+
+        :arg x: x coordinate
+        :arg y: y coordinate
+        :type x: float or numpy.array_like
+        :type y: float or numpy.array_like
+        :kwarg positive_lon: should positive longitude be enforced?
+        :return: longitude, latitude coordinates
+        """
+        lon, lat = proj_transform(x, y, trans=self.transformer_lonlat)
+        if positive_lon:
+            lon = numpy.mod(lon, 360.0)
+        return lon, lat

+
+
[docs]    def to_xy(self, lon, lat):
+        """
+        Convert (latitude, longitude) coordinates to (x, y)
+
+        :arg lon: longitude coordinate
+        :arg lat: latitude coordinate
+        :type longitude: float or numpy.array_like
+        :type latitude: float or numpy.array_like
+        :return: x, y coordinates
+        """
+        x, y = proj_transform(lon, lat, trans=self.transformer_xy)
+        return x, y

+
+
[docs]    def get_mesh_lonlat_function(self, mesh2d):
+        """
+        Construct a :class:`Function` holding the mesh coordinates in
+        longitude-latitude coordinates.
 
+        :arg mesh2d: the 2D mesh
+        """
+        dim = mesh2d.topological_dimension()
+        if dim != 2:
+            raise ValueError(f'Expected a mesh of dimension 2, not {dim}')
+        if mesh2d.geometric_dimension() != 2:
+            raise ValueError('Mesh must reside in 2-dimensional space')
+        x = mesh2d.coordinates.dat.data_ro[:, 0]
+        y = mesh2d.coordinates.dat.data_ro[:, 1]
+        lon, lat = self.transformer_lonlat.transform(x, y)
+        lonlat = fd.Function(mesh2d.coordinates.function_space())
+        lonlat.dat.data[:, 0] = lon
+        lonlat.dat.data[:, 1] = lat
+        return lonlat

+
+
[docs]    def get_vector_rotator(self, lon, lat):
+        """
+        Returns a vector rotator object.
+
+        The rotator converts vector-valued data from longitude, latitude
+        coordinates to mesh coordinate system.
+        """
+        return VectorCoordSysRotation(self.transformer_xy, lon, lat)

 
-
[docs]def get_vector_rotation_matrix(source_sys, target_sys, x, y, delta=None):
-    """
+
+
[docs]def get_vector_rotation_matrix(trans, x, y, delta=None):
+    """
     Estimate rotation matrix that converts vectors defined in source_sys to
-    target_sys.
+    target_sys. Conversion is carried out by the given `trans` object.
 
     Assume that we have a vector field defined in source_sys: vectors located at
     (x, y) define the x and y components. We can then rotate the vectors to
@@ -106,7 +191,8 @@ 
Source code for thetis.coordsys
 
     .. code-block:: python
 
-        R, theta = get_vector_rotation_matrix(source_sys, target_sys, x, lat)
+        trans = pyproj.Transformer.from_crs(source_sys, target_sys)
+        R, theta = get_vector_rotation_matrix(trans, x, y)
         v_xy = numpy.array([[v_x], [v_y]])
         v_new = numpy.matmul(R, v_xy)
         v_x2, v_y2 = v_new
@@ -114,9 +200,8 @@ 
Source code for thetis.coordsys
     """
     if delta is None:
         delta = 1e-6  # ~1 m in LL_WGS84
-    x1, y1 = pyproj.transform(source_sys, target_sys, x, y)
-
-    x2, y2 = pyproj.transform(source_sys, target_sys, x, y + delta)
+    x1, y1 = trans.transform(x, y)
+    x2, y2 = trans.transform(x, y + delta)
     dxdl = (x2 - x1) / delta
     dydl = (y2 - y1) / delta
     theta = numpy.arctan2(-dxdl, dydl)
@@ -129,24 +214,22 @@ 
Source code for thetis.coordsys
 
 
 
[docs]class VectorCoordSysRotation(object):
-    """
+    """
     Rotates vectors defined in source_sys coordinates to a different coordinate
     system.
 
     """
-    def __init__(self, source_sys, target_sys, x, y):
-        """
-        :arg source_sys: pyproj coordinate system where (x, y) are defined in
-        :arg target_sys: target pyproj coordinate system
-        :arg x: x coordinate
-        :arg y: y coordinate
+    def __init__(self, trans, x, y):
+        """
+        :arg trans: pyproj.Transformer object, maps from source to destination system
+        :arg x, y: coordinates in the source coordinate system
         """
-        R, theta = get_vector_rotation_matrix(source_sys, target_sys, x, y)
+        R, theta = get_vector_rotation_matrix(trans, x, y)
         self.rotation_sin = numpy.sin(theta)
         self.rotation_cos = numpy.cos(theta)
 
     def __call__(self, v_x, v_y, i_node=None):
-        """
+        """
         Rotate vectors defined by the `v_x` and `v_y` components.
 
         :arg v_x, v_y: vector x, y components
@@ -159,42 +242,6 @@ 
Source code for thetis.coordsys
         u = v_x * self.rotation_cos[f] - v_y * self.rotation_sin[f]
         v = v_x * self.rotation_sin[f] + v_y * self.rotation_cos[f]
         return u, v

-
-
-
[docs]def to_latlon(coord_system, x, y, positive_lon=False):
-    """
-    Convert model coordinates to latitude-longitude coordinates.
-
-    :arg coord_system: the local mesh coordinate system
-    :arg x: x coordinate
-    :arg y: y coordinate
-    :kwarg positive_lon: should positive longitude be enforced?
-    """
-    lon, lat = convert_coords(coord_system, LL_WGS84, x, y)
-    if positive_lon:
-        lon = numpy.mod(lon, 360.0)
-    return lat, lon

-
-
-
[docs]def get_lonlat_function(mesh2d, coord_system):
-    """
-    Construct a :class:`Function` holding the mesh coordinates in
-    longitude-latitude coordinates.
-
-    :arg mesh2d: the 2D mesh
-    :arg coord_system: the coordinate system of the mesh
-    """
-    dim = mesh2d.topological_dimension()
-    if dim != 2:
-        raise ValueError(f"Expected a mesh of dimension 2, not {dim}")
-    trans = pyproj.Transformer.from_crs(coord_system.srs, LL_WGS84.srs)
-    x = mesh2d.coordinates.dat.data_ro[:, 0]
-    y = mesh2d.coordinates.dat.data_ro[:, 1]
-    lon, lat = trans.transform(x, y)
-    lonlat = fd.Function(mesh2d.coordinates.function_space())
-    lonlat.dat.data[:, 0] = lon
-    lonlat.dat.data[:, 1] = lat
-    return lonlat

 

 
             

@@ -205,8 +252,8 @@ 
Source code for thetis.coordsys
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/coupled_timeintegrator.html b/_modules/thetis/coupled_timeintegrator.html
index c8f2378..3a5f6df 100644
--- a/_modules/thetis/coupled_timeintegrator.html
+++ b/_modules/thetis/coupled_timeintegrator.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.coupled_timeintegrator — Thetis 0+untagged.1878.g52658ff.dirty documentation
+    thetis.coupled_timeintegrator — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -58,18 +58,18 @@ 
Source code for thetis.coupled_timeintegrator
from . import timeintegrator
 from .log import *
 from . import rungekutta
-from abc import ABCMeta, abstractproperty
+from abc import ABC, abstractproperty
 import numpy
 
 
 
[docs]class CoupledTimeIntegratorBase(timeintegrator.TimeIntegratorBase):
-    """
+    """
     Base class for coupled 2D-3D time integrators
 
     Provides common functionality for updating diagnostic fields etc.
     """
     def __init__(self, solver):
-        """
+        """
         :arg solver: :class:`.FlowSolver` object
         """
         self.solver = solver
@@ -78,30 +78,30 @@ 
Source code for thetis.coupled_timeintegrator
self.timesteppers = AttrDict()
 
     def _update_3d_elevation(self):
-        """Projects elevation to 3D"""
+        """Projects elevation to 3D"""
         with timed_stage('aux_elev_3d'):
             self.solver.fields.elev_domain_2d.assign(self.solver.fields.elev_2d)
 
     def _update_vertical_velocity(self):
-        """Solve vertical velocity"""
+        """Solve vertical velocity"""
         with timed_stage('continuity_eq'):
             self.solver.w_solver.solve()
 
     def _update_moving_mesh(self):
-        """Updates 3D mesh to match elevation field"""
+        """Updates 3D mesh to match elevation field"""
         if self.options.use_ale_moving_mesh:
             with timed_stage('aux_mesh_ale'):
                 self.solver.mesh_updater.update_mesh_coordinates()
 
     def _update_2d_coupling(self):
-        """Does 2D-3D coupling for the velocity field"""
+        """Does 2D-3D coupling for the velocity field"""
         with timed_stage('aux_uv_coupling'):
             self._remove_depth_average_from_uv_3d()
             self._update_2d_coupling_term()
             self._copy_uv_2d_to_3d()
 
     def _remove_depth_average_from_uv_3d(self):
-        """Computes depth averaged velocity and removes it from the 3D velocity field"""
+        """Computes depth averaged velocity and removes it from the 3D velocity field"""
         with timed_stage('aux_uv_coupling'):
             # compute depth averaged 3D velocity
             self.solver.uv_averager.solve()  # uv -> uv_dav_3d
@@ -111,24 +111,24 @@ 
Source code for thetis.coupled_timeintegrator
self.fields.uv_3d -= self.fields.uv_dav_3d
 
     def _copy_uv_2d_to_3d(self):
-        """Copies uv_2d to uv_dav_3d"""
+        """Copies uv_2d to uv_dav_3d"""
         with timed_stage('aux_uv_coupling'):
             self.solver.copy_uv_to_uv_dav_3d.solve()
 
     def _update_2d_coupling_term(self):
-        """Update split_residual_2d field for 2D-3D coupling"""
+        """Update split_residual_2d field for 2D-3D coupling"""
         with timed_stage('aux_uv_coupling'):
             # scale dav uv 2D to be used as a forcing in 2D mom eq.
             self.fields.split_residual_2d.assign(self.fields.uv_dav_2d)
             self.fields.split_residual_2d /= self.timesteppers.mom_expl.dt_const
 
     def _update_baroclinicity(self):
-        """Computes baroclinic head"""
+        """Computes baroclinic head"""
         if self.options.use_baroclinic_formulation:
             compute_baroclinic_head(self.solver)
 
     def _update_turbulence(self, t):
-        """
+        """
         Updates turbulence related fields
 
         :arg t: simulation time
@@ -144,7 +144,7 @@ 
Source code for thetis.coupled_timeintegrator
self.solver.turbulence_model.postprocess()
 
     def _update_stabilization_params(self):
-        """
+        """
         Computes Smagorinsky viscosity etc fields
         """
         with timed_stage('aux_stability'):
@@ -157,7 +157,7 @@ 
Source code for thetis.coupled_timeintegrator
do_ale_update=False,
                                  do_stab_params=False,
                                  do_turbulence=False):
-        """Default routine for updating all dependent fields after a time step"""
+        """Default routine for updating all dependent fields after a time step"""
         self._update_3d_elevation()
         if do_ale_update:
             self._update_moving_mesh()
@@ -180,30 +180,28 @@ 
Source code for thetis.coupled_timeintegrator
self._update_stabilization_params()

 
 
-
[docs]class CoupledTimeIntegrator(CoupledTimeIntegratorBase):
-    """
+
[docs]class CoupledTimeIntegrator(CoupledTimeIntegratorBase, ABC):
+    """
     Base class of mode-split time integrators that use 2D, 3D and implicit 3D
     time integrators.
     """
-    __metaclass__ = ABCMeta
-
     @abstractproperty
     def integrator_2d(self):
-        """time integrator for 2D equations"""
+        """time integrator for 2D equations"""
         pass
 
     @abstractproperty
     def integrator_3d(self):
-        """time integrator for explicit 3D equations (momentum, tracers)"""
+        """time integrator for explicit 3D equations (momentum, tracers)"""
         pass
 
     @abstractproperty
     def integrator_vert_3d(self):
-        """time integrator for implicit 3D equations (vertical diffusion)"""
+        """time integrator for implicit 3D equations (vertical diffusion)"""
         pass
 
     def __init__(self, solver):
-        """
+        """
         :arg solver: :class:`.FlowSolver` object
         """
         super(CoupledTimeIntegrator, self).__init__(solver)
@@ -217,7 +215,7 @@ 
Source code for thetis.coupled_timeintegrator
self.n_stages = self.timesteppers.swe2d.n_stages
 
     def _get_vert_diffusivity_functions(self):
-        """
+        """
         Assign vertical viscosity/diffusivity to implicit/explicit parts
         """
         if self.options.use_implicit_vertical_diffusion:
@@ -234,7 +232,7 @@ 
Source code for thetis.coupled_timeintegrator
return impl_v_visc, expl_v_visc, impl_v_diff, expl_v_diff
 
     def _create_swe_integrator(self):
-        """
+        """
         Create time integrator for 2D system
         """
         solver = self.solver
@@ -254,7 +252,7 @@ 
Source code for thetis.coupled_timeintegrator
solver.bnd_functions['shallow_water'])
 
     def _create_mom_integrator(self):
-        """
+        """
         Create time integrator for 3D momentum equation
         """
         solver = self.solver
@@ -290,7 +288,7 @@ 
Source code for thetis.coupled_timeintegrator
self.options.timestepper_options.implicit_momentum_options, solver.bnd_functions['momentum'])
 
     def _create_salt_integrator(self):
-        """
+        """
         Create time integrator for salinity equation
         """
         solver = self.solver
@@ -319,7 +317,7 @@ 
Source code for thetis.coupled_timeintegrator
self.options.timestepper_options.implicit_tracer_options, solver.bnd_functions['salt'])
 
     def _create_temp_integrator(self):
-        """
+        """
         Create time integrator for temperature equation
         """
         solver = self.solver
@@ -348,7 +346,7 @@ 
Source code for thetis.coupled_timeintegrator
self.options.timestepper_options.implicit_tracer_options, solver.bnd_functions['temp'])
 
     def _create_turb_integrator(self):
-        """
+        """
         Create time integrators for turbulence equations
         """
         solver = self.solver
@@ -390,7 +388,7 @@ 
Source code for thetis.coupled_timeintegrator
self.options.timestepper_options.explicit_tracer_options, {})
 
     def _create_integrators(self):
-        """
+        """
         Creates all time integrators with the correct arguments
         """
         self._create_swe_integrator()
@@ -403,7 +401,7 @@ 
Source code for thetis.coupled_timeintegrator
self.cfl_coeff_2d = self.timesteppers.swe2d.cfl_coeff
 
 
[docs]    def set_dt(self, dt, dt_2d):
-        """
+        """
         Set time step for the coupled time integrator
 
         :arg float dt: Time step. This is the master (macro) time step used to
@@ -422,7 +420,7 @@ 
Source code for thetis.coupled_timeintegrator
self.timesteppers[stepper].set_dt(dt)

 
 
[docs]    def initialize(self):
-        """
+        """
         Assign initial conditions to all necessary fields
 
         Initial conditions are read from :attr:`fields` dictionary.
@@ -451,7 +449,7 @@ 
Source code for thetis.coupled_timeintegrator
[docs]class CoupledLeapFrogAM3(CoupledTimeIntegrator):
-    """
+    """
     Leap-Frog Adams-Moulton 3 time integrator for coupled 2D-3D problem
 
     This is an ALE time integrator.
@@ -474,7 +472,7 @@ 
Source code for thetis.coupled_timeintegrator
[docs]    @PETSc.Log.EventDecorator("thetis.CoupledLeapFrogAM3.advance")
     def advance(self, t, update_forcings=None, update_forcings3d=None):
-        """
+        """
         Advances the equations for one time step
 
         :arg float t: simulation time
@@ -616,7 +614,7 @@ 
Source code for thetis.coupled_timeintegrator
[docs]class CoupledTwoStageRK(CoupledTimeIntegrator):
-    """
+    """
     Coupled time integrator based on SSPRK(2,2) scheme
 
     This ALE time integration method uses SSPRK(2,2) scheme to advance the 3D
@@ -638,7 +636,7 @@ 
Source code for thetis.coupled_timeintegrator
[docs]    @PETSc.Log.EventDecorator("thetis.CoupledTwoStageRK.store_elevation")
     def store_elevation(self, istage):
-        """
+        """
         Store current elevation field for computing mesh velocity
 
         Must be called before updating the 2D mode.
@@ -652,7 +650,7 @@ 
Source code for thetis.coupled_timeintegrator
[docs]    @PETSc.Log.EventDecorator("thetis.CoupledTwoStageRK.compute_mesh_velocity")
     def compute_mesh_velocity(self, istage):
-        """
+        """
         Computes mesh velocity for stage i
 
         Must be called after updating the 2D mode.
@@ -674,7 +672,7 @@ 
Source code for thetis.coupled_timeintegrator
[docs]    @PETSc.Log.EventDecorator("thetis.CoupledTwoStageRK.advance")
     def advance(self, t, update_forcings=None, update_forcings3d=None):
-        """
+        """
         Advances the equations for one time step
 
         :arg float t: simulation time
@@ -773,8 +771,8 @@ 
Source code for thetis.coupled_timeintegrator

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/coupled_timeintegrator_2d.html b/_modules/thetis/coupled_timeintegrator_2d.html
index 9c4418b..65fd8ad 100644
--- a/_modules/thetis/coupled_timeintegrator_2d.html
+++ b/_modules/thetis/coupled_timeintegrator_2d.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.coupled_timeintegrator_2d — Thetis 0+untagged.1878.g52658ff.dirty documentation
+    thetis.coupled_timeintegrator_2d — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -57,30 +57,28 @@ 
Source code for thetis.coupled_timeintegrator_2d

 from .utility import *
 from . import timeintegrator
 from .log import *
-from abc import ABCMeta
+from abc import ABC
 
 
-
[docs]class CoupledTimeIntegrator2D(timeintegrator.TimeIntegratorBase):
-    """
+
[docs]class CoupledTimeIntegrator2D(timeintegrator.TimeIntegratorBase, ABC):
+    """
     Base class of time integrator for coupled shallow water and tracer/sediment equations and exner equation
     """
-    __metaclass__ = ABCMeta
-
 
[docs]    def swe_integrator(self):
-        """time integrator for the shallow water equations"""
+        """time integrator for the shallow water equations"""
         pass

 
 
[docs]    def tracer_integrator(self):
-        """time integrator for the tracer equation"""
+        """time integrator for the tracer equation"""
         pass

 
 
[docs]    def exner_integrator(self):
-        """time integrator for the exner equation"""
+        """time integrator for the exner equation"""
         pass

 
     @PETSc.Log.EventDecorator("thetis.CoupledTimeIntegrator2D.__init__")
     def __init__(self, solver):
-        """
+        """
         :arg solver: :class:`.FlowSolver` object
         """
         self.solver = solver
@@ -101,20 +99,20 @@ 
Source code for thetis.coupled_timeintegrator_2d

         self._create_integrators()
 
     def _create_integrators(self):
-        """
+        """
         Creates all time integrators with the correct arguments
         """
         if not self.options.tracer_only:
             self.timesteppers.swe2d = self.solver.get_swe_timestepper(self.swe_integrator)
-        for label in self.options.tracer:
-            self.timesteppers[label] = self.solver.get_tracer_timestepper(self.tracer_integrator, label)
+        for system in self.options.tracer_fields:
+            self.timesteppers[system] = self.solver.get_tracer_timestepper(self.tracer_integrator, system)
         if self.solver.options.sediment_model_options.solve_suspended_sediment:
             self.timesteppers.sediment = self.solver.get_sediment_timestepper(self.sediment_integrator)
         if self.solver.options.sediment_model_options.solve_exner:
             self.timesteppers.exner = self.solver.get_exner_timestepper(self.exner_integrator)
 
 
[docs]    def set_dt(self, dt):
-        """
+        """
         Set time step for the coupled time integrator
 
         :arg float dt: Time step.
@@ -124,7 +122,7 @@ 
Source code for thetis.coupled_timeintegrator_2d

 
 
[docs]    @PETSc.Log.EventDecorator("thetis.CoupledTimeIntegrator2D.initialize")
     def initialize(self, solution2d):
-        """
+        """
         Assign initial conditions to all necessary fields
 
         Initial conditions are read from :attr:`fields` dictionary.
@@ -136,8 +134,8 @@ 
Source code for thetis.coupled_timeintegrator_2d

 
         if not self.options.tracer_only:
             self.timesteppers.swe2d.initialize(self.fields.solution_2d)
-        for label in self.options.tracer:
-            self.timesteppers[label].initialize(self.fields[label])
+        for system in self.options.tracer_fields:
+            self.timesteppers[system].initialize(self.fields[system])
         if self.options.sediment_model_options.solve_suspended_sediment:
             self.timesteppers.sediment.initialize(self.fields.sediment_2d)
         if self.options.sediment_model_options.solve_exner:
@@ -152,10 +150,12 @@ 
Source code for thetis.coupled_timeintegrator_2d

         else:
             if not self.options.tracer_only:
                 self.timesteppers.swe2d.advance(t, update_forcings=update_forcings)
-            for label in self.options.tracer:
-                self.timesteppers[label].advance(t, update_forcings=update_forcings)
+            for system in self.options.tracer_fields:
+                self.timesteppers[system].advance(t, update_forcings=update_forcings)
                 if self.options.use_limiter_for_tracers:
-                    self.solver.tracer_limiter.apply(self.fields[label])
+                    if ',' in system:
+                        raise NotImplementedError("Slope limiters not supported for mixed systems of tracers")
+                    self.solver.tracer_limiter.apply(self.fields[system])
             if self.solver.sediment_model is not None:
                 self.solver.sediment_model.update()
             if self.options.sediment_model_options.solve_suspended_sediment:
@@ -172,10 +172,12 @@ 
Source code for thetis.coupled_timeintegrator_2d

         p = self.options.tracer_picard_iterations
         for i in range(p):
             kwargs = {'update_lagged': i == 0, 'update_fields': i == p-1}
-            for label in self.options.tracer:
-                self.timesteppers[label].advance_picard(t, update_forcings=update_forcings, **kwargs)
+            for system in self.options.tracer_fields:
+                self.timesteppers[system].advance_picard(t, update_forcings=update_forcings, **kwargs)
                 if self.options.use_limiter_for_tracers:
-                    self.solver.tracer_limiter.apply(self.fields[label])
+                    if ',' in system:
+                        raise NotImplementedError("Slope limiters not supported for mixed systems of tracers")
+                    self.solver.tracer_limiter.apply(self.fields[system])
         if self.solver.sediment_model is not None:
             self.solver.sediment_model.update()
         if self.options.sediment_model_options.solve_suspended_sediment:
@@ -187,13 +189,13 @@ 
Source code for thetis.coupled_timeintegrator_2d

 
 
 
[docs]class GeneralCoupledTimeIntegrator2D(CoupledTimeIntegrator2D):
-    """
+    """
     A :class:`CoupledTimeIntegrator2D` which supports
     a general set of time integrators for the different
     components.
     """
     def __init__(self, solver, integrators):
-        """
+        """
         :arg solver: the :class:`FlowSolver2d` object
         :arg integrators: dictionary of time integrators
             to be used for each equation
@@ -210,7 +212,7 @@ 
Source code for thetis.coupled_timeintegrator_2d

 
 
 
[docs]class NonHydrostaticTimeIntegrator2D(CoupledTimeIntegrator2D):
-    """
+    """
     2D non-hydrostatic time integrator based on Shallow Water time integrator
 
     This time integration method uses SWE time integrator to advance
@@ -240,7 +242,7 @@ 
Source code for thetis.coupled_timeintegrator_2d

 
 
[docs]    @PETSc.Log.EventDecorator("thetis.NonHydrostaticTimeIntegrator2D.initialize")
     def initialize(self, solution2d):
-        """
+        """
         Assign initial conditions to all necessary fields
 
         Initial conditions are read from :attr:`fields` dictionary.
@@ -254,7 +256,7 @@ 
Source code for thetis.coupled_timeintegrator_2d

 
 
[docs]    @PETSc.Log.EventDecorator("thetis.NonHydrostaticTimeIntegrator2D.advance")
     def advance(self, t, update_forcings=None):
-        """Advances equations for one time step."""
+        """Advances equations for one time step."""
         if self.serial_advancing:
             # --- advance in serial ---
             self.timesteppers.swe2d.advance(t, update_forcings=update_forcings)
@@ -296,8 +298,8 @@ 
Source code for thetis.coupled_timeintegrator_2d

       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/diagnostics.html b/_modules/thetis/diagnostics.html
index 50f943d..8edf064 100644
--- a/_modules/thetis/diagnostics.html
+++ b/_modules/thetis/diagnostics.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.diagnostics — Thetis 0+untagged.1888.gd2b834f documentation
+    thetis.diagnostics — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -56,13 +56,14 @@ 
Source code for thetis.diagnostics
 """
 from .utility import *
 from .configuration import *
+from abc import ABCMeta, abstractmethod
 
-
-__all__ = ["VorticityCalculator2D", "HessianRecoverer2D", "KineticEnergyCalculator"]
+__all__ = ["VorticityCalculator2D", "HessianRecoverer2D", "KineticEnergyCalculator",
+           "ShallowWaterDualWeightedResidual2D", "TracerDualWeightedResidual2D"]
 
 
 class DiagnosticCalculator(FrozenHasTraits):
-    """
+    """
     Base class that defines the API for all diagnostic calculators.
     """
     __metaclass__ = ABCMeta
@@ -76,7 +77,7 @@ 
Source code for thetis.diagnostics
 
 
 
[docs]class VorticityCalculator2D(DiagnosticCalculator):
-    r"""
+    r"""
     Linear solver for recovering fluid vorticity,
     interpreted as a scalar field:
 
@@ -94,7 +95,7 @@ 
Source code for thetis.diagnostics
     @unfrozen
     @PETSc.Log.EventDecorator("thetis.VorticityCalculator2D.__init__")
     def __init__(self, uv_2d, vorticity_2d, **kwargs):
-        """
+        """
         :arg uv_2d: vector expression for the horizontal velocity.
         :arg vorticity_2d: :class:`Function` to hold calculated vorticity.
         :kwargs: to be passed to the :class:`LinearVariationalSolver`.
@@ -130,7 +131,7 @@ 
Source code for thetis.diagnostics
 
 
 
[docs]class HessianRecoverer2D(DiagnosticCalculator):
-    r"""
+    r"""
     Linear solver for recovering Hessians.
 
     Hessians are recoved using double :math:`L^2`
@@ -147,7 +148,7 @@ 
Source code for thetis.diagnostics
     @unfrozen
     @PETSc.Log.EventDecorator("thetis.HessianRecoverer2D.__init__")
     def __init__(self, field_2d, hessian_2d, gradient_2d=None, **kwargs):
-        """
+        """
         :arg field_2d: scalar expression to recover the Hessian of.
         :arg hessian_2d: :class:`Function` to hold recovered Hessian.
         :kwarg gradient_2d: :class:`Function` to hold recovered gradient.
@@ -182,7 +183,7 @@ 
Source code for thetis.diagnostics
         g, H = TrialFunctions(W)
         phi, tau = TestFunctions(W)
         sol = Function(W)
-        self._gradient, self._hessian = sol.split()
+        self._gradient, self._hessian = sol.subfunctions
 
         # The formulation is chosen such that f does not need to have any
         # finite element derivatives
@@ -236,7 +237,7 @@ 
Source code for thetis.diagnostics
 
 
 
[docs]class KineticEnergyCalculator(DiagnosticCalculator):
-    r"""
+    r"""
     Class for calculating dynamic pressure (i.e. kinetic energy),
 
     .. math::
@@ -251,7 +252,7 @@ 
Source code for thetis.diagnostics
     @unfrozen
     @PETSc.Log.EventDecorator("thetis.KineticEnergyCalculator.__init__")
     def __init__(self, uv, ke, density=None, horizontal=False, project=False):
-        """
+        """
         :arg uv: scalar expression for the fluid velocity.
         :arg ke: :class:`Function` to hold calculated kinetic energy.
         :kwarg density: fluid density.
@@ -275,6 +276,115 @@ 
Source code for thetis.diagnostics
         else:
             assert hasattr(self, 'interpolator')
             self.interpolator.interpolate()

+
+
+class DualWeightedResidual2D(DiagnosticCalculator):
+    r"""
+    Class for computing contributions to dual weighted residual (DWR)
+    error indicators.
+
+    Suppose we have a weak formulation
+
+    .. math::
+        F(u_h; v) = 0,\quad\forall v\in V,
+
+    where :math:`F(u_h;\cdot)` is the weak residual of the forward PDE
+    and :math:`u_h` is its weak solution. The DWR is obtained by
+    replacing the test function :math:`v` with the (exact) adjoint
+    solution :math:`u^*`.
+
+    In practice, we do not have the exact adjoint solution, so it is
+    common practice to approximate it in some enriched finite element
+    space.
+    """
+    __metaclass__ = ABCMeta
+    error = None
+
+    @unfrozen
+    @PETSc.Log.EventDecorator("thetis.DualWeightedResidual.__init__")
+    def __init__(self, solver_obj, dual):
+        """
+        :arg solver_obj: :class:`FlowSolver2d` instance
+        :arg dual: a :class:`Function` that approximates the true adjoint solution,
+            which will replace the test function
+        """
+        mesh2d = solver_obj.mesh2d
+        if mesh2d.topological_dimension() != 2:
+            dim = mesh2d.topological_dimension()
+            raise ValueError(f"Expected a mesh of dimension 2, not {dim}")
+        if mesh2d != dual.ufl_domain():
+            raise ValueError(f"Mismatching meshes ({mesh2d} vs {func.ufl_domain()})")
+        self.F = replace(self.form, {TestFunction(self.space): dual})
+
+    @abstractmethod
+    def form(self):
+        pass
+
+    @abstractmethod
+    def space(self):
+        pass
+
+    @PETSc.Log.EventDecorator("thetis.DualWeightedResidual.solve")
+    def solve(self):
+        self.error = form2indicator(self.F)
+
+
+
[docs]class ShallowWaterDualWeightedResidual2D(DualWeightedResidual2D):
+    """
+    Class for computing dual weighted residual contributions
+    for the shallow water equations.
+    """
+
+    def __init__(self, solver_obj, dual):
+        """
+        :arg solver_obj: :class:`FlowSolver2d` instance
+        :arg dual: a :class:`Function` that approximates the true adjoint solution,
+            which will replace the test function
+        """
+        self.solver_obj = solver_obj
+        options = solver_obj.options
+        if options.swe_timestepper_type not in ("SteadyState", "CrankNicolson"):
+            typ = options.swe_timestepper_type
+            raise NotImplementedError(f"Error indication not yet supported for {typ}")
+        super().__init__(solver_obj, dual)
+
+    @property
+    def form(self):
+        ts = self.solver_obj.timestepper
+        return ts.F if not hasattr(ts, "timesteppers") else ts.timesteppers.swe2d.F
+
+    @property
+    def space(self):
+        return self.solver_obj.function_spaces.V_2d

+
+
+
[docs]class TracerDualWeightedResidual2D(DualWeightedResidual2D):
+    """
+    Class for computing dual weighted residual contributions
+    for 2D tracer transport problems.
+    """
+
+    def __init__(self, solver_obj, dual, label="tracer_2d"):
+        """
+        :arg solver_obj: :class:`FlowSolver2d` instance
+        :arg dual: a :class:`Function` that approximates the true adjoint solution,
+            which will replace the test function
+        """
+        self.solver_obj = solver_obj
+        self.label = label
+        options = solver_obj.options
+        if options.tracer_timestepper_type not in ("SteadyState", "CrankNicolson"):
+            typ = options.tracer_timestepper_type
+            raise NotImplementedError(f"Error indication not yet supported for {typ}")
+        super().__init__(solver_obj, dual)
+
+    @property
+    def form(self):
+        return self.solver_obj.timestepper.timesteppers[self.label].F
+
+    @property
+    def space(self):
+        return self.solver_obj.function_spaces.Q_2d

 

 
             

@@ -285,8 +395,8 @@ 
Source code for thetis.diagnostics
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/equation.html b/_modules/thetis/equation.html
index 7d37e54..e472d60 100644
--- a/_modules/thetis/equation.html
+++ b/_modules/thetis/equation.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.equation — Thetis 0+untagged.1888.gc4e776a documentation
+    thetis.equation — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -60,29 +60,30 @@ 
Source code for thetis.equation
 
 
 
[docs]class Term(object):
-    """
+    """
     Implements a single term of an equation.
 
     .. note::
         Sign convention: all terms are assumed to be on the right hand side of the equation: d(u)/dt = term.
     """
-    def __init__(self, function_space, test_function=None):
-        """
+    def __init__(self, function_space, test_function=None, trial_function=None):
+        """
         :arg function_space: the :class:`FunctionSpace` the solution belongs to
         :kwarg test_function: custom :class:`TestFunction`.
+        :kwarg trial_function: custom :class:`TrialFunction`.
         """
         # define bunch of members needed to construct forms
         self.function_space = function_space
         self.mesh = self.function_space.mesh()
         self.test = test_function or TestFunction(self.function_space)
-        self.tri = TrialFunction(self.function_space)
+        self.tri = trial_function or TrialFunction(self.function_space)
         self.normal = FacetNormal(self.mesh)
         # TODO construct them here from mesh ?
         self.boundary_markers = sorted(function_space.mesh().exterior_facets.unique_markers)
         self.boundary_len = function_space.mesh().boundary_len
 
 
[docs]    def residual(self, solution, solution_old, fields, fields_old, bnd_conditions):
-        """
+        """
         Returns an UFL form of the term.
 
         :arg solution: solution :class:`.Function` of the corresponding equation
@@ -96,7 +97,7 @@ 
Source code for thetis.equation
         raise NotImplementedError('Must be implemented in the derived class')

 
 
[docs]    def jacobian(self, solution, solution_old, fields, fields_old, bnd_conditions):
-        """
+        """
         Returns an UFL form of the Jacobian of the term.
 
         :arg solution: solution :class:`.Function` of the corresponding equation
@@ -112,11 +113,11 @@ 
Source code for thetis.equation
 
 
 
[docs]class Equation(object):
-    """
+    """
     Implements an equation, made out of terms.
     """
     SUPPORTED_LABELS = frozenset(['source', 'explicit', 'implicit', 'nonlinear'])
-    """
+    """
     Valid labels for terms, indicating how they should be treated in the time
     integrator.
 
@@ -134,7 +135,7 @@ 
Source code for thetis.equation
     """
     @PETSc.Log.EventDecorator("thetis.Equation.__init__")
     def __init__(self, function_space):
-        """
+        """
         :arg function_space: the :class:`FunctionSpace` the solution belongs to
         """
         self.terms = OrderedDict()
@@ -149,27 +150,30 @@ 
Source code for thetis.equation
         self.e_x, self.e_y, self.e_y = unit_vectors(3)
 
 
[docs]    def mass_term(self, solution):
-        """
+        """
         Returns default mass matrix term for the solution function space.
 
         :returns: UFL form of the mass term
         """
         return inner(solution, self.test) * dx

 
-
[docs]    def add_term(self, term, label):
-        """
+
[docs]    def add_term(self, term, label, suffix=None):
+        """
         Adds a term in the equation
 
         :arg term: :class:`.Term` object to add_term
         :arg string label: Assign a label to the term. Valid labels are given by
             :attr:`.SUPPORTED_LABELS`.
+        :arg suffix: optional custom key suffix
         """
         key = term.__class__.__name__
+        if suffix is not None:
+            key = '_'.join([key, suffix])
         self.terms[key] = term
         self.label_term(key, label)

 
 
[docs]    def label_term(self, term, label):
-        """
+        """
         Assings a label to the given term(s).
 
         :arg term: :class:`.Term` object, or a tuple of terms
@@ -184,7 +188,7 @@ 
Source code for thetis.equation
                 self.label_term(k, label)

 
 
[docs]    def select_terms(self, label):
-        """
+        """
         Generator function that selects terms by label(s).
 
         label can be a single label (e.g. 'explicit'), 'all' or a tuple of
@@ -203,7 +207,7 @@ 
Source code for thetis.equation
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.Equation.residual")
     def residual(self, label, solution, solution_old, fields, fields_old, bnd_conditions):
-        """
+        """
         Returns an UFL form of the residual by summing up all the terms with the desired label.
 
         Sign convention: all terms are assumed to be on the right hand side of the equation: d(u)/dt = term.
@@ -226,7 +230,7 @@ 
Source code for thetis.equation
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.Equation.jacobian")
     def jacobian(self, label, solution, solution_old, fields, fields_old, bnd_conditions):
-        """
+        """
         Returns an UFL form of the Jacobian by summing up all the Jacobians of the terms.
 
         Sign convention: all terms are assumed to be on the right hand side of the equation: d(u)/dt = term.
@@ -257,8 +261,8 @@ 
Source code for thetis.equation
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/exner_eq.html b/_modules/thetis/exner_eq.html
index 65a24da..66ae1f6 100644
--- a/_modules/thetis/exner_eq.html
+++ b/_modules/thetis/exner_eq.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.exner_eq — Thetis 0+untagged.1878.g52658ff.dirty documentation
+    thetis.exner_eq — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -68,8 +68,8 @@ 
Source code for thetis.exner_eq
 :math:`\nabla_h` denotes horizontal gradient, :math:`m` is the morphological scale factor,
 :math:`p` is the porosity and :math:`\textbf{Q_b}` is the bedload transport vector
 """
-from .equation import Term, Equation
 from .utility import *
+from .equation import Term, Equation
 
 __all__ = [
     'ExnerEquation',
@@ -80,12 +80,12 @@ 
Source code for thetis.exner_eq
 
 
 
[docs]class ExnerTerm(Term):
-    """
+    """
     Generic term in the Exner equations that provides commonly used members
     There are no boundary conditions for the Exner equation.
     """
     def __init__(self, function_space, depth, sediment_model, depth_integrated_sediment=False):
-        """
+        """
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :arg depth: :class:`DepthExpression` containing depth info
         :arg sediment_model: :class:`SedimentModel` containing sediment info
@@ -107,7 +107,7 @@ 
Source code for thetis.exner_eq
 
 
 
[docs]class ExnerSourceTerm(ExnerTerm):
-    r"""
+    r"""
     Source term accounting for suspended sediment transport
 
     The weak form reads
@@ -138,7 +138,7 @@ 
Source code for thetis.exner_eq
 
 
 
[docs]class ExnerBedloadTerm(ExnerTerm):
-    r"""
+    r"""
     Bedload transport term, \nabla_h \cdot \textbf{Q_b}
 
     The weak form is
@@ -183,7 +183,7 @@ 
Source code for thetis.exner_eq
 
 
 class ExnerSedimentSlideTerm(ExnerTerm):
-    r"""
+    r"""
     Term which adds component to bedload transport to ensure the slope angle does not exceed a certain value
     """
     def residual(self, solution, solution_old, fields, fields_old, bnd_conditions):
@@ -203,13 +203,13 @@ 
Source code for thetis.exner_eq
 
 
 
[docs]class ExnerEquation(Equation):
-    """
+    """
     Exner equation
 
     2D conservation of mass equation describing bed evolution due to sediment transport
     """
     def __init__(self, function_space, depth, sediment_model, depth_integrated_sediment):
-        """
+        """
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :arg depth: :class: `DepthExpression` containing depth info
         :arg sediment_model: :class: `SedimentModel` containing sediment info
@@ -237,8 +237,8 @@ 
Source code for thetis.exner_eq
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/exporter.html b/_modules/thetis/exporter.html
index 15b5195..4d92082 100644
--- a/_modules/thetis/exporter.html
+++ b/_modules/thetis/exporter.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.exporter — Thetis 0+untagged.1878.g52658ff.dirty documentation
+    thetis.exporter — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -61,13 +61,13 @@ 
Source code for thetis.exporter
 
 
 
[docs]def is_2d(fs):
-    """Tests wether a function space is 2D or 3D"""
+    """Tests wether a function space is 2D or 3D"""
     return fs.mesh().geometric_dimension() == 2

 
 
 
[docs]@PETSc.Log.EventDecorator("thetis.get_visu_space")
 def get_visu_space(fs):
-    """
+    """
     Returns an appropriate VTK visualization space for a function space
 
     :arg fs: function space
@@ -75,11 +75,11 @@ 
Source code for thetis.exporter
     """
     mesh = fs.mesh()
     family = 'Lagrange' if is_cg(fs) else 'Discontinuous Lagrange'
-    if len(fs.ufl_element().value_shape()) == 1:
-        dim = fs.ufl_element().value_shape()[0]
+    if len(fs.ufl_element().value_shape) == 1:
+        dim = fs.ufl_element().value_shape[0]
         visu_fs = get_functionspace(mesh, family, 1, family, 1,
                                     vector=True, dim=dim)
-    elif len(fs.ufl_element().value_shape()) == 2:
+    elif len(fs.ufl_element().value_shape) == 2:
         visu_fs = get_functionspace(mesh, family, 1, family, 1,
                                     tensor=True)
     else:
@@ -90,11 +90,11 @@ 
Source code for thetis.exporter
 
 
 
[docs]class ExporterBase(object):
-    """
+    """
     Base class for exporter objects.
     """
     def __init__(self, filename, outputdir, next_export_ix=0, verbose=False):
-        """
+        """
         :arg string filename: output file name (without directory)
         :arg string outputdir: directory where file is stored
         :kwarg int next_export_ix: set the index for next output
@@ -107,22 +107,22 @@ 
Source code for thetis.exporter
         self.next_export_ix = next_export_ix
 
 
[docs]    def set_next_export_ix(self, next_export_ix):
-        """Sets the index of next export"""
+        """Sets the index of next export"""
         self.next_export_ix = next_export_ix

 
 
[docs]    def export(self, function):
-        """Exports given function to disk"""
+        """Exports given function to disk"""
         raise NotImplementedError('This method must be implemented in the derived class')

 
 
 
[docs]class VTKExporter(ExporterBase):
-    """
+    """
     Class that handles Paraview VTK file exports
     """
     @PETSc.Log.EventDecorator("thetis.VTKExporter.__init__")
     def __init__(self, fs_visu, func_name, outputdir, filename,
                  next_export_ix=0, project_output=False, verbose=False):
-        """
+        """
         :arg fs_visu: function space where input function will be cast
             before exporting
         :arg func_name: name of the function
@@ -148,7 +148,7 @@ 
Source code for thetis.exporter
         self.cast_operators = {}
 
 
[docs]    def set_next_export_ix(self, next_export_ix):
-        """Sets the index of next export"""
+        """Sets the index of next export"""
         # NOTE vtk io objects store current export index not next
         super(VTKExporter, self).set_next_export_ix(next_export_ix)
         # FIXME hack to change correct output file count
@@ -156,7 +156,7 @@ 
Source code for thetis.exporter
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.VTKExporter.export")
     def export(self, function):
-        """Exports given function to disk"""
+        """Exports given function to disk"""
         assert self.fs_visu.max_work_functions == 1
         tmp_proj_func = self.fs_visu.get_work_function()
         # NOTE tmp function must be invariant as the projector is built only once
@@ -182,13 +182,13 @@ 
Source code for thetis.exporter
 
 
 
[docs]class HDF5Exporter(ExporterBase):
-    """
+    """
     Stores fields in disk in native discretization using HDF5 containers
     """
     @PETSc.Log.EventDecorator("thetis.HDF5Exporter.__init__")
     def __init__(self, function_space, outputdir, filename_prefix,
-                 next_export_ix=0, verbose=False):
-        """
+                 next_export_ix=0, legacy_mode=False, verbose=False):
+        """
         Create exporter object for given function.
 
         :arg function_space: space where the exported functions belong
@@ -197,40 +197,50 @@ 
Source code for thetis.exporter
         :arg string filename_prefix: prefix of output filename. Filename is
             prefix_nnnnn.h5 where nnnnn is the export number.
         :kwarg int next_export_ix: index for next export (default 0)
+        :kwarg bool legacy_mode: use legacy DumbCheckpoint format
         :kwarg bool verbose: print debug info to stdout
         """
         super(HDF5Exporter, self).__init__(filename_prefix, outputdir,
                                            next_export_ix, verbose)
         self.function_space = function_space
+        self.dumb_checkpoint = legacy_mode
 
 
[docs]    def gen_filename(self, iexport):
-        """
+        """
         Generate file name 'prefix_nnnnn.h5' for i-th export
 
         :arg int iexport: export index >= 0
         """
         filename = '{0:s}_{1:05d}'.format(self.filename, iexport)
+        if not self.dumb_checkpoint:
+            filename += '.h5'
         return os.path.join(self.outputdir, filename)

 
 
[docs]    def export_as_index(self, iexport, function):
-        """
+        """
         Export function to disk using the specified export index number
 
         :arg int iexport: export index >= 0
         :arg function: :class:`Function` to export
         """
-        assert function.function_space() == self.function_space,\
+        assert function.function_space() == self.function_space, \
             'Function space does not match'
         filename = self.gen_filename(iexport)
         if self.verbose:
             print_output('saving {:} state to {:}'.format(function.name(), filename))
-        with DumbCheckpoint(filename, mode=FILE_CREATE, comm=function.comm) as f:
-            f.store(function)
+        if self.dumb_checkpoint:
+            with DumbCheckpoint(filename, mode=FILE_CREATE, comm=function.comm) as f:
+                f.store(function)
+        else:
+            with CheckpointFile(filename, 'w') as f:
+                mesh = function.function_space().mesh()
+                f.save_mesh(mesh)
+                f.save_function(function)
         self.next_export_ix = iexport + 1

 
 
[docs]    @PETSc.Log.EventDecorator("thetis.HDF5Exporter.export")
     def export(self, function):
-        """
+        """
         Export function to disk.
 
         Increments export index by 1.
@@ -241,23 +251,31 @@ 
Source code for thetis.exporter
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.HDF5Exporter.load")
     def load(self, iexport, function):
-        """
+        """
         Loads nodal values from disk and assigns to the given function
 
         :arg int iexport: export index >= 0
         :arg function: target :class:`Function`
         """
-        assert function.function_space() == self.function_space,\
+        assert function.function_space() == self.function_space, \
             'Function space does not match'
         filename = self.gen_filename(iexport)
         if self.verbose:
             print_output('loading {:} state from {:}'.format(function.name(), filename))
-        with DumbCheckpoint(filename, mode=FILE_READ, comm=function.comm) as f:
-            f.load(function)

+        if self.dumb_checkpoint:
+            with DumbCheckpoint(filename, mode=FILE_READ, comm=function.comm) as f:
+                f.load(function)
+        else:
+            with CheckpointFile(filename, 'r') as f:
+                mesh = function.function_space().mesh()
+                if not hasattr(mesh, 'sfXC'):
+                    raise IOError('When loading fields from hdf5 checkpoint files, you should also read the mesh from checkpoint. See the documentation for `read_mesh_from_checkpoint()`')
+                g = f.load_function(mesh, function.name())
+                function.assign(g)

 
 
 
[docs]class ExportManager(object):
-    """
+    """
     Helper object for exporting multiple fields simultaneously
 
     .. code-block:: python
@@ -274,8 +292,9 @@ 
Source code for thetis.exporter
     """
     def __init__(self, outputdir, fields_to_export, functions, field_metadata,
                  export_type='vtk', next_export_ix=0, verbose=False,
+                 legacy_mode=False,
                  preproc_funcs={}):
-        """
+        """
         :arg string outputdir: directory where files are stored
         :arg fields_to_export: list of fields to export
         :type fields_to_export: list of strings
@@ -285,6 +304,7 @@ 
Source code for thetis.exporter
         :kwarg str export_type: export format, either 'vtk' or 'hdf5'
         :kwarg int next_export_ix: index for next export (default 0)
         :kwarg bool verbose: print debug info to stdout
+        :kwarg bool legacy_mode: use legacy `DumbCheckpoint` hdf5 format
         """
         self.outputdir = outputdir
         self.fields_to_export = fields_to_export
@@ -300,12 +320,14 @@ 
Source code for thetis.exporter
             field = self.functions.get(key)
             if field is not None and isinstance(field, Function):
                 self.add_export(key, field, export_type,
+                                legacy_mode=legacy_mode,
                                 next_export_ix=next_export_ix)
 
 
[docs]    def add_export(self, fieldname, function,
                    export_type='vtk', next_export_ix=0, outputdir=None,
-                   shortname=None, filename=None, preproc_func=None):
-        """
+                   shortname=None, filename=None, legacy_mode=False,
+                   preproc_func=None):
+        """
         Adds a new field exporter in the manager.
 
         This method allows exporting both default Thetis fields and user
@@ -319,6 +341,7 @@ 
Source code for thetis.exporter
         :kwarg string outputdir: optional directory where files are stored
         :kwarg string shortname: override shortname defined in field_metadata
         :kwarg string filename: override filename defined in field_metadata
+        :kwarg bool legacy_mode: use legacy `DumbCheckpoint` hdf5 format
         :kwarg preproc_func: optional funtion that will be called prior to
             exporting. E.g. for computing diagnostic fields.
         """
@@ -345,15 +368,16 @@ 
Source code for thetis.exporter
             elif export_type.lower() == 'hdf5':
                 self.exporters[fieldname] = HDF5Exporter(native_space,
                                                          outputdir, filename,
+                                                         legacy_mode=legacy_mode,
                                                          next_export_ix=next_export_ix)

 
 
[docs]    def set_next_export_ix(self, next_export_ix):
-        """Set export index to all child exporters"""
+        """Set export index to all child exporters"""
         for k in self.exporters:
             self.exporters[k].set_next_export_ix(next_export_ix)

 
 
[docs]    def export(self):
-        """
+        """
         Export all designated functions to disk
 
         Increments export index by 1.
@@ -375,7 +399,7 @@ 
Source code for thetis.exporter
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.ExportManager.export_bathymetry")
     def export_bathymetry(self, bathymetry_2d):
-        """
+        """
         Special function to export 2D bathymetry data to disk
 
         :arg bathymetry_2d: 2D bathymetry :class:`Function`
@@ -393,8 +417,8 @@ 
Source code for thetis.exporter
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/implicitexplicit.html b/_modules/thetis/implicitexplicit.html
index 1842efc..73c24d3 100644
--- a/_modules/thetis/implicitexplicit.html
+++ b/_modules/thetis/implicitexplicit.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.implicitexplicit — Thetis 0+untagged.1888.gc4e776a documentation
+    thetis.implicitexplicit — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -58,8 +58,8 @@ 
Source code for thetis.implicitexplicit
 from .rungekutta import *
 
 
-
[docs]class IMEXGeneric(TimeIntegrator):
-    """
+
[docs]class IMEXGeneric(TimeIntegrator, ABC):
+    """
     Generic implementation of Runge-Kutta Implicit-Explicit schemes
 
     Derived classes must define the implicit :attr:`dirk_class` and explicit
@@ -70,21 +70,19 @@ 
Source code for thetis.implicitexplicit
     solver. In case of non-linear terms proper linearization must defined in the
     equation using the two solution functions (solution, solution_old)
     """
-    __metaclass__ = ABCMeta
-
     @abstractproperty
     def dirk_class(self):
-        """Implicit DIRK class"""
+        """Implicit DIRK class"""
         pass
 
     @abstractproperty
     def erk_class(self):
-        """Explicit Runge-Kutta class"""
+        """Explicit Runge-Kutta class"""
         pass
 
     @PETSc.Log.EventDecorator("thetis.IMEXGeneric.__init__")
     def __init__(self, equation, solution, fields, dt, options, bnd_conditions):
-        """
+        """
         :arg equation: equation to solve
         :type equation: :class:`Equation` object
         :arg solution: :class:`Function` where solution will be stored
@@ -115,26 +113,26 @@ 
Source code for thetis.implicitexplicit
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.IMEXGeneric.update_solver")
     def update_solver(self):
-        """Create solver objects"""
+        """Create solver objects"""
         self.dirk.update_solver()
         self.erk.update_solver()

 
 
[docs]    @PETSc.Log.EventDecorator("thetis.IMEXGeneric.initialize")
     def initialize(self, solution):
-        """Assigns initial conditions to all required fields."""
+        """Assigns initial conditions to all required fields."""
         self.dirk.initialize(solution)
         self.erk.initialize(solution)

 
 
[docs]    @PETSc.Log.EventDecorator("thetis.IMEXGeneric.advance")
     def advance(self, t, update_forcings=None):
-        """Advances equations for one time step."""
+        """Advances equations for one time step."""
         for i in range(self.n_stages):
             self.solve_stage(i, t, update_forcings)
         self.get_final_solution()

 
 
[docs]    @PETSc.Log.EventDecorator("thetis.IMEXGeneric.solve_stage")
     def solve_stage(self, i_stage, t, update_forcings=None):
-        """
+        """
         Solves i-th stage
         """
         # set solution to u_n + dt*sum(a*k_erk)
@@ -150,7 +148,7 @@ 
Source code for thetis.implicitexplicit
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.IMEXGeneric.get_final_solution")
     def get_final_solution(self):
-        """
+        """
         Evaluates the final solution.
         """
         # set solution to u_n + sum(b*k_erk)
@@ -163,7 +161,7 @@ 
Source code for thetis.implicitexplicit
         self.erk.solution_old.assign(self.dirk.solution)

 
 
[docs]    def set_dt(self, dt):
-        """
+        """
         Update time step
 
         :arg float dt: time step
@@ -173,7 +171,7 @@ 
Source code for thetis.implicitexplicit
 
 
 
[docs]class IMEXLPUM2(IMEXGeneric):
-    """
+    """
     SSP-IMEX RK scheme (20) in Higureras et al. (2014)
 
     CFL coefficient is 2.0
@@ -187,7 +185,7 @@ 
Source code for thetis.implicitexplicit
 
 
 
[docs]class IMEXLSPUM2(IMEXGeneric):
-    """
+    """
     SSP-IMEX RK scheme (17) in Higureras et al. (2014)
 
     CFL coefficient is 2.0
@@ -201,7 +199,7 @@ 
Source code for thetis.implicitexplicit
 
 
 
[docs]class IMEXMidpoint(IMEXGeneric):
-    """
+    """
     Implicit-explicit midpoint scheme (1, 2, 2) from Ascher et al. (1997)
 
     Ascher et al. (1997). Implicit-explicit Runge-Kutta methods for
@@ -213,7 +211,7 @@ 
Source code for thetis.implicitexplicit
 
 
 
[docs]class IMEXEuler(IMEXGeneric):
-    """
+    """
     Forward-Backward Euler
     """
     erk_class = ERKEuler
@@ -228,8 +226,8 @@ 
Source code for thetis.implicitexplicit
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/interpolation.html b/_modules/thetis/interpolation.html
index 28db785..e04036c 100644
--- a/_modules/thetis/interpolation.html
+++ b/_modules/thetis/interpolation.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.interpolation — Thetis 0+untagged.1891.g4528d31 documentation
+    thetis.interpolation — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -93,12 +93,11 @@ 
Source code for thetis.interpolation
 """
 import glob
 import os
-from .coordsys import to_latlon
 from .timezone import *
 from .log import *
 import scipy.spatial.qhull as qhull
 import netCDF4
-from abc import ABCMeta, abstractmethod
+from abc import ABC, abstractmethod
 from firedrake import *
 from firedrake.petsc import PETSc
 import re
@@ -110,7 +109,7 @@ 
Source code for thetis.interpolation
 
 
 
[docs]def get_ncvar_name(ncfile, standard_name=None, long_name=None, var_name=None):
-    """
+    """
     Look for variables that match either CF standard_name or long_name
     attributes.
 
@@ -165,7 +164,7 @@ 
Source code for thetis.interpolation
 
 
 
[docs]class GridInterpolator(object):
-    """
+    """
     A reuseable griddata interpolator object.
 
     Usage:
@@ -195,7 +194,7 @@ 
Source code for thetis.interpolation
     @PETSc.Log.EventDecorator("thetis.GridInterpolator.__init__")
     def __init__(self, grid_xyz, target_xyz, fill_mode=None, fill_value=numpy.nan,
                  normalize=False, dont_raise=False):
-        """
+        """
         :arg grid_xyz: Array of source grid coordinates, shape (npoints, 2) or
             (npoints, 3)
         :arg target_xyz: Array of target grid coordinates, shape (n, 2) or
@@ -274,7 +273,7 @@ 
Source code for thetis.interpolation
 
     @PETSc.Log.EventDecorator("thetis.GridInterpolator.__call__")
     def __call__(self, values):
-        """
+        """
         Interpolate values defined on grid_xyz to target_xyz.
 
         :arg values: Array of source values to interpolate, shape (npoints, )
@@ -295,12 +294,12 @@ 
Source code for thetis.interpolation
 
 
 
[docs]class FileTreeReader(object):
-    """
+    """
     Abstract base class of file tree reader object
     """
     @abstractmethod
     def __call__(self, filename, time_index):
-        """
+        """
         Reads a data for one time step from the file
 
         :arg str filename: a filename where to find the data (e.g. filename)
@@ -311,7 +310,7 @@ 
Source code for thetis.interpolation
 
 
 
[docs]class NetCDFTimeSeriesReader(FileTreeReader):
-    """
+    """
     A simple netCDF reader that returns a time slice of the given variable.
 
     This class does not interpolate the data in any way. Useful for
@@ -331,7 +330,7 @@ 
Source code for thetis.interpolation
         self.ndims = len(nc_var.dimensions)
 
     def _get_slice(self, time_index):
-        """
+        """
         Returns a slice object that extracts a single time index
         """
         if self.ndims == 1:
@@ -341,7 +340,7 @@ 
Source code for thetis.interpolation
         return slice_list
 
     def __call__(self, filename, time_index):
-        """
+        """
         Reads a time_index from the data base
 
         :arg str filename: netcdf file where to find the data
@@ -360,7 +359,7 @@ 
Source code for thetis.interpolation
 
 
 def _get_subset_nodes(grid_x, grid_y, target_x, target_y):
-    """
+    """
     Retuns grid nodes that are necessary for intepolating onto target_x,y
     """
     orig_shape = grid_x.shape
@@ -379,46 +378,42 @@ 
Source code for thetis.interpolation
     return nodes, ind_x, ind_y
 
 
-
[docs]class SpatialInterpolator():
-    """
+
[docs]class SpatialInterpolator(ABC):
+    """
     Abstract base class for spatial interpolators that read data from disk
     """
-    __metaclass__ = ABCMeta
-
     @abstractmethod
     def __init__(self, function_space, coord_system):
-        """
+        """
         :arg function_space: target Firedrake FunctionSpace
-        :arg coord_system: the local mesh coordinate system
+        :arg coord_system: :class:`CoordinateSystem` object
         """
         pass
 
 
[docs]    @abstractmethod
     def interpolate(self, filename, variable_list, itime):
-        """
+        """
         Interpolates data from the given file at given time step
         """
         pass

 
 
-
[docs]class SpatialInterpolator2d(SpatialInterpolator):
-    """
+
[docs]class SpatialInterpolator2d(SpatialInterpolator, ABC):
+    """
     Abstract spatial interpolator class that can interpolate onto a 2D Function
     """
-    __metaclass__ = ABCMeta
-
     @PETSc.Log.EventDecorator("thetis.SpatialInterpolator2d.__init__")
     def __init__(self, function_space, coord_system, fill_mode=None,
                  fill_value=numpy.nan):
-        """
+        """
         :arg function_space: target Firedrake FunctionSpace
-        :arg coord_system: the local mesh coordinate system
+        :arg coord_system: :class:`CoordinateSystem` object
         :kwarg fill_mode: Determines how points outside the source grid will be
             treated. If 'nearest', value of the nearest source point will be
             used. Otherwise a constant fill value will be used (default).
         :kwarg float fill_value: Set the fill value (default: NaN)
         """
-        assert function_space.ufl_element().value_shape() == ()
+        assert function_space.ufl_element().value_shape == ()
 
         # construct local coordinates
         on_sphere = function_space.mesh().geometric_dimension() == 3
@@ -435,7 +430,7 @@ 
Source code for thetis.interpolation
             fsy = Function(function_space).interpolate(y).dat.data_with_halos
             coords = (fsx, fsy)
 
-        lat, lon = to_latlon(coord_system, *coords)
+        lon, lat = coord_system.to_lonlat(*coords)
         self.mesh_lonlat = numpy.array([lon, lat]).T
 
         self.fill_mode = fill_mode
@@ -444,7 +439,7 @@ 
Source code for thetis.interpolation
 
     @PETSc.Log.EventDecorator("thetis.SpatialInterpolator2d._create_interpolator")
     def _create_interpolator(self, lat_array, lon_array):
-        """
+        """
         Create compact interpolator by finding the minimal necessary support
         """
         assert len(lat_array.shape) == 2, 'Latitude must be two dimensional array.'
@@ -473,14 +468,14 @@ 
Source code for thetis.interpolation
 
 
[docs]    @abstractmethod
     def interpolate(self, filename, variable_list, time):
-        """
+        """
         Calls the interpolator object
         """
         pass

 
 
 
[docs]class NetCDFLatLonInterpolator2d(SpatialInterpolator2d):
-    """
+    """
     Interpolates netCDF data on a local 2D unstructured mesh
 
     The intepolator is constructed for a single netCDF file that defines the
@@ -502,7 +497,7 @@ 
Source code for thetis.interpolation
     """
 
[docs]    @PETSc.Log.EventDecorator("thetis.NetCDFLatLonInterpolator2d.interpolate")
     def interpolate(self, nc_filename, variable_list, itime):
-        """
+        """
         Interpolates data from a netCDF file onto Firedrake function space.
 
         :arg str nc_filename: netCDF file to read
@@ -536,7 +531,7 @@ 
Source code for thetis.interpolation
 
 
 
[docs]class NetCDFSpatialInterpolator(FileTreeReader):
-    """
+    """
     Wrapper class that provides FileTreeReader API for grid interpolators
     """
     def __init__(self, grid_interpolator, variable_list):
@@ -548,7 +543,7 @@ 
Source code for thetis.interpolation
 
 
 
[docs]class TimeParser(object):
-    """
+    """
     Abstract base class for time definition objects.
 
     Defines the time span that a file (or data set) covers and provides a time
@@ -556,17 +551,17 @@ 
Source code for thetis.interpolation
     """
 
[docs]    @abstractmethod
     def get_start_time(self):
-        """Returns the first time stamp in the file/data set"""
+        """Returns the first time stamp in the file/data set"""
         pass

 
 
[docs]    @abstractmethod
     def get_end_time(self):
-        """Returns the last time stamp in the file/data set"""
+        """Returns the last time stamp in the file/data set"""
         pass

 
 
[docs]    @abstractmethod
     def find_time_stamp(self, t, previous=False):
-        """
+        """
         Given time t, returns index of the next (previous) time stamp
 
         raises IndexError if t is out of range, i.e.
@@ -576,12 +571,12 @@ 
Source code for thetis.interpolation
 
 
 
[docs]class NetCDFTimeParser(TimeParser):
-    """
+    """
     Describes the time stamps stored in a netCDF file.
     """
     def __init__(self, filename, time_variable_name='time', allow_gaps=False,
                  verbose=False):
-        """
+        """
         Construct a new object by scraping data from the given netcdf file.
 
         :arg str filename: name of the netCDF file to read
@@ -594,7 +589,7 @@ 
Source code for thetis.interpolation
         self.time_variable_name = time_variable_name
 
         def get_datetime(time, units, calendar):
-            """
+            """
             Convert netcdf time value to datetime.
             """
             d = cftime.num2pydate(time, units, calendar)
@@ -644,12 +639,12 @@ 
Source code for thetis.interpolation
 
 
 
[docs]class TimeSearch(object):
-    """
+    """
     Base class for searching nearest time steps in a file tree or database
     """
 
[docs]    @abstractmethod
     def find(self, time, previous=False):
-        """
+        """
         Find a next (previous) time stamp from a given time
 
         :arg float time: input time stamp
@@ -661,7 +656,7 @@ 
Source code for thetis.interpolation
 
 
 
[docs]class NetCDFTimeSearch(TimeSearch):
-    """
+    """
     Finds a nearest time stamp in a collection of netCDF files.
     """
     @PETSc.Log.EventDecorator("thetis.NetCDFTimeSearch.__init__")
@@ -701,7 +696,7 @@ 
Source code for thetis.interpolation
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.NetCDFTimeSearch.find")
     def find(self, simulation_time, previous=False):
-        """
+        """
         Find file that contains the given simulation time
 
         :arg float simulation_time: simulation time in seconds
@@ -732,7 +727,7 @@ 
Source code for thetis.interpolation
 
 
 
[docs]class DailyFileTimeSearch(TimeSearch):
-    """
+    """
     Treats a list of daily files as a time series.
 
     File name pattern must be given as a string where the 4-digit year is
@@ -775,7 +770,7 @@ 
Source code for thetis.interpolation
                 print_output('  {:}'.format(self.start_datetime[i]))
 
     def _find_files(self):
-        """Finds all files that match the given pattern."""
+        """Finds all files that match the given pattern."""
         search_pattern = str(self.file_pattern)
         search_pattern = search_pattern.replace(':02d}', ':}')
         search_pattern = search_pattern.replace(':04d}', ':}')
@@ -785,7 +780,7 @@ 
Source code for thetis.interpolation
         return all_files
 
     def _parse_date(self, filename):
-        """
+        """
         Parse year, month, day from filename using the given pattern.
         """
         re_pattern = str(self.file_pattern)
@@ -804,7 +799,7 @@ 
Source code for thetis.interpolation
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.DailyFileTimeSearch.find")
     def find(self, simulation_time, previous=False):
-        """
+        """
         Find file that contains the given simulation time
 
         :arg float simulation_time: simulation time in seconds
@@ -823,7 +818,7 @@ 
Source code for thetis.interpolation
 
 
 
[docs]class LinearTimeInterpolator(object):
-    """
+    """
     Interpolates time series in time
 
     User must provide timesearch_obj that finds time stamps from
@@ -833,7 +828,7 @@ 
Source code for thetis.interpolation
     time.
     """
     def __init__(self, timesearch_obj, reader):
-        """
+        """
         :arg timesearch_obj: TimeSearch object
         :arg reader: FileTreeReader object
         """
@@ -842,7 +837,7 @@ 
Source code for thetis.interpolation
         self.cache = {}
 
     def _get_from_cache(self, key):
-        """
+        """
         Fetch data set from cache, read if not present
         """
         if key not in self.cache:
@@ -850,7 +845,7 @@ 
Source code for thetis.interpolation
         return self.cache[key]
 
     def _clean_cache(self, keys_to_keep):
-        """
+        """
         Remove cached data sets that are no longer needed
         """
         for key in list(self.cache.keys()):
@@ -858,7 +853,7 @@ 
Source code for thetis.interpolation
                 self.cache.pop(key)
 
     def __call__(self, t):
-        """
+        """
         Interpolate at time t
 
         :retuns: list of numpy arrays
@@ -883,13 +878,13 @@ 
Source code for thetis.interpolation
 
 
 
[docs]class NetCDFTimeSeriesInterpolator(object):
-    """
+    """
     Reads and interpolates scalar time series from a sequence of netCDF files.
     """
     @PETSc.Log.EventDecorator("thetis.NetCDFTimeSeriesInterpolator.__init__")
     def __init__(self, ncfile_pattern, variable_list, init_date,
                  time_variable_name='time', scalars=None, allow_gaps=False):
-        """
+        """
         :arg str ncfile_pattern: file search pattern, e.g. "mydir/foo_*.nc"
         :arg variable_list: list if netCDF variable names to read
         :arg datetime.datetime init_date: simulation start time
@@ -913,7 +908,7 @@ 
Source code for thetis.interpolation
 
     @PETSc.Log.EventDecorator("thetis.NetCDFTimeSeriesInterpolator.__call__")
     def __call__(self, time):
-        """
+        """
         Time series at the given time
 
         :returns: list of scalars or numpy.arrays
@@ -933,8 +928,8 @@ 
Source code for thetis.interpolation
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/inversion_tools.html b/_modules/thetis/inversion_tools.html
index 63eede9..c9ce83f 100644
--- a/_modules/thetis/inversion_tools.html
+++ b/_modules/thetis/inversion_tools.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.inversion_tools — Thetis 0+untagged.1888.gc4e776a documentation
+    thetis.inversion_tools — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -52,12 +52,15 @@
             
   
Source code for thetis.inversion_tools
 import firedrake as fd
-from firedrake_adjoint import *
+from firedrake.adjoint import *
+import ufl
+from .configuration import FrozenHasTraits
 from .solver2d import FlowSolver2d
-from .utility import create_directory, print_function_value_range
-from .utility import get_functionspace
+from .utility import create_directory, print_function_value_range, get_functionspace, unfrozen
 from .log import print_output
 from .diagnostics import HessianRecoverer2D
+from .exporter import HDF5Exporter
+import abc
 import numpy
 import h5py
 from scipy.interpolate import interp1d
@@ -65,37 +68,67 @@ 
Source code for thetis.inversion_tools
 import os
 
 
-
[docs]class OptimisationProgress(object):
-    """
-    Class for stashing progress of the optimisation routine.
+
[docs]class InversionManager(FrozenHasTraits):
+    """
+    Class for handling inversion problems and stashing
+    the progress of the associated optimization routines.
     """
-    J = 0  # cost function value (float)
-    dJdm_list = None  # cost function gradient (Function)
-    m_list = None  # control (Function)
-    m_progress = []
-    J_progress = []
-    dJdm_progress = []
-    i = 0
-    tic = None
-    nb_grad_evals = 0
-    control_coeff_list = []
-    control_list = []
-
-    def __init__(self, output_dir='outputs', no_exports=False, real=False):
-        """
+
+    @unfrozen
+    def __init__(self, sta_manager, output_dir='outputs', no_exports=False, real=False,
+                 penalty_parameters=[], cost_function_scaling=None,
+                 test_consistency=True, test_gradient=True):
+        """
+        :arg sta_manager: the :class:`StationManager` instance
         :kwarg output_dir: model output directory
-        :kwarg no_exports: toggle exports to vtu
+        :kwarg no_exports: if True, nothing will be written to disk
         :kwarg real: is the inversion in the Real space?
+        :kwarg penalty_parameters: a list of penalty parameters to pass
+            to the :class:`ControlRegularizationManager`
+        :kwarg cost_function_scaling: global scaling for the cost function.
+            As rule of thumb, it's good to scale the functional to J < 1.
+        :kwarg test_consistency: toggle testing the correctness with
+            which the :class:`ReducedFunctional` can recompute values
+        :kwarg test_gradient: toggle testing the correctness with
+            which the :class:`ReducedFunctional` can recompute gradients
         """
+        assert isinstance(sta_manager, StationObservationManager)
+        self.sta_manager = sta_manager
+        self.reg_manager = None
         self.output_dir = output_dir
         self.no_exports = no_exports or real
         self.real = real
+        self.penalty_parameters = penalty_parameters
+        self.cost_function_scaling = cost_function_scaling or fd.Constant(1.0)
+        self.sta_manager.cost_function_scaling = self.cost_function_scaling
+        self.test_consistency = test_consistency
+        self.test_gradient = test_gradient
         self.outfiles_m = []
         self.outfiles_dJdm = []
+        self.control_exporters = []
         self.initialized = False
 
-
[docs]    def initialize(self):
+        self.J = 0  # cost function value (float)
+        self.J_reg = 0  # regularization term value (float)
+        self.J_misfit = 0  # misfit term value (float)
+        self.dJdm_list = None  # cost function gradient (Function)
+        self.m_list = None  # control (Function)
+        self.Jhat = None
+        self.m_progress = []
+        self.J_progress = []
+        self.J_reg_progress = []
+        self.J_misfit_progress = []
+        self.dJdm_progress = []
+        self.i = 0
+        self.tic = None
+        self.nb_grad_evals = 0
+        self.control_coeff_list = []
+        self.control_list = []
+
+
[docs]    def initialize(self):
         if not self.no_exports:
+            if self.real:
+                raise ValueError("Exports are not supported in Real mode.")
             create_directory(self.output_dir)
             create_directory(self.output_dir + '/hdf5')
             for i in range(len(self.control_coeff_list)):
@@ -105,23 +138,29 @@ 
Source code for thetis.inversion_tools
                     fd.File(f'{self.output_dir}/gradient_progress_{i:02d}.pvd'))
         self.initialized = True

 
-
[docs]    def add_control(self, f):
-        """
+
[docs]    def add_control(self, f):
+        """
         Add a control field.
 
-        Can be called multiple times in case of multiparameter optimisation.
+        Can be called multiple times in case of multiparameter optimization.
 
         :arg f: Function or Constant to be used as a control variable.
         """
         self.control_coeff_list.append(f)
-        self.control_list.append(Control(f))

-
-
[docs]    def reset_counters(self):
+        self.control_list.append(Control(f))
+        if isinstance(f, fd.Function) and not self.no_exports:
+            j = len(self.control_coeff_list) - 1
+            prefix = f'control_{j:02d}'
+            self.control_exporters.append(
+                HDF5Exporter(f.function_space(), self.output_dir + '/hdf5', prefix)
+            )

+
+
[docs]    def reset_counters(self):
         self.nb_grad_evals = 0

 
-
[docs]    def set_control_state(self, j, djdm_list, m_list):
-        """
-        Stores optimisation state.
+
[docs]    def set_control_state(self, j, djdm_list, m_list):
+        """
+        Stores optimization state.
 
         To call whenever variables are updated.
 
@@ -131,22 +170,28 @@ 
Source code for thetis.inversion_tools
         """
         self.J = j
         self.dJdm_list = djdm_list
-        self.m_list = m_list

+        self.m_list = m_list
 
-
[docs]    def start_clock(self):
+        tape = get_working_tape()
+        reg_blocks = tape.get_blocks(tag="reg_eval")
+        self.J_reg = sum([b.get_outputs()[0].saved_output for b in reg_blocks])
+        misfit_blocks = tape.get_blocks(tag="misfit_eval")
+        self.J_misfit = sum([b.get_outputs()[0].saved_output for b in misfit_blocks])

+
+
[docs]    def start_clock(self):
         self.tic = time_mod.perf_counter()

 
-
[docs]    def stop_clock(self):
+
[docs]    def stop_clock(self):
         toc = time_mod.perf_counter()
         return toc

 
-
[docs]    def set_initial_state(self, *state):
+
[docs]    def set_initial_state(self, *state):
         self.set_control_state(*state)
         self.update_progress()

 
-
[docs]    def update_progress(self):
-        """
-        Updates optimisation progress and stores variables to disk.
+
[docs]    def update_progress(self):
+        """
+        Updates optimization progress and stores variables to disk.
 
         To call after successful line searches.
         """
@@ -155,7 +200,7 @@ 
Source code for thetis.inversion_tools
             for f in self.control_coeff_list:
                 print_function_value_range(f, prefix='Initial')
 
-        elapsed = '-' if self.tic is None else f'{toc - self.tic:.1f} s'
+        elapsed = '-' if self.tic is None else f'{toc - self.tic:.1f} s'
         self.tic = toc
 
         if not self.initialized:
@@ -167,12 +212,16 @@ 
Source code for thetis.inversion_tools
             controls = [m.dat.data[0] for m in self.m_list]
             self.m_progress.append(controls)
         self.J_progress.append(self.J)
+        self.J_reg_progress.append(self.J_reg)
+        self.J_misfit_progress.append(self.J_misfit)
         self.dJdm_progress.append(djdm)
         comm = self.control_coeff_list[0].comm
-        if comm.rank == 0:
+        if comm.rank == 0 and not self.no_exports:
             if self.real:
                 numpy.save(f'{self.output_dir}/m_progress', self.m_progress)
             numpy.save(f'{self.output_dir}/J_progress', self.J_progress)
+            numpy.save(f'{self.output_dir}/J_reg_progress', self.J_reg_progress)
+            numpy.save(f'{self.output_dir}/J_misfit_progress', self.J_misfit_progress)
             numpy.save(f'{self.output_dir}/dJdm_progress', self.dJdm_progress)
         if len(djdm) > 10:
             djdm = f"[{numpy.min(djdm):.4e} .. {numpy.max(djdm):.4e}]"
@@ -191,9 +240,8 @@ 
Source code for thetis.inversion_tools
                 m.rename(self.control_coeff_list[j].name())
                 o.write(m)
                 # hdf5 format
-                h5_filename = f'{self.output_dir}/hdf5/control_{j:02d}_{self.i:04d}'
-                with fd.DumbCheckpoint(h5_filename, mode=fd.FILE_CREATE) as chk:
-                    chk.store(m)
+                e = self.control_exporters[j]
+                e.export(m)
             # gradient output
             for f, o in zip(self.dJdm_list, self.outfiles_dJdm):
                 # store gradient in vtk format
@@ -205,22 +253,151 @@ 
Source code for thetis.inversion_tools
 
     @property
     def rf_kwargs(self):
-        """
+        """
         Default keyword arguments to pass to the
         :class:`ReducedFunctional` class.
         """
         def gradient_eval_cb(j, djdm, m):
             self.set_control_state(j, djdm, m)
             self.nb_grad_evals += 1
+            return djdm
 
         params = {
             'derivative_cb_post': gradient_eval_cb,
         }
-        return params

+        return params
+
+
[docs]    def get_cost_function(self, solver_obj, weight_by_variance=False):
+        r"""
+        Get a sum of square errors cost function for the problem:
+
+      ..math::
+            J(u) = \sum_{i=1}^{n_{ts}} \sum_{j=1}^{n_{sta}} (u_j^{(i)} - u_{j,o}^{(i)})^2,
+
+        where :math:`u_{j,o}^{(i)}` and :math:`u_j^{(i)}` denote the
+        observed and computed values at timestep :math:`i`, and
+        :math:`n_{ts}` and :math:`n_{sta}` are the numbers of timesteps
+        and stations, respectively.
+
+        Regularization terms are included if a
+        :class:`RegularizationManager` instance is provided.
+
+        :arg solver_obj: the :class:`FlowSolver2d` instance
+        :kwarg weight_by_variance: should the observation data be
+            weighted by the variance at each station?
+        """
+        assert isinstance(solver_obj, FlowSolver2d)
+        if len(self.penalty_parameters) > 0:
+            self.reg_manager = ControlRegularizationManager(
+                self.control_coeff_list,
+                self.penalty_parameters,
+                self.cost_function_scaling,
+                RSpaceRegularizationCalculator if self.real else HessianRegularizationCalculator)
+        self.J_reg = 0
+        self.J_misfit = 0
+        if self.reg_manager is not None:
+            self.J_reg = self.reg_manager.eval_cost_function()
+        self.J = self.J_reg
+
+        if weight_by_variance:
+            var = fd.Function(self.sta_manager.fs_points_0d)
+            for i, j in enumerate(self.sta_manager.local_station_index):
+                var.dat.data[i] = numpy.var(self.sta_manager.observation_values[j])
+            self.sta_manager.station_weight_0d.interpolate(1/var)
+
+        def cost_fn(t):
+            misfit = self.sta_manager.eval_cost_function(t)
+            self.J_misfit += misfit
+            self.J += misfit
+
+        return cost_fn

+
+    @property
+    def reduced_functional(self):
+        """
+        Create a Pyadjoint :class:`ReducedFunctional` for the optimization.
+        """
+        if self.Jhat is None:
+            self.Jhat = ReducedFunctional(self.J, self.control_list, **self.rf_kwargs)
+        return self.Jhat
+
+
[docs]    def stop_annotating(self):
+        """
+        Stop recording operations for the adjoint solver.
+
+        This method should be called after the :meth:`iterate`
+        method of :class:`FlowSolver2d`.
+        """
+        assert self.reduced_functional is not None
+        if self.test_consistency:
+            self.consistency_test()
+        if self.test_gradient:
+            self.taylor_test()
+        pause_annotation()

+
+
[docs]    def get_optimization_callback(self):
+        """
+        Get a callback for stashing optimization progress
+        after successful line search.
+        """
+
+        def optimization_callback(m):
+            self.update_progress()
+            if not self.no_exports:
+                self.sta_manager.dump_time_series()
+
+        return optimization_callback

+
+
[docs]    def minimize(self, opt_method="BFGS", bounds=None, **opt_options):
+        """
+        Minimize the reduced functional using a given optimization routine.
+
+        :kwarg opt_method: the optimization routine
+        :kwarg bounds: a list of bounds to pass to the optimization routine
+        :kwarg opt_options: other optimization parameters to pass
+        """
+        print_output(f'Running {opt_method} optimization')
+        self.reset_counters()
+        self.start_clock()
+        J = float(self.reduced_functional(self.control_coeff_list))
+        self.set_initial_state(J, self.reduced_functional.derivative(), self.control_coeff_list)
+        if not self.no_exports:
+            self.sta_manager.dump_time_series()
+        return minimize(
+            self.reduced_functional, method=opt_method, bounds=bounds,
+            callback=self.get_optimization_callback(), options=opt_options)

+
+
[docs]    def consistency_test(self):
+        """
+        Test that :attr:`reduced_functional` can correctly recompute the
+        objective value, assuming that none of the controls have changed
+        since it was created.
+        """
+        print_output("Running consistency test")
+        J = self.reduced_functional(self.control_coeff_list)
+        if not numpy.isclose(J, self.J):
+            raise ValueError(f"Consistency test failed (expected {self.J}, got {J})")
+        print_output("Consistency test passed!")

+
+
[docs]    def taylor_test(self):
+        """
+        Run a Taylor test to check that the :attr:`reduced_functional` can
+        correctly compute consistent gradients.
+
+        Note that the Taylor test is applied on the current control values.
+        """
+        func_list = []
+        for f in self.control_coeff_list:
+            dc = f.copy(deepcopy=True)
+            func_list.append(dc)
+        minconv = taylor_test(self.reduced_functional, self.control_coeff_list, func_list)
+        if minconv < 1.9:
+            raise ValueError("Taylor test failed")  # NOTE: Pyadjoint already prints the testing
+        print_output("Taylor test passed!")

 
 
 
[docs]class StationObservationManager:
-    """
+    """
     Implements error functional based on observation time series.
 
     The functional is the squared sum of error between the model and
@@ -230,23 +407,20 @@ 
Source code for thetis.inversion_tools
     interpolates the observations time series to the model time, computes the
     error functional, and also stores the model's time series data to disk.
     """
-    def __init__(self, mesh, J_scalar=None, output_directory='outputs'):
-        """
+    def __init__(self, mesh, output_directory='outputs'):
+        """
         :arg mesh: the 2D mesh object.
-        :kwarg J_scalar: Optional factor to scale the error functional. As a
-            rule of thumb, it's good to scale the functional to J < 1.
         :kwarg output_directory: directory where model time series are stored.
         """
         self.mesh = mesh
         on_sphere = self.mesh.geometric_dimension() == 3
         if on_sphere:
             raise NotImplementedError('Sphere meshes are not supported yet.')
-        self.J_scalar = J_scalar if J_scalar else fd.Constant(1.0)
+        self.cost_function_scaling = fd.Constant(1.0)
         self.output_directory = output_directory
-        create_directory(self.output_directory)
         # keep observation time series in memory
         self.obs_func_list = []
-        # keep model time series in memory during optimisation progress
+        # keep model time series in memory during optimization progress
         self.station_value_progress = []
         # model time when cost function was evaluated
         self.simulation_time = []
@@ -255,7 +429,7 @@ 
Source code for thetis.inversion_tools
 
 
[docs]    def register_observation_data(self, station_names, variable, time,
                                   values, x, y, start_times=None, end_times=None):
-        """
+        """
         Add station time series data to the object.
 
         The `x`, and `y` coordinates must be such that
@@ -281,14 +455,14 @@ 
Source code for thetis.inversion_tools
         self._end_times = end_times or numpy.ones(num_stations)*numpy.inf

 
 
[docs]    def set_model_field(self, function):
-        """
+        """
         Set the model field that will be evaluated.
         """
         self.model_observation_field = function

 
 
[docs]    def load_observation_data(self, observation_data_dir, station_names, variable,
                               start_times=None, end_times=None):
-        """
+        """
         Load observation data from disk.
 
         Assumes that observation data were stored with
@@ -329,7 +503,7 @@ 
Source code for thetis.inversion_tools
         self.construct_evaluator()

 
 
[docs]    def update_stations_in_use(self, t):
-        """
+        """
         Indicate which stations are in use at the current time.
 
         An entry of unity indicates use, whereas zero indicates disuse.
@@ -344,7 +518,7 @@ 
Source code for thetis.inversion_tools
         self.indicator_0d.assign(in_use)

 
 
[docs]    def construct_evaluator(self):
-        """
+        """
         Builds evaluators needed to compute the error functional.
         """
         # Create 0D mesh for station evaluation
@@ -355,6 +529,8 @@ 
Source code for thetis.inversion_tools
         self.mod_values_0d = fd.Function(self.fs_points_0d, name='model values')
         self.indicator_0d = fd.Function(self.fs_points_0d, name='station use indicator')
         self.indicator_0d.assign(1.0)
+        self.cost_function_scaling_0d = fd.Constant(0.0, domain=mesh0d)
+        self.cost_function_scaling_0d.assign(self.cost_function_scaling)
         self.station_weight_0d = fd.Function(self.fs_points_0d, name='station-wise weighting')
         self.station_weight_0d.assign(1.0)
         interp_kw = {}
@@ -398,7 +574,7 @@ 
Source code for thetis.inversion_tools
         self.initialized = True

 
 
[docs]    def eval_observation_at_time(self, t):
-        """
+        """
         Evaluate observation time series at the given time.
 
         :arg t: model simulation time
@@ -408,7 +584,7 @@ 
Source code for thetis.inversion_tools
         return [float(ip(t)) for ip in self.station_interpolators]

 
 
[docs]    def eval_cost_function(self, t):
-        """
+        """
         Evaluate the cost function.
 
         Should be called at every export of the forward model.
@@ -424,14 +600,15 @@ 
Source code for thetis.inversion_tools
         # compute square error
         self.obs_values_0d.assign(obs_func)
         self.mod_values_0d.interpolate(self.model_observation_field, ad_block_tag='observation')
-        J_misfit = fd.assemble(self.J_scalar*self.indicator_0d*self.station_weight_0d*self.misfit_expr**2*fd.dx)
-        return J_misfit

+        s = self.cost_function_scaling_0d * self.indicator_0d * self.station_weight_0d
+        self.J_misfit = fd.assemble(s * self.misfit_expr ** 2 * fd.dx, ad_block_tag='misfit_eval')
+        return self.J_misfit

 
 
[docs]    def dump_time_series(self):
-        """
+        """
         Stores model time series to disk.
 
-        Obtains station time series from the last optimisation iteration,
+        Obtains station time series from the last optimization iteration,
         and stores the data to disk.
 
         The output files are have the format
@@ -443,6 +620,7 @@ 
Source code for thetis.inversion_tools
         """
         assert self.station_names is not None
 
+        create_directory(self.output_directory)
         tape = get_working_tape()
         blocks = tape.get_blocks(tag='observation')
         ts_data = [b.get_outputs()[0].saved_output.dat.data for b in blocks]
@@ -468,123 +646,117 @@ 
Source code for thetis.inversion_tools
                 hdf5file.attrs.update(attrs)

 
 
-
[docs]class ControlRegularizationCalculator:
-    r"""
-    Computes regularization cost function for a control `Function`.
+
[docs]class RegularizationCalculator(abc.ABC):
+    """
+    Base class for computing regularization terms.
+
+    A derived class should set :attr:`regularization_expr` in
+    :meth:`__init__`. Whenever the cost function is evaluated,
+    the ratio of this expression and the total mesh area will
+    be added.
+    """
+    @abc.abstractmethod
+    def __init__(self, function, scaling=1.0):
+        """
+        :arg function: Control :class:`Function`
+        """
+        self.scaling = scaling
+        self.regularization_expr = 0
+        self.mesh = function.function_space().mesh()
+        # calculate mesh area (to scale the cost function)
+        self.mesh_area = fd.assemble(fd.Constant(1.0, domain=self.mesh) * fd.dx)
+        self.name = function.name()
+
+
[docs]    def eval_cost_function(self):
+        expr = self.scaling * self.regularization_expr / self.mesh_area * fd.dx
+        return fd.assemble(expr, ad_block_tag="reg_eval")

+
+
+
[docs]class HessianRegularizationCalculator(RegularizationCalculator):
+    r"""
+    Computes the following regularization term for a cost function
+    involving a control :class:`Function` :math:`f`:
 
     .. math::
-        J = \gamma | H(f) |^2
+        J = \gamma \| (\Delta x)^2 H(f) \|^2,
 
-    where :math:`H` is the Hessian of field math:`f`:.
+    where :math:`H` is the Hessian of field :math:`f`.
     """
-    def __init__(self, function, gamma_hessian):
-        """
-        :arg function: Control `Function`
-        :arg gamma_Hessian: Hessian penalty coefficient
+    def __init__(self, function, gamma, scaling=1.0):
+        """
+        :arg function: Control :class:`Function`
+        :arg gamma: Hessian penalty coefficient
         """
-        self.function = function
-        self.gamma_hessian = gamma_hessian
+        super().__init__(function, scaling=scaling)
         # solvers to evaluate the gradient of the control
-        mesh = function.function_space().mesh()
-        P1v_2d = get_functionspace(mesh, 'CG', 1, vector=True)
-        P1t_2d = get_functionspace(mesh, 'CG', 1, tensor=True)
-        name = function.name()
-        gradient_2d = fd.Function(P1v_2d, name=f'{name} gradient')
-        hessian_2d = fd.Function(P1t_2d, name=f'{name} hessian')
+        P1v_2d = get_functionspace(self.mesh, "CG", 1, vector=True)
+        P1t_2d = get_functionspace(self.mesh, "CG", 1, tensor=True)
+        gradient_2d = fd.Function(P1v_2d, name=f"{self.name} gradient")
+        hessian_2d = fd.Function(P1t_2d, name=f"{self.name} hessian")
         self.hessian_calculator = HessianRecoverer2D(
             function, hessian_2d, gradient_2d)
 
-        h = fd.CellSize(mesh)
+        h = fd.CellSize(self.mesh)
         # regularization expression |hessian|^2
         # NOTE this is normalized by the mesh element size
         # d^2 u/dx^2 * dx^2 ~ du^2
-        self.regularization_hess_expr = gamma_hessian * fd.inner(hessian_2d, hessian_2d)*h**4
-        # calculate mesh area (to scale the cost function)
-        self.mesh_area = fd.assemble(fd.Constant(1.0, domain=mesh)*fd.dx)
+        self.regularization_expr = gamma * fd.inner(hessian_2d, hessian_2d) * h**4
 
-
[docs]    def eval_cost_function(self):
+
[docs]    def eval_cost_function(self):
         self.hessian_calculator.solve()
-        J_regularization = fd.assemble(
-            self.regularization_hess_expr / self.mesh_area * fd.dx
-        )
-        return J_regularization

+        return super().eval_cost_function()

+
+
+
[docs]class RSpaceRegularizationCalculator(RegularizationCalculator):
+    r"""
+    Computes the following regularization term for a cost function
+    involving a control :class:`Function` :math:`f` from an R-space:
+
+    .. math::
+        J = \gamma (f - f_0)^2,
+
+    where :math:`f_0` is a prior, taken to be the initial value of
+    :math:`f`.
+    """
+    def __init__(self, function, gamma, eps=1.0e-03, scaling=1.0):
+        """
+        :arg function: Control :class:`Function`
+        :arg gamma: penalty coefficient
+        :kwarg eps: tolerance for normalising by near-zero priors
+        """
+        super().__init__(function, scaling=scaling)
+        R = function.function_space()
+        if R.ufl_element().family() != "Real":
+            raise ValueError("function must live in R-space")
+        prior = fd.Function(R, name=f"{self.name} prior")
+        prior.assign(function, annotate=False)  # Set the prior to the initial value
+        self.regularization_expr = gamma * (function - prior) ** 2 / ufl.max_value(abs(prior), eps)

+        # NOTE: If the prior is small then dividing by prior**2 puts too much emphasis
+        #       on the regularization. Therefore, we divide by abs(prior) instead.
 
 
 
[docs]class ControlRegularizationManager:
-    """
+    """
     Handles regularization of multiple control fields
     """
-    def __init__(self, function_list, gamma_list, J_scalar=None):
-        """
+    def __init__(self, function_list, gamma_list, penalty_term_scaling=None,
+                 calculator=HessianRegularizationCalculator):
+        """
         :arg function_list: list of control functions
         :arg gamma_list: list of penalty parameters
-        :kwarg J_scalar: Penalty term scaling factor
+        :kwarg penalty_term_scaling: Penalty term scaling factor
+        :kwarg calculator: class used for obtaining regularization
         """
-        self.J_scalar = J_scalar
         self.reg_calculators = []
         assert len(function_list) == len(gamma_list), \
             'Number of control functions and parameters must match'
-        for f, g in zip(function_list, gamma_list):
-            r = ControlRegularizationCalculator(f, g)
-            self.reg_calculators.append(r)
+        self.reg_calculators = [
+            calculator(f, g, scaling=penalty_term_scaling)
+            for f, g in zip(function_list, gamma_list)
+        ]
 
 
[docs]    def eval_cost_function(self):
-        v = 0
-        for r in self.reg_calculators:
-            u = r.eval_cost_function()
-            if self.J_scalar is not None:
-                u *= float(self.J_scalar)
-            v += u
-        return v

-
-
-
[docs]def get_cost_function(solver_obj, op, stationmanager,
-                      reg_manager=None, weight_by_variance=False):
-    r"""
-    Get a sum of square errors cost function for the problem:
-
-  ..math::
-        J(u) = \sum_{i=1}^{n_{ts}} \sum_{j=1}^{n_{sta}} (u_j^{(i)} - u_{j,o}^{(i)})^2,
-
-    where :math:`u_{j,o}^{(i)}` and :math:`u_j^{(i)}` denote the
-    observed and computed values at timestep :math:`i`, and
-    :math:`n_{ts}` and :math:`n_{sta}` are the numbers of timesteps
-    and stations, respectively.
-
-    Regularization terms are included if a
-    :class:`RegularizationManager` instance is provided.
-
-    Note that the current value of the cost function is
-    stashed on the :class:`OptimisationProgress` object.
-
-    :arg solver_obj: the :class:`FlowSolver2d` instance
-    :arg op: the :class:`OptimisationProgress` instance
-    :arg stationmanager: the :class:`StationManager` instance
-    :kwarg reg_manager: the :class:`RegularizationManager` instance
-    :kwarg weight_by_variance: should the observation data be
-        weighted by the variance at each station?
-    """
-    assert isinstance(solver_obj, FlowSolver2d)
-    assert isinstance(op, OptimisationProgress)
-    assert isinstance(stationmanager, StationObservationManager)
-    if reg_manager is None:
-        op.J = 0
-    else:
-        assert isinstance(reg_manager, ControlRegularizationManager)
-        op.J = reg_manager.eval_cost_function()
-
-    if weight_by_variance:
-        var = fd.Function(stationmanager.fs_points_0d)
-        for i, j in enumerate(stationmanager.local_station_index):
-            var.dat.data[i] = numpy.var(stationmanager.observation_values[j])
-        stationmanager.station_weight_0d.assign(1/var)
-
-    def cost_fn():
-        t = solver_obj.simulation_time
-        J_misfit = stationmanager.eval_cost_function(t)
-        op.J += J_misfit
-
-    return cost_fn

+        return sum([r.eval_cost_function() for r in self.reg_calculators])

 

 
             

@@ -595,8 +767,8 @@ 
Source code for thetis.inversion_tools
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/limiter.html b/_modules/thetis/limiter.html
index 9f29eea..2cc8579 100644
--- a/_modules/thetis/limiter.html
+++ b/_modules/thetis/limiter.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.limiter — Thetis 0+untagged.1888.gc4e776a documentation
+    thetis.limiter — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -56,13 +56,12 @@ 
Source code for thetis.limiter
 """
 from .utility import *
 from firedrake import VertexBasedLimiter
-import ufl
-from pyop2.profiling import timed_region, timed_function, timed_stage  # NOQA
+from pyop2.profiling import timed_stage
 import numpy
 
 
 
[docs]def assert_function_space(fs, family, degree):
-    """
+    """
     Checks the family and degree of function space.
 
     Raises AssertionError if function space differs.
@@ -77,30 +76,30 @@ 
Source code for thetis.limiter
     if not isinstance(family, list):
         fam_list = [family]
     ufl_elem = fs.ufl_element()
-    if isinstance(ufl_elem, ufl.VectorElement):
-        ufl_elem = ufl_elem.sub_elements()[0]
+    if isinstance(ufl_elem, firedrake.VectorElement):
+        ufl_elem = ufl_elem.sub_elements[0]
 
     if ufl_elem.family() == 'TensorProductElement':
         # extruded mesh
-        A, B = ufl_elem.sub_elements()
-        assert A.family() in fam_list,\
+        A, B = ufl_elem.sub_elements
+        assert A.family() in fam_list, \
             'horizontal space must be one of {0:s}'.format(fam_list)
-        assert B.family() in fam_list,\
+        assert B.family() in fam_list, \
             'vertical space must be {0:s}'.format(fam_list)
-        assert A.degree() == degree,\
+        assert A.degree() == degree, \
             'degree of horizontal space must be {0:d}'.format(degree)
-        assert B.degree() == degree,\
+        assert B.degree() == degree, \
             'degree of vertical space must be {0:d}'.format(degree)
     else:
         # assume 2D mesh
-        assert ufl_elem.family() in fam_list,\
+        assert ufl_elem.family() in fam_list, \
             'function space must be one of {0:s}'.format(fam_list)
-        assert ufl_elem.degree() == degree,\
+        assert ufl_elem.degree() == degree, \
             'degree of function space must be {0:d}'.format(degree)

 
 
 
[docs]class VertexBasedP1DGLimiter(VertexBasedLimiter):
-    """
+    """
     Vertex based limiter for P1DG tracer fields, see Kuzmin (2010)
 
     .. note::
@@ -112,7 +111,7 @@ 
Source code for thetis.limiter
     http://dx.doi.org/10.1016/j.cam.2009.05.028
     """
     def __init__(self, p1dg_space, time_dependent_mesh=True):
-        """
+        """
         :arg p1dg_space: P1DG function space
         """
 
@@ -128,7 +127,7 @@ 
Source code for thetis.limiter
         self.time_dependent_mesh = time_dependent_mesh
 
     def _construct_centroid_solver(self):
-        """
+        """
         Constructs a linear problem for computing the centroids
 
         :return: LinearSolver instance
@@ -142,7 +141,7 @@ 
Source code for thetis.limiter
                                                   'sub_pc_type': 'ilu'})
 
     def _update_centroids(self, field):
-        """
+        """
         Update centroid values
         """
         b = assemble(TestFunction(self.P0) * field * dx)
@@ -152,7 +151,7 @@ 
Source code for thetis.limiter
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.VertexBasedP1DGLimiter.compute_bounds")
     def compute_bounds(self, field):
-        """
+        """
         Re-compute min/max values of all neighbouring centroids
 
         :arg field: :class:`Function` to limit
@@ -234,7 +233,7 @@ 
Source code for thetis.limiter
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.VertexBasedP1DGLimiter.apply")
     def apply(self, field):
-        """
+        """
         Applies the limiter on the given field (in place)
 
         :arg field: :class:`Function` to limit
@@ -260,8 +259,8 @@ 
Source code for thetis.limiter
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/log.html b/_modules/thetis/log.html
index a5830a0..e2b904e 100644
--- a/_modules/thetis/log.html
+++ b/_modules/thetis/log.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.log — Thetis 0+untagged.1888.gc4e776a documentation
+    thetis.log — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -81,7 +81,7 @@ 
Source code for thetis.log
 
 
 class ThetisLogConfig:
-    """Module-wide config object"""
+    """Module-wide config object"""
     filename = None
     mem_buffer = None
 
@@ -94,7 +94,7 @@ 
Source code for thetis.log
 
 
 
[docs]def set_thetis_loggers(comm=COMM_WORLD):
-    """Set stream handlers for log messages.
+    """Set stream handlers for log messages.
 
     :kwarg comm: The communicator the handler should be collective
          over. Only rank-0 on that communicator will write to the log, other
@@ -125,7 +125,7 @@ 
Source code for thetis.log
 
 
 
[docs]def set_log_directory(output_directory, comm=COMM_WORLD, mode='w'):
-    """
+    """
     Forward all log output to `output_directory/log` file.
 
     When called, a new empty log file is created.
@@ -200,7 +200,7 @@ 
Source code for thetis.log
 
 
 
[docs]def thetis_log_level(level):
-    """Set the log level for Thetis logger.
+    """Set the log level for Thetis logger.
 
     This controls what level of logging messages are printed to
     stderr. The higher the level, the fewer the number of messages.
@@ -235,8 +235,8 @@ 
Source code for thetis.log
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/momentum_eq.html b/_modules/thetis/momentum_eq.html
index 0eb9e1c..2c03c1b 100644
--- a/_modules/thetis/momentum_eq.html
+++ b/_modules/thetis/momentum_eq.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.momentum_eq — Thetis 0+untagged.1878.g52658ff.dirty documentation
+    thetis.momentum_eq — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -133,7 +133,7 @@ 
Source code for thetis.momentum_eq
 
 
 
[docs]class MomentumTerm(Term):
-    """
+    """
     Generic term for momentum equation that provides commonly used members and
     mapping for boundary functions.
     """
@@ -142,7 +142,7 @@ 
Source code for thetis.momentum_eq
                  use_nonlinear_equations=True, use_lax_friedrichs=True,
                  use_bottom_friction=False, sipg_factor=Constant(1.0),
                  sipg_factor_vertical=Constant(1.0)):
-        """
+        """
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :kwarg bathymetry: bathymetry of the domain
         :type bathymetry: 3D :class:`Function` or :class:`Constant`
@@ -181,7 +181,7 @@ 
Source code for thetis.momentum_eq
 
 
 
[docs]class PressureGradientTerm(MomentumTerm):
-    r"""
+    r"""
     Internal pressure gradient term, :math:`g\nabla_h r`
 
     where :math:`r` is the baroclinic head :eq:`baroc_head`. Let :math:`s`
@@ -214,7 +214,7 @@ 
Source code for thetis.momentum_eq
 
 
 
[docs]class HorizontalAdvectionTerm(MomentumTerm):
-    r"""
+    r"""
     Horizontal advection term, :math:`\nabla_h \cdot (\textbf{u} \textbf{u})`
 
     The weak form reads
@@ -320,7 +320,7 @@ 
Source code for thetis.momentum_eq
 
 
 
[docs]class VerticalAdvectionTerm(MomentumTerm):
-    r"""
+    r"""
     Vertical advection term, :math:`\partial \left(w\textbf{u} \right)/(\partial z)`
 
     The weak form reads
@@ -368,7 +368,7 @@ 
Source code for thetis.momentum_eq
 
 
 
[docs]class HorizontalViscosityTerm(MomentumTerm):
-    r"""
+    r"""
     Horizontal viscosity term, :math:`- \nabla_h \cdot \left( \nu_h \nabla_h \textbf{u} \right)`
 
     Using the symmetric interior penalty method the weak form becomes
@@ -443,7 +443,7 @@ 
Source code for thetis.momentum_eq
 
 
 
[docs]class VerticalViscosityTerm(MomentumTerm):
-    r"""
+    r"""
     Vertical viscosity term, :math:`- \frac{\partial }{\partial z}\left( \nu \frac{\partial \textbf{u}}{\partial z}\right)`
 
     Using the symmetric interior penalty method the weak form becomes
@@ -494,7 +494,7 @@ 
Source code for thetis.momentum_eq
 
 
 
[docs]class BottomFrictionTerm(MomentumTerm):
-    r"""
+    r"""
     Quadratic bottom friction term, :math:`\tau_b = C_D \| \textbf{u}_b \| \textbf{u}_b`
 
     The weak formulation reads
@@ -551,7 +551,7 @@ 
Source code for thetis.momentum_eq
 
 
 
[docs]class LinearDragTerm(MomentumTerm):
-    r"""
+    r"""
     Linear drag term, :math:`\tau_b = D \textbf{u}_b`
 
     where :math:`D` is the drag coefficient, read from ``linear_drag_coefficient`` field.
@@ -572,7 +572,7 @@ 
Source code for thetis.momentum_eq
 
 
 
[docs]class CoriolisTerm(MomentumTerm):
-    r"""
+    r"""
     Coriolis term, :math:`f\textbf{e}_z\wedge \bar{\textbf{u}}`
     """
 
[docs]    def residual(self, solution, solution_old, fields, fields_old, bnd_conditions):
@@ -585,7 +585,7 @@ 
Source code for thetis.momentum_eq
 
 
 
[docs]class SourceTerm(MomentumTerm):
-    r"""
+    r"""
     Generic momentum source term
 
     The weak form reads
@@ -623,7 +623,7 @@ 
Source code for thetis.momentum_eq
 
 
 
[docs]class MomentumEquation(Equation):
-    """
+    """
     Hydrostatic 3D momentum equation :eq:`mom_eq_split` for mode split models
     """
     def __init__(self, function_space,
@@ -631,7 +631,7 @@ 
Source code for thetis.momentum_eq
                  use_nonlinear_equations=True, use_lax_friedrichs=True,
                  use_bottom_friction=False, sipg_factor=Constant(1.0),
                  sipg_factor_vertical=Constant(1.0)):
-        """
+        """
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :kwarg bathymetry: bathymetry of the domain
         :type bathymetry: 3D :class:`Function` or :class:`Constant`
@@ -663,7 +663,7 @@ 
Source code for thetis.momentum_eq
 
 
 
[docs]class InternalPressureGradientCalculator(MomentumTerm):
-    r"""
+    r"""
     Computes the internal pressure gradient term, :math:`g\nabla_h r`
 
     where :math:`r` is the baroclinic head :eq:`baroc_head`.
@@ -681,7 +681,7 @@ 
Source code for thetis.momentum_eq
     """
     def __init__(self, fields, bathymetry, bnd_functions,
                  internal_pg_scalar=None, solver_parameters=None):
-        """
+        """
         :arg solver: `class`FlowSolver` object
         :kwarg dict solver_parameters: PETSc solver options
         """
@@ -705,7 +705,7 @@ 
Source code for thetis.momentum_eq
         self.lin_solver = LinearVariationalSolver(prob, solver_parameters=solver_parameters)
 
 
[docs]    def solve(self):
-        """
+        """
         Computes internal pressure gradient and stores it in int_pg_3d field
         """
         self.lin_solver.solve()

@@ -760,8 +760,8 @@ 
Source code for thetis.momentum_eq
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/optimisation.html b/_modules/thetis/optimisation.html
index e8a9b0a..7bd9e31 100644
--- a/_modules/thetis/optimisation.html
+++ b/_modules/thetis/optimisation.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.optimisation — Thetis 0+untagged.1878.g52658ff.dirty documentation
+    thetis.optimisation — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -68,12 +68,38 @@ 
Source code for thetis.optimisation
 from .callback import DiagnosticCallback
 from .exporter import ExportManager
 import thetis.field_defs as field_defs
-from abc import abstractmethod
+from abc import abstractmethod, ABC
 import numpy
 
 
+
[docs]class OptimisationCallback(ABC):
+    """Base class for callback that can be used as callbacks of a :class:`ReducedFunctional`
+
+    Called at various stages during the optimisation process:
+    - eval_cb_pre(controls) and eval_cb_post(functional, controls)                    called before and after (re)evaluation of the forward model
+    - derivative_cb_pre(controls) and eval_cb_post(functional, derivative, controls)  called before and after the gradient computation using the adjoint of the model
+    - hessian_cb_pre(controls) and eval_cb_post(functional, derivative, controls)     called before and after the hessian computation
+    OptimisationCallbacks that (can) use controls, functional and derivative information, work out
+    what is provided by the number of arguments: current control values are always in the last argument;
+    if more than 2 arguments are provided, the first is the latest evaluated functional value.
+    """
+
[docs]    @abstractmethod
+    def callback(self, *args):
+        pass

+
+    def __call__(self, *args):
+        self.callback(*args)
+
+        # if used as derivative callback, we need to return the derivatives (args[1]) or the controls (args[0]) for resp. the post or pre callback
+        # for other callbacks it doesn't matter what we return
+        if len(args) > 1:
+            return args[1]
+        else:
+            return args[0]

+
+
 
[docs]class UserExportManager(ExportManager):
-    """
+    """
     ExportManager for user provided functions (not necessarily known to Thetis)
 
     In the standard :class:`.ExportManager` all provided functions need to have standard names
@@ -85,7 +111,7 @@ 
Source code for thetis.optimisation
     def __init__(self, solver_obj_or_outputdir, functions_to_export,
                  filenames=None, filename_prefix='',
                  shortnames=None, **kwargs):
-        """
+        """
         :arg solver_obj_or_outputdir: a :class:`.FlowSolver2d` object, used to determine the output directory. Alternatively, the
               outputdir can be specified with a string as the first argument.
         :arg functions_to_export: a list of :class:`Function` s
@@ -110,14 +136,14 @@ 
Source code for thetis.optimisation
         for field_name, function in zip(field_name_list, functions_to_export):
             field_dict[field_name] = function
             if shortnames is None and field_name in field_defs.field_metadata:
-                field_metadata[field_name] = {'shortname': field_defs.field_metadata[field_name]}
+                field_metadata[field_name] = {'shortname': field_defs.field_metadata[field_name]['shortname']}
             else:
                 field_metadata[field_name] = {'shortname': field_name}
 
         if filenames is None:
             for field_name in field_name_list:
                 if field_name in field_defs.field_metadata:
-                    field_metadata[field_name]['filename'] = field_defs.field_metadata['filename']
+                    field_metadata[field_name]['filename'] = filename_prefix + field_defs.field_metadata[field_name]['filename']
                 else:
                     field_metadata[field_name]['filename'] = filename_prefix + field_name
         else:
@@ -128,7 +154,7 @@ 
Source code for thetis.optimisation
 
 
 
[docs]class DeferredExportManager(object):
-    """
+    """
     A wrapper around a UserExportManager that is only created on the first export() call.
 
     In addition the functions provided in the export call are copied into a fixed set of functions,
@@ -136,7 +162,7 @@ 
Source code for thetis.optimisation
     function space). This is used in the :class:`.ControlsExportOptimisationCallback`
     and :class:`.DerivativesExportOptimisationCallback`."""
     def __init__(self, solver_obj_or_outputdir, **kwargs):
-        """
+        """
         :arg solver_obj_or_outputdir: a :class:`.FlowSolver2d` object, used to determine the output directory. Alternatively, the
               outputdir can be specified with a string as the first argument.
         :arg kwargs: any further keyword arguments are passed on to :class:`.UserExportManager`"""
@@ -145,7 +171,7 @@ 
Source code for thetis.optimisation
         self.export_manager = None
 
 
[docs]    def export(self, functions, suggested_names=None):
-        """
+        """
         Create the :class:`.UserExportManager` (first call only), and call its export() method.
 
         :arg functions: a list of :class:`Function` s that the :class:`.UserExportManager` will be based on. Their values
@@ -169,12 +195,12 @@ 
Source code for thetis.optimisation
         self.export_manager.export()

 
 
-
[docs]class UserExportOptimisationCallback(UserExportManager):
-    """A :class:`.UserExportManager` that can be used as a :class:`ReducedFunctional` callback
+
[docs]class UserExportOptimisationCallback(UserExportManager, OptimisationCallback):
+    """A :class:`.UserExportManager` that can be used as a :class:`ReducedFunctional` callback
 
     Any callback arguments (functional value, derivatives, controls) are ignored"""
     def __init__(self, solver_obj_or_outputdir, functions_to_export, **kwargs):
-        """
+        """
         :arg solver_obj_or_outputdir: a :class:`.FlowSolver2d` object, used to determine the output directory. Alternatively, the
               outputdir can be specified with a string as the first argument.
         :arg functions_to_export: a list of :class:`Function` s
@@ -186,70 +212,70 @@ 
Source code for thetis.optimisation
         # that determine what will be written
         self.orig_functions = self.functions.copy()
 
-    def __call__(self, *args):
-        """
+
[docs]    def callback(self, *args):
+        """
         Ensure the :class:`.UserExportManager` uses the checkpointed values and call its export().
 
         :args: these are ignored"""
         for name in self.fields_to_export:
             self.functions[name] = self.orig_functions[name].block_variable.saved_output
-        self.export()

+        self.export()

 
 
-
[docs]class ControlsExportOptimisationCallback(DeferredExportManager):
-    """A callback that exports the current control values (assumed to all be :class:`Function` s)
+
[docs]class ControlsExportOptimisationCallback(DeferredExportManager, OptimisationCallback):
+    """A callback that exports the current control values (assumed to all be :class:`Function` s)
 
     The control values are assumed to be the last argument in the callback (as for all :class:`ReducedFunctional` callbacks)."""
     def __init__(self, solver_obj_or_outputdir, **kwargs):
-        """
+        """
         :arg solver_obj_or_outputdir: a :class:`.FlowSolver2d` object, used to determine the output directory. Alternatively, the
               outputdir can be specified with a string as the first argument.
         :arg kwargs: any further keyword arguments are passed on to :class:`.UserExportManager`"""
         kwargs.setdefault('filename_prefix', 'control_')
         super().__init__(solver_obj_or_outputdir, **kwargs)
 
-    def __call__(self, *args):
-        self.export(args[-1])

+
[docs]    def callback(self, *args):
+        self.export(args[-1])

 
 
-
[docs]class DerivativesExportOptimisationCallback(DeferredExportManager):
-    """A callback that exports the derivatives calculated by the adjoint.
+
[docs]class DerivativesExportOptimisationCallback(DeferredExportManager, OptimisationCallback):
+    """A callback that exports the derivatives calculated by the adjoint.
 
     The derivatives are assumed to be the second argument in the callback. This can therefore
     be used as a derivative_cb_post callback in a :class:`ReducedFunctional`"""
     def __init__(self, solver_obj_or_outputdir, **kwargs):
-        """
+        """
         :arg solver_obj_or_outputdir: a :class:`.FlowSolver2d` object, used to determine the output directory. Alternatively, the
               outputdir can be specified with a string as the first argument.
         :arg kwargs: any further keyword arguments are passed on to :class:`.UserExportManager`"""
         kwargs.setdefault('filename_prefix', 'derivative_')
         super().__init__(solver_obj_or_outputdir, **kwargs)
 
-    def __call__(self, *args):
+
[docs]    def callback(self, *args):
         if len(args) != 3:
-            raise TypeError("DerivativesExportOptimsationCallback called with wrong number of arguments: should be used for derivative_cb_post callback only.")
+            raise TypeError("DerivativesExportOptimisationCallback called with wrong number of arguments: should be used for derivative_cb_post callback only.")
         try:
             # get name from controls args[-1]
             names = [function.name() for function in args[-1]]
         except (TypeError, NotImplementedError):
             # args[-1] is not a list but a single control
             names = [args[-1].name()]
-        self.export(args[1], suggested_names=names)

+        self.export(args[1], suggested_names=names)

 
 
-
[docs]class OptimisationCallbackList(list):
-    """
+
[docs]class OptimisationCallbackList(list, OptimisationCallback):
+    """
     A list of callbacks that can be used as a single callback itself.
 
     Calls all callbacks in order."""
-    def __call__(self, *args):
-        for callback in self:
-            callback(*args)

+
[docs]    def callback(self, *args):
+        for cb in self:
+            cb(*args)

 
 
-
[docs]class DiagnosticOptimisationCallback(DiagnosticCallback):
-    """
-    An OptimsationCallback similar to :class:`.DiagnosticCallback` that can be used as callback in a :class:`ReducedFunctional`.
+
[docs]class DiagnosticOptimisationCallback(OptimisationCallback, DiagnosticCallback):
+    """
+    An OptimisationCallback similar to :class:`.DiagnosticCallback` that can be used as callback in a :class:`ReducedFunctional`.
 
     Note that in this case the computing of the values needs to be defined in the compute_values method,
     not in the __call__ method (as this one is directly called from the :class:`ReducedFunctional`). In addition,
@@ -257,7 +283,7 @@ 
Source code for thetis.optimisation
     """
 
     def __init__(self, solver_obj, **kwargs):
-        """
+        """
         :arg solver_obj: Thetis solver object
         :arg kwargs: keyword arguments passed to :class:`.DiagnosticCallback`
         """
@@ -266,7 +292,7 @@ 
Source code for thetis.optimisation
 
 
[docs]    @abstractmethod
     def compute_values(self, *args):
-        """
+        """
         Compute diagnostic values.
 
         This method is to be implemented in concrete subclasses of a :class:`.DiagnosticOptimisationCallback`.
@@ -278,7 +304,7 @@ 
Source code for thetis.optimisation
         pass

 
 
[docs]    def evaluate(self, *args, index=None):
-        """Evaluates callback and pushes values to log and hdf file (if enabled)"""
+        """Evaluates callback and pushes values to log and hdf file (if enabled)"""
         values = self.compute_values(*args)
         if len(args) > 0:
             functional = args[0]
@@ -290,12 +316,12 @@ 
Source code for thetis.optimisation
         if self.append_to_hdf5:
             self.push_to_hdf5(functional, values)

 
-    def __call__(self, *args):
-        self.evaluate(*args)

+
[docs]    def callback(self, *args):
+        self.evaluate(*args)

 
 
 
[docs]class FunctionalOptimisationCallback(DiagnosticOptimisationCallback):
-    """
+    """
     A simple OptimisationCallback that records the functional value in the log and/or hdf5 file."""
     variable_names = ['functional']
     name = 'functional'
@@ -307,6 +333,40 @@ 
Source code for thetis.optimisation
 
 
[docs]    def message_str(self, functional):
         return 'Functional value: {}'.format(functional)

+
+
+
[docs]class ConstantControlOptimisationCallback(DiagnosticOptimisationCallback):
+    """
+    OptimisationCallback that records the control values (which are assumed to be a list of Constants) in the log and/or hdf5 file."""
+    variable_names = ['controls']
+    name = 'controls'
+
+
[docs]    def compute_values(self, *args):
+        controls = args[-1]
+        if self.array_dim != len(controls):
+            raise ValueError("Need array_dim argument in ConstantControlOptimisationCallback set to the number of controls")
+        return [[float(c) for c in controls]]

+
+
[docs]    def message_str(self, *controls):
+        return 'Controls value: {}'.format(controls)

+
+
+
[docs]class DerivativeConstantControlOptimisationCallback(DiagnosticOptimisationCallback):
+    """
+    OptimisationCallback that records the derivatives with respect to the controls, assumed to be a list of Constants, in the log and/or hdf5 file."""
+    variable_names = ['derivatives']
+    name = 'derivatives'
+
+
[docs]    def compute_values(self, *args):
+        if len(args) != 3:
+            raise TypeError("DerivativesExportOptimisationCallback called with wrong number of arguments: should be used for derivative_cb_post callback only.")
+        derivatives = args[1]
+        if self.array_dim != len(derivatives):
+            raise ValueError("Need array_dim argument in ConstantControlOptimisationCallback set to the number of controls")
+        return [[float(d) for d in derivatives]]

+
+
[docs]    def message_str(self, *derivatives):
+        return 'Derivatives: {}'.format(derivatives)

 

 
             

@@ -317,8 +377,8 @@ 
Source code for thetis.optimisation
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/options.html b/_modules/thetis/options.html
index 6a781a4..77c9a9f 100644
--- a/_modules/thetis/options.html
+++ b/_modules/thetis/options.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.options — Thetis 0+untagged.1888.gc4e776a documentation
+    thetis.options — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -58,13 +58,14 @@ 
Source code for thetis.options
 objects.
 """
 from .configuration import *
+from traitlets import Integer, Float, TraitType, Dict, Enum, Unicode, Bool, Instance, List, Type
 from firedrake import Constant
 from .sediment_model import SedimentModel
 from collections import OrderedDict
 
 
 
[docs]class TimeStepperOptions(FrozenHasTraits):
-    """Base class for all time stepper options"""
+    """Base class for all time stepper options"""
     name = 'Time stepper'
     solver_parameters = PETScSolverParameters({}).tag(config=True)
     ad_block_tag = Unicode(default_value=None, allow_none=True, help="""
@@ -74,22 +75,22 @@ 
Source code for thetis.options
 
 
 
[docs]class ExplicitTimeStepperOptions(TimeStepperOptions):
-    """Options for explicit time integrator"""
+    """Options for explicit time integrator"""
     use_automatic_timestep = Bool(True, help='Set time step automatically based on local CFL conditions.').tag(config=True)

 
 
 
[docs]class ExplicitTimeStepperOptions2d(ExplicitTimeStepperOptions):
-    """Options for 2d explicit time integrator"""

+    """Options for 2d explicit time integrator"""

 
 
 
[docs]class SemiImplicitTimeStepperOptions2d(TimeStepperOptions):
-    """Options for 2d semi-implicit time integrator"""
+    """Options for 2d semi-implicit time integrator"""
     use_semi_implicit_linearization = Bool(
         False, help="Use linearized semi-implicit time integration").tag(config=True)

 
 
 
[docs]class SemiImplicitSWETimeStepperOptions2d(SemiImplicitTimeStepperOptions2d):
-    """
+    """
     Options for 2d semi-implicit time integrator
     applied to shallow water equations
     """
@@ -101,7 +102,7 @@ 
Source code for thetis.options
 
 
 
[docs]class SemiImplicitTracerTimeStepperOptions2d(SemiImplicitTimeStepperOptions2d):
-    """
+    """
     Options for 2d semi-implicit time integrator
     applied to tracer equations
     """
@@ -112,16 +113,17 @@ 
Source code for thetis.options
 
 
 
[docs]class SteadyStateTimeStepperOptions2d(TimeStepperOptions):
-    """Options for 2d steady state solver"""
+    """Options for 2d steady state solver"""
     solver_parameters = PETScSolverParameters({
         'ksp_type': 'preonly',
         'pc_type': 'lu',
+        'pc_factor_mat_solver_type': 'mumps',
         'mat_type': 'aij'
     }).tag(config=True)

 
 
 
[docs]class CrankNicolsonSWETimeStepperOptions2d(SemiImplicitSWETimeStepperOptions2d):
-    """
+    """
     Options for 2d Crank-Nicolson time integrator
     applied to shallow water equations
     """
@@ -131,7 +133,7 @@ 
Source code for thetis.options
 
 
 
[docs]class CrankNicolsonTracerTimeStepperOptions2d(SemiImplicitTracerTimeStepperOptions2d):
-    """
+    """
     Options for 2d Crank-Nicolson time integrator
     applied to tracer equations
     """
@@ -141,7 +143,7 @@ 
Source code for thetis.options
 
 
 
[docs]class PressureProjectionSWETimeStepperOptions2d(TimeStepperOptions):
-    """
+    """
     Options for 2d pressure-projection time integrator
     applied to shallow water equations
     """
@@ -189,7 +191,7 @@ 
Source code for thetis.options
 
 
 
[docs]class ExplicitSWETimeStepperOptions2d(ExplicitTimeStepperOptions2d):
-    """
+    """
     Options for 2d explicit time integrator
     applied to shallow water equations
     """
@@ -204,7 +206,7 @@ 
Source code for thetis.options
 
 
 
[docs]class ExplicitTracerTimeStepperOptions2d(ExplicitTimeStepperOptions2d):
-    """
+    """
     Options for 2d explicit time integrator
     applied to tracer equations
     """
@@ -215,7 +217,7 @@ 
Source code for thetis.options
 
 
 
[docs]class IMEXSWETimeStepperOptions2d(SemiImplicitTimeStepperOptions2d):
-    """
+    """
     Options for 2d implicit-explicit time integrators
     applied to shallow water equations
     """
@@ -228,7 +230,7 @@ 
Source code for thetis.options
 
 
 
[docs]class TimeStepperOptions3d(FrozenHasTraits):
-    """Options for 3d explicit time integrator"""
+    """Options for 3d explicit time integrator"""
     ad_block_tag = Unicode(default_value=None, allow_none=True, help="""
         Custom tag for the Pyadjoint blocks generated by
         this timestepper.
@@ -236,7 +238,7 @@ 
Source code for thetis.options
 
 
 
[docs]class SWETimeStepperOptions3d(TimeStepperOptions3d):
-    """
+    """
     Options for 3d time integrator
     shallow water component
     """
@@ -248,7 +250,7 @@ 
Source code for thetis.options
 
 
 
[docs]class ImplicitSWETimeStepperOptions3d(SWETimeStepperOptions3d):
-    """
+    """
     Options for 3d implicit time integrator
     shallow water component
     """
@@ -258,7 +260,7 @@ 
Source code for thetis.options
 
 
 
[docs]class ExplicitMomentumTimeStepperOptions3d(TimeStepperOptions3d):
-    """
+    """
     Options for 3d explicit time integrator
     momentum component
     """
@@ -272,7 +274,7 @@ 
Source code for thetis.options
 
 
 
[docs]class ImplicitMomentumTimeStepperOptions3d(TimeStepperOptions3d):
-    """
+    """
     Options for 3d implicit time integrator
     momentum component
     """
@@ -286,7 +288,7 @@ 
Source code for thetis.options
 
 
 
[docs]class ExplicitTracerTimeStepperOptions3d(TimeStepperOptions3d):
-    """
+    """
     Options for 3d explicit time integrator
     applied to tracer equations
     """
@@ -300,7 +302,7 @@ 
Source code for thetis.options
 
 
 
[docs]class ImplicitTracerTimeStepperOptions3d(TimeStepperOptions3d):
-    """
+    """
     Options for 3d implicit time integrator
     applied to tracer equations
     """
@@ -314,7 +316,7 @@ 
Source code for thetis.options
 
 
 
[docs]class SSPRKTimeStepperOptions3d(TimeStepperOptions3d):
-    """Options for 3d SSPRK coupled time integrator"""
+    """Options for 3d SSPRK coupled time integrator"""
     use_automatic_timestep = Bool(True, help='Set time step automatically based on local CFL conditions.').tag(config=True)
     swe_options = SWETimeStepperOptions3d()
     implicit_momentum_options = ImplicitMomentumTimeStepperOptions3d()
@@ -324,17 +326,17 @@ 
Source code for thetis.options
 
 
 
[docs]class LeapFrogTimeStepperOptions3d(SSPRKTimeStepperOptions3d):
-    """Options for 3d LeapFrog coupled time integrator"""
+    """Options for 3d LeapFrog coupled time integrator"""
     swe_options = ImplicitSWETimeStepperOptions3d()

 
 
 
[docs]class TurbulenceModelOptions(FrozenHasTraits):
-    """Abstract base class for all turbulence model options"""
+    """Abstract base class for all turbulence model options"""
     name = 'Turbulence closure model'

 
 
 
[docs]class PacanowskiPhilanderModelOptions(TurbulenceModelOptions):
-    """Options for Pacanowski-Philander turbulence model"""
+    """Options for Pacanowski-Philander turbulence model"""
     name = 'Pacanowski-Philander turbulence closure model'
     max_viscosity = PositiveFloat(5e-2, help=r"float: Constant maximum viscosity :math:`\nu_{max}`").tag(config=True)
     alpha = PositiveFloat(10.0, help="float: Richardson number multiplier").tag(config=True)
@@ -342,7 +344,7 @@ 
Source code for thetis.options
 
 
 
[docs]class GLSModelOptions(TurbulenceModelOptions):
-    """Options for Generic Length Scale turbulence model"""
+    """Options for Generic Length Scale turbulence model"""
     name = 'Generic Length Scale turbulence closure model'
     closure_name = Enum(
         ['k-epsilon', 'k-omega', 'Generic Length Scale'],
@@ -418,7 +420,7 @@ 
Source code for thetis.options
                          help='bool: limit minimum turbulent length scale to len_min').tag(config=True)
 
 
[docs]    def apply_defaults(self, closure_name):
-        """
+        """
         Applies default parameters for given closure name
 
         :arg closure_name: name of the turbulence closure model
@@ -492,12 +494,12 @@ 
Source code for thetis.options
 
 
 
[docs]class EquationOfStateOptions(FrozenHasTraits):
-    """Base class of equation of state options"""
+    """Base class of equation of state options"""
     name = 'Equation of State'

 
 
 
[docs]class LinearEquationOfStateOptions(EquationOfStateOptions):
-    """Linear equation of state options"""
+    """Linear equation of state options"""
     # TODO more human readable parameter names
     # TODO document the actual equation somewhere
     name = 'Linear Equation of State'
@@ -509,25 +511,62 @@ 
Source code for thetis.options
 
 
 
[docs]class TidalTurbineOptions(FrozenHasTraits):
-    """Tidal turbine parameters"""
-    thrust_coefficient = PositiveFloat(
-        0.8, help='Thrust coefficient C_T').tag(config=True)
+    """Tidal turbine parameters"""
+    name = 'Tidal turbine options'
     diameter = PositiveFloat(
-        18., help='Turbine diameter').tag(config=True)

+        18., help='Turbine diameter').tag(config=True)
+    C_support = NonNegativeFloat(
+        0., help='Thrust coefficient for support structure').tag(config=True)
+    A_support = NonNegativeFloat(
+        0., help='Cross section of support structure').tag(config=True)

+
+
+
[docs]class ConstantTidalTurbineOptions(TidalTurbineOptions):
+    """Options for tidal turbine with constant thrust"""
+    name = 'Constant tidal turbine options'
+    thrust_coefficient = PositiveFloat(
+        0.8, help='Thrust coefficient C_T').tag(config=True)

+
+
+
[docs]class TabulatedTidalTurbineOptions(TidalTurbineOptions):
+    """Options for tidal turbine with tabulated thrust coefficient"""
+    name = 'Tabulated tidal turbine options'
+    thrust_coefficients = List([3.0], help='Table of thrust coefficients')
+    thrust_speeds = List(
+        [0.8, 0.8],
+        help="""List of speeds at which thrust_coefficients are applied.
+        First and last entry function as cut-in and cut-out speeds respectively""")

 
 
-
[docs]class TidalTurbineFarmOptions(FrozenHasTraits, TraitType):
-    """Tidal turbine farm options"""
+
[docs]@attach_paired_options("turbine_type",
+                       PairedEnum([('constant', ConstantTidalTurbineOptions),
+                                   ('table', TabulatedTidalTurbineOptions),
+                                   ],
+                                  "turbine_options",
+                                  default_value='constant',
+                                  help='Type of turbine thrust specification').tag(config=True),
+                       Instance(TidalTurbineOptions, args=()).tag(config=True))
+class TidalTurbineFarmOptions(FrozenHasTraits, TraitType):
+    """Tidal turbine farm options"""
     name = 'Farm options'
-    turbine_options = TidalTurbineOptions()
     turbine_density = FiredrakeScalarExpression(
         Constant(0.0), help='Density of turbines within the farm')
     break_even_wattage = NonNegativeFloat(
         0.0, help='Average power production per turbine required to break even')

 
 
+
[docs]class DiscreteTidalTurbineFarmOptions(TidalTurbineFarmOptions):
+    """Discrete Tidal turbine farm options - defaults to 0 turbines in the field"""
+    name = 'Discrete Farm options'
+    turbine_coordinates = List(default_value=[], help="Coordinates for turbines").tag(config=True)
+    upwind_correction = Bool(True,
+                             help='bool: Apply flow correction to correct for upwind velocity').tag(config=True)
+    quadrature_degree = PositiveInteger(10,
+                                        help='Quadrature degree for thrust force and power output integral').tag(config=True)

+
+
 
[docs]class TracerFieldOptions(FrozenHasTraits):
-    """Tracer field options"""
+    """Tracer field options"""
     name = 'Tracer options'
     function = FiredrakeScalarExpression(
         None, allow_none=True, help='Firedrake Function representing the tracer')
@@ -559,7 +598,7 @@ 
Source code for thetis.options
                                   help='Name of the free surface time integrator').tag(config=True),
                        Instance(TimeStepperOptions, args=()).tag(config=True))
 class NonhydrostaticModelOptions(FrozenHasTraits):
-    """Options for non-hydrostatic models"""
+    """Options for non-hydrostatic models"""
     name = 'Non-hydrostatic 2D/3D models'
     solve_nonhydrostatic_pressure = Bool(False, help='Solve equations with the non-hydrostatic pressure').tag(config=True)
     q_degree = NonNegativeInteger(
@@ -576,7 +615,7 @@ 
Source code for thetis.options
 
 
 
[docs]class CommonModelOptions(FrozenConfigurable):
-    """Options that are common for both 2d and 3d models"""
+    """Options that are common for both 2d and 3d models"""
     name = 'Model options'
     nh_model_options = Instance(NonhydrostaticModelOptions, args=()).tag(config=True)
     polynomial_degree = NonNegativeInteger(1, help='Polynomial degree of elements').tag(config=True)
@@ -626,6 +665,12 @@ 
Source code for thetis.options
         1.0, help="Factor to scale the 2d time step OBSOLETE").tag(config=True)  # TODO OBSOLETE
     cfl_3d = PositiveFloat(
         1.0, help="Factor to scale the 2d time step OBSOLETE").tag(config=True)  # TODO OBSOLETE
+    simulation_initial_date = DatetimeTraitlet(
+        default_value=None, allow_none=True,
+        help="Model initialization date. Corresponds to zero in simulation time.").tag(config=True)
+    simulation_end_date = DatetimeTraitlet(
+        default_value=None, allow_none=True,
+        help="Simulation end date").tag(config=True)
     simulation_export_time = PositiveFloat(
         100.0, help="""
         Export interval in seconds
@@ -634,7 +679,7 @@ 
Source code for thetis.options
         diagnostics will be computed
         """).tag(config=True)
     simulation_end_time = PositiveFloat(
-        1000.0, help="Simulation duration in seconds").tag(config=True)
+        None, allow_none=True, help="Simulation duration in seconds").tag(config=True)
     horizontal_velocity_scale = FiredrakeConstantTraitlet(
         Constant(0.1), help="""
         Maximum horizontal velocity magnitude
@@ -853,7 +898,7 @@ 
Source code for thetis.options
                                   help='Name of the tracer time integrator').tag(config=True),
                        Instance(TimeStepperOptions, args=()).tag(config=True))
 class ModelOptions2d(CommonModelOptions):
-    """Options for 2D depth-averaged shallow water model"""
+    """Options for 2D depth-averaged shallow water model"""
     name = 'Depth-averaged 2D model'
     sediment_model_options = Instance(SedimentModelOptions, args=()).tag(config=True)
     use_tracer_conservative_form = Bool(False, help='Solve 2D tracer transport in the conservative form').tag(config=True)
@@ -891,6 +936,11 @@ 
Source code for thetis.options
         """).tag(config=True)
     tidal_turbine_farms = Dict(trait=TidalTurbineFarmOptions(),
                                default_value={}, help='Dictionary mapping subdomain ids to the options of the corresponding farm')
+
+    discrete_tidal_turbine_farms = Dict(trait=DiscreteTidalTurbineFarmOptions(),
+                                        default_value={},
+                                        help='Dictionary mapping subdomain ids to the options of the corresponding farm')
+
     check_tracer_conservation = Bool(
         False, help="""
         Compute total tracer mass at every export
@@ -928,10 +978,11 @@ 
Source code for thetis.options
 
     def __init__(self, *args, **kwargs):
         self.tracer = OrderedDict()
+        self.tracer_fields = OrderedDict()
         super().__init__(*args, **kwargs)
 
 
[docs]    def add_tracer_2d(self, label, name, filename, shortname=None, unit='-', **kwargs):
-        """
+        """
         Add a 2D tracer field to :attr:`tracer`.
 
         Note that the tracer equations will be solved in the order in which they are added.
@@ -953,6 +1004,7 @@ 
Source code for thetis.options
         assert isinstance(unit, str)
         assert label not in self.tracer, f"Field '{label}' already exists."
         assert ' ' not in label, "Labels cannot contain spaces"
+        assert ',' not in label, "Labels cannot contain commas"
         assert ' ' not in filename, "Filenames cannot contain spaces"
         self.tracer[label] = TracerFieldOptions()
         self.tracer[label].metadata = {
@@ -964,10 +1016,42 @@ 
Source code for thetis.options
         self.tracer[label].function = kwargs.get('function')
         self.tracer[label].source = kwargs.get('source')
         self.tracer[label].diffusivity = kwargs.get('diffusivity')
-        self.tracer[label].use_conservative_form = kwargs.get('use_conservative_form', False)

+        self.tracer[label].use_conservative_form = kwargs.get('use_conservative_form', False)
+        if not kwargs.get('mixed', False):
+            self.tracer_fields[label] = self.tracer[label].function

+
+
[docs]    def add_tracer_system_2d(self, labels, names, filenames, shortnames=None, units=None, function=None, **kwargs):
+        """
+        Add multiple 2D tracer fields to :attr:`tracer` at once.
+
+        The equations associated with these fields will be solved as a mixed system.
+
+        :arg labels: a list of field labels used internally by Thetis
+        :arg names: a list of human readable names for the tracer fields
+        :arg filenames: a list of file names for outputs
+        :kwarg shortnames: a list of short versions of the names
+        :kwarg units: a list of units for fields
+        :kwarg function: :class:`Function` to use for the mixed tracer
+        :kwargs: other parameters passed to :meth:`add_tracer_2d`
+        """
+        N = len(labels)
+        shortnames = shortnames or names
+        units = units or ['-']*N
+        assert len(names) == len(filenames) == len(shortnames) == len(units) == N
+        if kwargs == {}:
+            kwargs = {label: {} for label in labels}
+        assert set(kwargs.keys()).issubset(set(labels))
+        for label in labels:
+            kwargs[label]['mixed'] = True
+        self.tracer_fields[','.join(labels)] = function
+        if function is not None:
+            for label, child in zip(labels, function.subfunctions):
+                kwargs[label]['function'] = child
+        for label, name, filename, shortname, unit in zip(labels, names, filenames, shortnames, units):
+            self.add_tracer_2d(label, name, filename, shortname=shortname, unit=unit, **kwargs[label])

 
 
[docs]    def set_timestepper_type(self, timestepper_type, **kwargs):
-        """
+        """
         Set the same timestepper type for all components.
 
         Any keyword arguments are passed through to the
@@ -1014,7 +1098,7 @@ 
Source code for thetis.options
                                   help='Type of equation of state').tag(config=True),
                        Instance(EquationOfStateOptions, args=()).tag(config=True))
 class ModelOptions3d(CommonModelOptions):
-    """Options for 3D hydrostatic model"""
+    """Options for 3D hydrostatic model"""
     name = '3D hydrostatic model'
     solve_salinity = Bool(True, help='Solve salinity transport').tag(config=True)
     solve_temperature = Bool(True, help='Solve temperature transport').tag(config=True)
@@ -1138,8 +1222,8 @@ 
Source code for thetis.options
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/rungekutta.html b/_modules/thetis/rungekutta.html
index f5b11fa..7f61736 100644
--- a/_modules/thetis/rungekutta.html
+++ b/_modules/thetis/rungekutta.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.rungekutta — Thetis 0+untagged.1888.gc4e776a documentation
+    thetis.rungekutta — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -58,13 +58,13 @@ 
Source code for thetis.rungekutta
 coefficients, and can be used to implement generic time integrators.
 """
 from .timeintegrator import *
-from abc import ABCMeta, abstractproperty, abstractmethod
+from abc import ABC, abstractproperty, abstractmethod
 import operator
 import numpy
 
 
 
[docs]def butcher_to_shuosher_form(a, b):
-    r"""
+    r"""
     Converts Butcher tableau to Shu-Osher form.
 
     The Shu-Osher form of a s-stage scheme is defined by two s+1 by s+1 arrays
@@ -140,8 +140,8 @@ 
Source code for thetis.rungekutta
     return alpha, beta

 
 
-
[docs]class AbstractRKScheme(object):
-    """
+
[docs]class AbstractRKScheme(ABC):
+    """
     Abstract class for defining Runge-Kutta schemes.
 
     Derived classes must define the Butcher tableau (arrays :attr:`a`, :attr:`b`,
@@ -149,26 +149,24 @@ 
Source code for thetis.rungekutta
 
     Currently only explicit or diagonally implicit schemes are supported.
     """
-    __metaclass__ = ABCMeta
-
     @abstractproperty
     def a(self):
-        """Runge-Kutta matrix :math:`a_{i,j}` of the Butcher tableau"""
+        """Runge-Kutta matrix :math:`a_{i,j}` of the Butcher tableau"""
         pass
 
     @abstractproperty
     def b(self):
-        """weights :math:`b_{i}` of the Butcher tableau"""
+        """weights :math:`b_{i}` of the Butcher tableau"""
         pass
 
     @abstractproperty
     def c(self):
-        """nodes :math:`c_{i}` of the Butcher tableau"""
+        """nodes :math:`c_{i}` of the Butcher tableau"""
         pass
 
     @abstractproperty
     def cfl_coeff(self):
-        """
+        """
         CFL number of the scheme
 
         Value 1.0 corresponds to Forward Euler time step.
@@ -195,7 +193,7 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class ForwardEulerAbstract(AbstractRKScheme):
-    """
+    """
     Forward Euler method
     """
     a = [[0]]
@@ -205,7 +203,7 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class BackwardEulerAbstract(AbstractRKScheme):
-    """
+    """
     Backward Euler method
     """
     a = [[1.0]]
@@ -215,7 +213,7 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class ImplicitMidpointAbstract(AbstractRKScheme):
-    r"""
+    r"""
     Implicit midpoint method, second order.
 
     This method has the Butcher tableau
@@ -234,7 +232,7 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class CrankNicolsonAbstract(AbstractRKScheme):
-    """
+    """
     Crack-Nicolson scheme
     """
     a = [[0.0, 0.0],
@@ -245,7 +243,7 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class DIRK22Abstract(AbstractRKScheme):
-    r"""
+    r"""
     2-stage, 2nd order, L-stable Diagonally Implicit Runge Kutta method
 
     This method has the Butcher tableau
@@ -274,7 +272,7 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class DIRK23Abstract(AbstractRKScheme):
-    r"""
+    r"""
     2-stage, 3rd order Diagonally Implicit Runge Kutta method
 
     This method has the Butcher tableau
@@ -303,7 +301,7 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class DIRK33Abstract(AbstractRKScheme):
-    """
+    """
     3-stage, 3rd order, L-stable Diagonally Implicit Runge Kutta method
 
     From DIRK(3,4,3) IMEX scheme in Ascher et al. (1997)
@@ -324,7 +322,7 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class DIRK43Abstract(AbstractRKScheme):
-    """
+    """
     4-stage, 3rd order, L-stable Diagonally Implicit Runge Kutta method
 
     From DIRK(4,4,3) IMEX scheme in Ascher et al. (1997)
@@ -343,7 +341,7 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class DIRKLSPUM2Abstract(AbstractRKScheme):
-    """
+    """
     DIRKLSPUM2, 3-stage, 2nd order, L-stable Diagonally Implicit Runge Kutta method
 
     From IMEX RK scheme (17) in Higureras et al. (2014).
@@ -361,7 +359,7 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class DIRKLPUM2Abstract(AbstractRKScheme):
-    """
+    """
     DIRKLPUM2, 3-stage, 2nd order, L-stable Diagonally Implicit Runge Kutta method
 
     From IMEX RK scheme (20) in Higureras et al. (2014).
@@ -379,7 +377,7 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class SSPRK33Abstract(AbstractRKScheme):
-    r"""
+    r"""
     3rd order Strong Stability Preserving Runge-Kutta scheme, SSP(3,3).
 
     This scheme has Butcher tableau
@@ -403,7 +401,7 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class ERKLSPUM2Abstract(AbstractRKScheme):
-    """
+    """
     ERKLSPUM2, 3-stage, 2nd order Explicit Runge Kutta method
 
     From IMEX RK scheme (17) in Higureras et al. (2014).
@@ -421,7 +419,7 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class ERKLPUM2Abstract(AbstractRKScheme):
-    """
+    """
     ERKLPUM2, 3-stage, 2nd order
     Explicit Runge Kutta method
 
@@ -463,20 +461,18 @@ 
Source code for thetis.rungekutta
     cfl_coeff = CFL_UNCONDITIONALLY_STABLE

 
 
-
[docs]class RungeKuttaTimeIntegrator(TimeIntegrator):
-    """Abstract base class for all Runge-Kutta time integrators"""
-    __metaclass__ = ABCMeta
-
+
[docs]class RungeKuttaTimeIntegrator(TimeIntegrator, ABC):
+    """Abstract base class for all Runge-Kutta time integrators"""
 
[docs]    @abstractmethod
     def get_final_solution(self, additive=False):
-        """
+        """
         Evaluates the final solution
         """
         pass

 
 
[docs]    @abstractmethod
     def solve_stage(self, i_stage, t, update_forcings=None):
-        """
+        """
         Solves a single stage of step from t to t+dt.
         All functions that the equation depends on must be at right state
         corresponding to each sub-step.
@@ -484,7 +480,7 @@ 
Source code for thetis.rungekutta
         pass

 
 
[docs]    def advance(self, t, update_forcings=None):
-        """Advances equations for one time step."""
+        """Advances equations for one time step."""
         if not self._initialized:
             self.initialize(self.solution)
         for i in range(self.n_stages):
@@ -493,7 +489,7 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class DIRKGeneric(RungeKuttaTimeIntegrator):
-    """
+    """
     Generic implementation of Diagonally Implicit Runge Kutta schemes.
 
     All derived classes must define the Butcher tableau coefficients :attr:`a`,
@@ -501,7 +497,7 @@ 
Source code for thetis.rungekutta
     """
     @PETSc.Log.EventDecorator("thetis.DIRKGeneric.__init__")
     def __init__(self, equation, solution, fields, dt, options, bnd_conditions, terms_to_add='all'):
-        """
+        """
         :arg equation: the equation to solve
         :type equation: :class:`Equation` object
         :arg solution: :class:`Function` where solution will be stored
@@ -582,7 +578,7 @@ 
Source code for thetis.rungekutta
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.DIRKGeneric.update_solver")
     def update_solver(self):
-        """Create solver objects"""
+        """Create solver objects"""
         self.solver = []
         for i in range(self.n_stages):
             p = NonlinearVariationalProblem(self.F[i], self.k[i])
@@ -595,13 +591,13 @@ 
Source code for thetis.rungekutta
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.DIRKGeneric.initialize")
     def initialize(self, init_cond):
-        """Assigns initial conditions to all required fields."""
+        """Assigns initial conditions to all required fields."""
         self.solution_old.assign(init_cond)
         self._initialized = True

 
 
[docs]    @PETSc.Log.EventDecorator("thetis.DIRKGeneric.update_solution")
     def update_solution(self, i_stage):
-        """
+        """
         Updates solution to i_stage sub-stage.
 
         Tendencies must have been evaluated first.
@@ -610,7 +606,7 @@ 
Source code for thetis.rungekutta
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.DIRKGeneric.solve_tendency")
     def solve_tendency(self, i_stage, t, update_forcings=None):
-        """
+        """
         Evaluates the tendency of i-th stage.
         """
         if i_stage == 0:
@@ -624,12 +620,12 @@ 
Source code for thetis.rungekutta
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.DIRKGeneric.get_final_solution")
     def get_final_solution(self):
-        """Assign final solution to :attr:`self.solution`"""
+        """Assign final solution to :attr:`self.solution`"""
         self.solution.assign(self.final_sol_expr)

 
 
[docs]    @PETSc.Log.EventDecorator("thetis.DIRKGeneric.solve_stage")
     def solve_stage(self, i_stage, t, update_forcings=None):
-        """Solve i-th stage and assign solution to :attr:`self.solution`."""
+        """Solve i-th stage and assign solution to :attr:`self.solution`."""
         self.solve_tendency(i_stage, t, update_forcings)
         self.update_solution(i_stage)

 
@@ -639,7 +635,7 @@ 
Source code for thetis.rungekutta
 
     @PETSc.Log.EventDecorator("thetis.DIRKGenericUForm.__init__")
     def __init__(self, equation, solution, fields, dt, options, bnd_conditions, terms_to_add='all'):
-        """
+        """
         :arg equation: the equation to solve
         :type equation: :class:`Equation` object
         :arg solution: :class:`Function` where solution will be stored
@@ -709,7 +705,7 @@ 
Source code for thetis.rungekutta
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.DIRKGenericUForm.update_solver")
     def update_solver(self):
-        """Create solver objects"""
+        """Create solver objects"""
         # Ensure LU assembles monolithic matrices
         if self.solver_parameters.get('pc_type') == 'lu':
             self.solver_parameters['mat_type'] = 'aij'
@@ -739,19 +735,19 @@ 
Source code for thetis.rungekutta
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.DIRKGenericUForm.initialize")
     def initialize(self, init_cond):
-        """Assigns initial conditions to all required fields."""
+        """Assigns initial conditions to all required fields."""
         self.solution_old.assign(init_cond)
         self._initialized = True

 
 
[docs]    def get_final_solution(self, additive=False):
-        """
+        """
         Evaluates the final solution
         """
         pass

 
 
[docs]    @PETSc.Log.EventDecorator("thetis.DIRKGenericUForm.solve_stage")
     def solve_stage(self, i_stage, t, update_forcings=None):
-        """
+        """
         Solves a single stage of step from t to t+dt.
         All functions that the equation depends on must be at right state
         corresponding to each sub-step.
@@ -817,14 +813,14 @@ 
Source code for thetis.rungekutta
 
 
 
[docs]class ERKGeneric(RungeKuttaTimeIntegrator):
-    """
+    """
     Generic explicit Runge-Kutta time integrator.
 
     Implements the Butcher form. All terms in the equation are treated explicitly.
     """
     @PETSc.Log.EventDecorator("thetis.ERKGeneric.__init__")
     def __init__(self, equation, solution, fields, dt, options, bnd_conditions, terms_to_add='all'):
-        """
+        """
         :arg equation: the equation to solve
         :type equation: :class:`Equation` object
         :arg solution: :class:`Function` where solution will be stored
@@ -875,13 +871,13 @@ 
Source code for thetis.rungekutta
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.ERKGeneric.initialize")
     def initialize(self, solution):
-        """Assigns initial conditions to all required fields."""
+        """Assigns initial conditions to all required fields."""
         self.solution_old.assign(solution)
         self._initialized = True

 
 
[docs]    @PETSc.Log.EventDecorator("thetis.ERKGeneric.update_solution")
     def update_solution(self, i_stage, additive=False):
-        """
+        """
         Computes the solution of the i-th stage
 
         Tendencies must have been evaluated first.
@@ -896,7 +892,7 @@ 
Source code for thetis.rungekutta
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.ERKGeneric.solve_tendency")
     def solve_tendency(self, i_stage, t, update_forcings=None):
-        """
+        """
         Evaluates the tendency of i-th stage
         """
         if self._nontrivial:
@@ -906,7 +902,7 @@ 
Source code for thetis.rungekutta
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.ERKGeneric.get_final_solution")
     def get_final_solution(self, additive=False):
-        """Assign final solution to :attr:`self.solution`
+        """Assign final solution to :attr:`self.solution`
 
         If additive=False, will overwrite :attr:`solution` function, otherwise
         will add to it.
@@ -919,20 +915,20 @@ 
Source code for thetis.rungekutta
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.ERKGeneric.solve_stage")
     def solve_stage(self, i_stage, t, update_forcings=None):
-        """Solve i-th stage and assign solution to :attr:`self.solution`."""
+        """Solve i-th stage and assign solution to :attr:`self.solution`."""
         self.update_solution(i_stage)
         self.solve_tendency(i_stage, t, update_forcings)

 
 
 
[docs]class ERKGenericShuOsher(TimeIntegrator):
-    """
+    """
     Generic explicit Runge-Kutta time integrator.
 
     Implements the Shu-Osher form.
     """
     @PETSc.Log.EventDecorator("thetis.ERKGenericShuOsher.__init__")
     def __init__(self, equation, solution, fields, dt, options, bnd_conditions, terms_to_add='all'):
-        """
+        """
         :arg equation: the equation to solve
         :type equation: :class:`Equation` object
         :arg solution: :class:`Function` where solution will be stored
@@ -985,7 +981,7 @@ 
Source code for thetis.rungekutta
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.ERKGenericShuOsher.solve_stage")
     def solve_stage(self, i_stage, t, update_forcings=None):
-        """Solve i-th stage and assign solution to :attr:`self.solution`."""
+        """Solve i-th stage and assign solution to :attr:`self.solution`."""
         if self._nontrivial:
             if update_forcings is not None:
                 update_forcings(t + self.c[i_stage]*self.dt)
@@ -1004,7 +1000,7 @@ 
Source code for thetis.rungekutta
 
 
[docs]    @PETSc.Log.EventDecorator("thetis.ERKGenericShuOsher.advance")
     def advance(self, t, update_forcings=None):
-        """Advances equations for one time step."""
+        """Advances equations for one time step."""
         for i in range(self.n_stages):
             self.solve_stage(i, t, update_forcings)

 
@@ -1045,8 +1041,8 @@ 
Source code for thetis.rungekutta
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/sediment_eq_2d.html b/_modules/thetis/sediment_eq_2d.html
index 7d251bc..9c7b65b 100644
--- a/_modules/thetis/sediment_eq_2d.html
+++ b/_modules/thetis/sediment_eq_2d.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.sediment_eq_2d — Thetis 0+untagged.1878.g52658ff.dirty documentation
+    thetis.sediment_eq_2d — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -68,8 +68,7 @@ 
Source code for thetis.sediment_eq_2d
 velocities, and :math:`\mu_h` denotes horizontal diffusivity.
 """
 from .equation import Equation
-from .tracer_eq_2d import HorizontalDiffusionTerm, HorizontalAdvectionTerm, TracerTerm
-from .conservative_tracer_eq_2d import ConservativeHorizontalAdvectionTerm
+from .tracer_eq_2d import HorizontalDiffusionTerm, HorizontalAdvectionTerm, ConservativeHorizontalAdvectionTerm, TracerTerm
 
 __all__ = [
     'SedimentEquation2D',
@@ -82,17 +81,17 @@ 
Source code for thetis.sediment_eq_2d
 
 
 
[docs]class SedimentTerm(TracerTerm):
-    """
+    """
     Generic sediment term that provides commonly used members.
     """
     def __init__(self, function_space, depth, options, sediment_model, conservative=False):
-        """
+        """
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :arg depth: :class:`DepthExpression` containing depth info
         :arg options: :class`ModelOptions2d` containing parameters
         :kwarg bool conservative: whether to use conservative tracer
         """
-        super(SedimentTerm, self).__init__(function_space, depth, options)
+        super(SedimentTerm, self).__init__(0, 'sediment_2d', function_space, depth, options)
         self.sediment_model = sediment_model
         self.conservative = conservative
 
@@ -109,7 +108,7 @@ 
Source code for thetis.sediment_eq_2d
 
 
 
[docs]class ConservativeSedimentAdvectionTerm(SedimentTerm, ConservativeHorizontalAdvectionTerm):
-    """
+    """
     Advection term for sediment equation
 
     Same as :class:`ConservativeHorizontalAdvectionTerm` but allows for equilibrium boundary condition
@@ -118,7 +117,7 @@ 
Source code for thetis.sediment_eq_2d
 
 
 
[docs]class SedimentAdvectionTerm(SedimentTerm, HorizontalAdvectionTerm):
-    """
+    """
     Advection term for sediment equation
 
     Same as :class:`HorizontalAdvectionTerm` but allows for equilibrium boundary condition
@@ -127,7 +126,7 @@ 
Source code for thetis.sediment_eq_2d
 
 
 class SedimentDiffusionTerm(SedimentTerm, HorizontalDiffusionTerm):
-    """
+    """
     Diffusion term for sediment equation
 
     Same as :class:`HorizontalDiffusionTerm` but allows for equilibrium boundary condition
@@ -136,7 +135,7 @@ 
Source code for thetis.sediment_eq_2d
 
 
 
[docs]class SedimentErosionTerm(SedimentTerm):
-    """
+    """
     Erosion term for sediment equation"""
 
[docs]    def residual(self, solution, solution_old, fields, fields_old, bnd_conditions):
         ero = self.sediment_model.get_erosion_term()
@@ -148,7 +147,7 @@ 
Source code for thetis.sediment_eq_2d
 
 
 
[docs]class SedimentDepositionTerm(SedimentTerm):
-    """
+    """
     Deposition term for sediment equation"""
 
[docs]    def residual(self, solution, solution_old, fields, fields_old, bnd_conditions):
         depo = self.sediment_model.get_deposition_coefficient()
@@ -159,12 +158,12 @@ 
Source code for thetis.sediment_eq_2d
 
 
 
[docs]class SedimentEquation2D(Equation):
-    """
+    """
     2D sediment advection-diffusion equation: :eq:`tracer_eq_2d` or :eq:`cons_tracer_eq_2d`
     with sediment source and sink term
     """
     def __init__(self, function_space, depth, options, sediment_model, conservative=False):
-        """
+        """
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :arg depth: :class:`DepthExpression` containing depth info
         :arg options: :class`ModelOptions2d` containing parameters
@@ -189,8 +188,8 @@ 
Source code for thetis.sediment_eq_2d
       

     

     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/sediment_model.html b/_modules/thetis/sediment_model.html
index 0d833cd..b06b6aa 100644
--- a/_modules/thetis/sediment_model.html
+++ b/_modules/thetis/sediment_model.html
@@ -5,10 +5,10 @@
   
     
     
-    thetis.sediment_model — Thetis 0+untagged.1718.gf836c9a.dirty documentation
-    
-    
-    
+    thetis.sediment_model — Thetis 0+untagged.2009.gd7af522 documentation
+    
+    
+    
     
     
     
@@ -37,7 +37,6 @@
         
  • Funding
    • 
         
  • Contact
    • 
         
  • GitHub
    • 
-        
  • Jenkins
    • 
       
     
    
   
    
@@ -58,7 +57,7 @@ 
    Source code for thetis.sediment_model
     
 
    [docs]class CorrectiveVelocityFactor:
     def __init__(self, depth, ksp, settling_velocity, ustar, a):
-        """
+        """
         Set up advective velocity factor `self.velocity_correction_factor`
         which accounts for mismatch between depth-averaged product of
         velocity with sediment and product of depth-averaged velocity
@@ -102,7 +101,7 @@ 
    Source code for thetis.sediment_model
             self.update()
 
 
    [docs]    def update(self):
-        """
+        """
         Update `self.velocity_correction_factor` using the updated values for velocity
         """
         # final correction factor
@@ -112,7 +111,7 @@ 
    Source code for thetis.sediment_model
     
    [docs]class SedimentModel(object):
     def __init__(self, options, mesh2d, uv, elev, depth):
 
-        """
+        """
         Set up a full morphological model simulation based on provided velocity and elevation functions.
 
         :arg options: Model options.
@@ -193,10 +192,10 @@ 
    Source code for thetis.sediment_model
             self.dstar = Function(self.P1_2d).interpolate(self.average_size*((self.g*self.R)/(self.viscosity**2))**(1/3))
         if float(max(self.dstar.dat.data[:])) < 1:
             raise ValueError('dstar value less than 1')
-        self.thetacr = Function(self.P1_2d).assign(conditional(self.dstar < 4, 0.24*(self.dstar**(-1)),
-                                                   conditional(self.dstar < 10, 0.14*(self.dstar**(-0.64)),
-                                                   conditional(self.dstar < 20, 0.04*(self.dstar**(-0.1)),
-                                                               conditional(self.dstar < 150, 0.013*(self.dstar**(0.29)), 0.055)))))
+        self.thetacr = Function(self.P1_2d).interpolate(conditional(self.dstar < 4, 0.24*(self.dstar**(-1)),
+                                                        conditional(self.dstar < 10, 0.14*(self.dstar**(-0.64)),
+                                                        conditional(self.dstar < 20, 0.04*(self.dstar**(-0.1)),
+                                                                    conditional(self.dstar < 150, 0.013*(self.dstar**(0.29)), 0.055)))))
 
         # critical bed shear stress
         self.taucr = Function(self.P1_2d).interpolate((self.rhos-self.rhow)*self.g*self.average_size*self.thetacr)
@@ -274,7 +273,7 @@ 
    Source code for thetis.sediment_model
                     self.stress = Function(self.P1DG_2d).interpolate(self.rhow*Constant(0.5)*self.qfc*self.unorm)
 
 
    [docs]    def get_bedload_term(self, bathymetry):
-        """
+        """
         Returns expression for bedload transport :math:`(qbx, qby)` to be used in the Exner equation.
         Note bathymetry is the function which is solved for in the exner equation.
 
@@ -408,7 +407,7 @@ 
    Source code for thetis.sediment_model
             return diff_tensor
    
 
 
    [docs]    def get_deposition_coefficient(self):
-        """Returns coefficient :math:`C` such that :math:`C/H*sediment` is deposition term in sediment equation
+        """Returns coefficient :math:`C` such that :math:`C/H*sediment` is deposition term in sediment equation
 
         If sediment field is depth-averaged, :math:`C*sediment` is (total) deposition (over the column)
         as it appears in the Exner equation, but deposition term in sediment equation needs
@@ -418,15 +417,15 @@ 
    Source code for thetis.sediment_model
             return self._deposition
    
 
 
    [docs]    def get_erosion_term(self):
-        """Returns expression for (depth-integrated) erosion."""
+        """Returns expression for (depth-integrated) erosion."""
         return self._erosion
    
 
 
    [docs]    def get_equilibrium_tracer(self):
-        """Returns expression for (depth-averaged) equilibrium tracer."""
+        """Returns expression for (depth-averaged) equilibrium tracer."""
         return self.equilibrium_tracer
    
 
 
    [docs]    def get_advective_velocity_correction_factor(self):
-        """Returns correction factor for the advective velocity in the sediment equations
+        """Returns correction factor for the advective velocity in the sediment equations
 
         With :attr:`.SedimentModelOptions.use_advective_velocity_correction`, this applies a correction to
         the supplied velocity solution `uv` to take into account the mismatch between
@@ -439,7 +438,7 @@ 
    Source code for thetis.sediment_model
                 return 1
    
 
 
    [docs]    def update(self):
-        """Update all functions used by :class:`.SedimentModel`
+        """Update all functions used by :class:`.SedimentModel`
 
         This repeats all projection and interpolations steps based on the current values
         of the `uv` and `elev` functions, provided in __init__."""
@@ -479,8 +478,8 @@ 
    Source code for thetis.sediment_model
           
    
     
    
     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/shallowwater_eq.html b/_modules/thetis/shallowwater_eq.html
index 56099ee..80afc52 100644
--- a/_modules/thetis/shallowwater_eq.html
+++ b/_modules/thetis/shallowwater_eq.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.shallowwater_eq — Thetis 0+untagged.1878.g52658ff.dirty documentation
+    thetis.shallowwater_eq — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -257,7 +257,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class ShallowWaterTerm(Term):
-    """
+    """
     Generic term in the shallow water equations that provides commonly used
     members and mapping for boundary functions.
     """
@@ -283,7 +283,7 @@ 
    Source code for thetis.shallowwater_eq
                          domain=self.function_space.ufl_domain())
 
 
    [docs]    def get_bnd_functions(self, eta_in, uv_in, bnd_id, bnd_conditions):
-        """
+        """
         Returns external values of elev and uv for all supported
         boundary conditions.
 
@@ -325,7 +325,7 @@ 
    Source code for thetis.shallowwater_eq
             return eta_ext, uv_ext
    
 
 
    [docs]    def impose_dynamic_bnd(self, bnd_funcs, bnd_id):
-        """
+        """
         Check if the prognostic variables have been specified on the boundary.
 
         If any prognostic value has been specified, dynamic boundary terms
@@ -350,7 +350,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class ShallowWaterMomentumTerm(ShallowWaterTerm):
-    """
+    """
     Generic term in the shallow water momentum equation that provides commonly used
     members and mapping for boundary functions.
     """
@@ -369,7 +369,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class ShallowWaterContinuityTerm(ShallowWaterTerm):
-    """
+    """
     Generic term in the depth-integrated continuity equation that provides commonly used
     members and mapping for boundary functions.
     """
@@ -386,7 +386,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class ExternalPressureGradientTerm(ShallowWaterMomentumTerm):
-    r"""
+    r"""
     External pressure gradient term, :math:`g \nabla \eta`
 
     The weak form reads
@@ -447,7 +447,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class HUDivTerm(ShallowWaterContinuityTerm):
-    r"""
+    r"""
     Divergence term, :math:`\nabla \cdot (H \bar{\textbf{u}})`
 
     The weak form reads
@@ -504,7 +504,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class HorizontalAdvectionTerm(ShallowWaterMomentumTerm):
-    r"""
+    r"""
     Advection of momentum term, :math:`\bar{\textbf{u}} \cdot \nabla\bar{\textbf{u}}`
 
     The weak form is
@@ -564,7 +564,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class HorizontalViscosityTerm(ShallowWaterMomentumTerm):
-    r"""
+    r"""
     Viscosity of momentum term
 
     If option :attr:`.ModelOptions.use_grad_div_viscosity_term` is ``True``, we
@@ -670,7 +670,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class CoriolisTerm(ShallowWaterMomentumTerm):
-    r"""
+    r"""
     Coriolis term, :math:`f\textbf{e}_z\wedge \bar{\textbf{u}}`
     """
 
    [docs]    def residual(self, uv, eta, uv_old, eta_old, fields, fields_old, bnd_conditions):
@@ -688,7 +688,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class WindStressTerm(ShallowWaterMomentumTerm):
-    r"""
+    r"""
     Wind stress term, :math:`-\tau_w/(H \rho_0)`
 
     Here :math:`\tau_w` is a user-defined wind stress :class:`Function`.
@@ -703,7 +703,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class AtmosphericPressureTerm(ShallowWaterMomentumTerm):
-    r"""
+    r"""
     Atmospheric pressure term, :math:`\nabla (p_a / \rho_0)`
 
     Here :math:`p_a` is a user-defined atmospheric pressure :class:`Function`.
@@ -717,7 +717,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class QuadraticDragTerm(ShallowWaterMomentumTerm):
-    r"""
+    r"""
     Quadratic Manning bottom friction term
     :math:`C_D \| \bar{\textbf{u}} \| \bar{\textbf{u}}`
 
@@ -755,7 +755,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 class BoundaryDragTerm(ShallowWaterMomentumTerm):
-    r"""
+    r"""
     Quadratic friction term on the boundary
     :math:`C_D \| \bar{\textbf{u}}_t \| \bar{\textbf{u}}_t`
 
@@ -779,7 +779,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class LinearDragTerm(ShallowWaterMomentumTerm):
-    r"""
+    r"""
     Linear friction term, :math:`C \bar{\textbf{u}}`
 
     Here :math:`C` is a user-defined drag coefficient.
@@ -794,7 +794,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class BottomDrag3DTerm(ShallowWaterMomentumTerm):
-    r"""
+    r"""
     Bottom drag term consistent with the 3D mode,
     :math:`C_D \| \textbf{u}_b \| \textbf{u}_b`
 
@@ -816,7 +816,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 class TurbineDragTerm(ShallowWaterMomentumTerm):
-    r"""
+    r"""
     Turbine drag parameterisation implemented through quadratic drag term
     :math:`c_t \| \bar{\textbf{u}} \| \bar{\textbf{u}}`
 
@@ -828,21 +828,24 @@ 
    Source code for thetis.shallowwater_eq
             c_t = (C_T A_T d)/2
 
     """
+    def __init__(self, u_test, u_space, eta_space,
+                 depth, options=None, tidal_farms=None):
+        super().__init__(u_test, u_space, eta_space, depth, options=options)
+        self.tidal_farms = tidal_farms
+
     def residual(self, uv, eta, uv_old, eta_old, fields, fields_old, bnd_conditions):
         total_h = self.depth.get_total_depth(eta_old)
         f = 0
-        for subdomain_id, farm_options in self.options.tidal_turbine_farms.items():
-            density = farm_options.turbine_density
-            C_T = farm_options.turbine_options.thrust_coefficient
-            A_T = pi * (farm_options.turbine_options.diameter/2.)**2
-            C_D = (C_T * A_T * density)/2.
+        for farm in self.tidal_farms:
+            density = farm.turbine_density
+            c_t = farm.friction_coefficient(uv_old, total_h)
             unorm = sqrt(dot(uv_old, uv_old))
-            f += C_D * unorm * inner(self.u_test, uv) / total_h * self.dx(subdomain_id)
+            f += c_t * density * unorm * inner(self.u_test, uv) / total_h * farm.dx
         return -f
 
 
 
    [docs]class MomentumSourceTerm(ShallowWaterMomentumTerm):
-    r"""
+    r"""
     Generic source term in the shallow water momentum equation
 
     The weak form reads
@@ -862,7 +865,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class ContinuitySourceTerm(ShallowWaterContinuityTerm):
-    r"""
+    r"""
     Generic source term in the depth-averaged continuity equation
 
     The weak form reads
@@ -882,7 +885,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 class BathymetryDisplacementMassTerm(ShallowWaterContinuityTerm):
-    r"""
+    r"""
     Bathmetry mass displacement term, :math:`\partial \eta / \partial t + \partial \tilde{h} / \partial t`
 
     The weak form reads
@@ -901,7 +904,7 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class BaseShallowWaterEquation(Equation):
-    """
+    """
     Abstract base class for ShallowWaterEquations, ShallowWaterMomentumEquation
     and FreeSurfaceEquation.
 
@@ -914,7 +917,7 @@ 
    Source code for thetis.shallowwater_eq
             self.depth = depth
         self.options = options
 
-
    [docs]    def add_momentum_terms(self, *args):
+
    [docs]    def add_momentum_terms(self, *args, tidal_farms=None):
         self.add_term(ExternalPressureGradientTerm(*args), 'implicit')
         self.add_term(HorizontalAdvectionTerm(*args), 'implicit')
         self.add_term(HorizontalViscosityTerm(*args), 'explicit')
@@ -925,8 +928,9 @@ 
    Source code for thetis.shallowwater_eq
             self.add_term(LinearDragTerm(*args), 'implicit')
         self.add_term(BoundaryDragTerm(*args), 'implicit')
         self.add_term(BottomDrag3DTerm(*args), 'source')
-        self.add_term(TurbineDragTerm(*args), 'implicit')
-        self.add_term(MomentumSourceTerm(*args), 'source')
    
+        self.add_term(MomentumSourceTerm(*args), 'source')
+        if tidal_farms:
+            self.add_term(TurbineDragTerm(*args, tidal_farms), 'implicit')
    
 
 
    [docs]    def add_continuity_terms(self, *args):
         self.add_term(HUDivTerm(*args), 'implicit')
@@ -940,14 +944,14 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class ShallowWaterEquations(BaseShallowWaterEquation):
-    """
+    """
     2D depth-averaged shallow water equations in non-conservative form.
 
     This defines the full 2D SWE equations :eq:`swe_freesurf` -
     :eq:`swe_momentum`.
     """
-    def __init__(self, function_space, depth, options):
-        """
+    def __init__(self, function_space, depth, options, tidal_farms=None):
+        """
         :arg function_space: Mixed function space where the solution belongs
         :arg depth: :class: `DepthExpression` containing depth info
         :arg options: :class:`.AttrDict` object containing all circulation model options
@@ -955,9 +959,9 @@ 
    Source code for thetis.shallowwater_eq
             super(ShallowWaterEquations, self).__init__(function_space, depth, options)
 
         u_test, eta_test = TestFunctions(function_space)
-        u_space, eta_space = function_space.split()
+        u_space, eta_space = function_space.subfunctions
 
-        self.add_momentum_terms(u_test, u_space, eta_space, depth, options)
+        self.add_momentum_terms(u_test, u_space, eta_space, depth, options, tidal_farms=tidal_farms)
 
         self.add_continuity_terms(eta_test, eta_space, u_space, depth, options)
         self.bathymetry_displacement_mass_term = BathymetryDisplacementMassTerm(
@@ -978,14 +982,14 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class ModeSplit2DEquations(BaseShallowWaterEquation):
-    r"""
+    r"""
     2D depth-averaged shallow water equations for mode splitting schemes.
 
     Defines the equations :eq:`swe_freesurf_modesplit` -
     :eq:`swe_momentum_modesplit`.
     """
     def __init__(self, function_space, depth, options):
-        """
+        """
         :arg function_space: Mixed function space where the solution belongs
         :arg depth: :class: `DepthExpression` containing depth info
         :arg options: :class:`.AttrDict` object containing all circulation model options
@@ -994,7 +998,7 @@ 
    Source code for thetis.shallowwater_eq
             super(ModeSplit2DEquations, self).__init__(function_space, depth, options)
 
         u_test, eta_test = TestFunctions(function_space)
-        u_space, eta_space = function_space.split()
+        u_space, eta_space = function_space.subfunctions
 
         self.add_momentum_terms(u_test, u_space, eta_space, depth, options)
 
@@ -1016,11 +1020,11 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class FreeSurfaceEquation(BaseShallowWaterEquation):
-    """
+    """
     2D free surface equation :eq:`swe_freesurf` in non-conservative form.
     """
     def __init__(self, eta_test, eta_space, u_space, depth, options):
-        """
+        """
         :arg eta_test: test function of the elevation function space
         :arg eta_space: elevation function space
         :arg u_space: velocity function space
@@ -1047,12 +1051,12 @@ 
    Source code for thetis.shallowwater_eq
     
 
 
    [docs]class ShallowWaterMomentumEquation(BaseShallowWaterEquation):
-    """
+    """
     2D depth averaged momentum equation :eq:`swe_momentum` in non-conservative
     form.
     """
-    def __init__(self, u_test, u_space, eta_space, depth, options):
-        """
+    def __init__(self, u_test, u_space, eta_space, depth, options, tidal_farms=None):
+        """
         :arg u_test: test function of the velocity function space
         :arg u_space: velocity function space
         :arg eta_space: elevation function space
@@ -1060,7 +1064,7 @@ 
    Source code for thetis.shallowwater_eq
             :arg options: :class:`.AttrDict` object containing all circulation model options
         """
         super(ShallowWaterMomentumEquation, self).__init__(u_space, depth, options)
-        self.add_momentum_terms(u_test, u_space, eta_space, depth, options)
+        self.add_momentum_terms(u_test, u_space, eta_space, depth, options, tidal_farms=tidal_farms)
 
 
    [docs]    def residual(self, label, solution, solution_old, fields, fields_old, bnd_conditions):
         uv = solution
@@ -1078,8 +1082,8 @@ 
    Source code for thetis.shallowwater_eq
           
    
     
    
     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/solver.html b/_modules/thetis/solver.html
index 3654290..4fb1bae 100644
--- a/_modules/thetis/solver.html
+++ b/_modules/thetis/solver.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.solver — Thetis 0+untagged.1888.gc4e776a documentation
+    thetis.solver — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -72,10 +72,11 @@ 
    Source code for thetis.solver
     from .log import *
 from collections import OrderedDict
 import numpy
+import datetime
 
 
 
    [docs]class FlowSolver(FrozenClass):
-    """
+    """
     Main object for 3D solver
 
     **Example**
@@ -132,8 +133,8 @@ 
    Source code for thetis.solver
         @unfrozen
     @PETSc.Log.EventDecorator("thetis.FlowSolver.__init__")
     def __init__(self, mesh2d, bathymetry_2d, n_layers,
-                 options=None, extrude_options=None):
-        """
+                 options=None, extrude_options=None, keep_log=False):
+        """
         :arg mesh2d: :class:`Mesh` object of the 2D mesh
         :arg bathymetry_2d: Bathymetry of the domain. Bathymetry stands for
             the mean water depth (positive downwards).
@@ -143,17 +144,18 @@ 
    Source code for thetis.solver
             :kwarg options: Model options (optional). Model options can also be
             changed directly via the :attr:`.options` class property.
         :type options: :class:`.ModelOptions3d` instance
+        :kwarg bool keep_log: append to an existing log file, or overwrite it?
         """
         self._initialized = False
 
         self.bathymetry_cg_2d = bathymetry_2d
 
         self.mesh2d = mesh2d
-        """2D :class`Mesh`"""
+        """2D :class`Mesh`"""
         if extrude_options is None:
             extrude_options = {}
         self.mesh = extrude_mesh_sigma(mesh2d, n_layers, bathymetry_2d, **extrude_options)
-        """3D :class`Mesh`"""
+        """3D :class`Mesh`"""
         self.comm = mesh2d.comm
 
         # add boundary length info
@@ -162,15 +164,15 @@ 
    Source code for thetis.solver
             self.mesh.boundary_len = bnd_len
 
         self.dt = None
-        """Time step"""
+        """Time step"""
         self.dt_2d = None
-        """Time of the 2D solver"""
+        """Time of the 2D solver"""
         self.M_modesplit = None
-        """Mode split ratio (int)"""
+        """Mode split ratio (int)"""
 
         # override default options
         self.options = ModelOptions3d()
-        """
+        """
         Dictionary of all options. A :class:`.ModelOptions3d` object.
         """
         if options is not None:
@@ -189,30 +191,31 @@ 
    Source code for thetis.solver
                                   }
 
         self.callbacks = callback.CallbackManager()
-        """
+        """
         :class:`.CallbackManager` object that stores all callbacks
         """
 
         self.fields = FieldDict()
-        """
+        """
         :class:`.FieldDict` that holds all functions needed by the solver
         object
         """
 
         self.function_spaces = AttrDict()
-        """
+        """
         :class:`.AttrDict` that holds all function spaces needed by the
         solver object
         """
 
         self.export_initial_state = True
-        """Do export initial state. False if continuing a simulation"""
+        """Do export initial state. False if continuing a simulation"""
 
         self._simulation_continued = False
+        self.keep_log = keep_log
         self._field_preproc_funcs = {}
 
 
    [docs]    def compute_dx_factor(self):
-        """
+        """
         Computes normalized distance between nodes in the horizontal direction
 
         The factor depends on the finite element space and its polynomial
@@ -228,7 +231,7 @@ 
    Source code for thetis.solver
             return factor
    
 
 
    [docs]    def compute_dz_factor(self):
-        """
+        """
         Computes a normalized distance between nodes in the vertical direction
 
         The factor depends on the finite element space and its polynomial
@@ -242,7 +245,7 @@ 
    Source code for thetis.solver
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver.compute_dt_2d")
     def compute_dt_2d(self, u_scale):
-        r"""
+        r"""
         Computes maximum explicit time step from CFL condition.
 
         .. math :: \Delta t = \frac{\Delta x}{U}
@@ -280,7 +283,7 @@ 
    Source code for thetis.solver
             return dt
    
 
 
    [docs]    def compute_dt_h_advection(self, u_scale):
-        r"""
+        r"""
         Computes maximum explicit time step for horizontal advection
 
         .. math :: \Delta t = \frac{\Delta x}{U_{scale}}
@@ -302,7 +305,7 @@ 
    Source code for thetis.solver
             return dt
    
 
 
    [docs]    def compute_dt_v_advection(self, w_scale):
-        r"""
+        r"""
         Computes maximum explicit time step for vertical advection
 
         .. math :: \Delta t = \frac{\Delta z}{W_{scale}}
@@ -324,7 +327,7 @@ 
    Source code for thetis.solver
             return dt
    
 
 
    [docs]    def compute_dt_diffusion(self, nu_scale):
-        r"""
+        r"""
         Computes maximum explicit time step for horizontal diffusion.
 
         .. math :: \Delta t = \alpha \frac{(\Delta x)^2}{\nu_{scale}}
@@ -346,7 +349,7 @@ 
    Source code for thetis.solver
             return dt
    
 
 
    [docs]    def compute_mesh_stats(self):
-        """
+        """
         Computes number of elements, nodes etc and prints to sdtout
         """
         nnodes = self.function_spaces.P1_2d.dim()
@@ -381,7 +384,7 @@ 
    Source code for thetis.solver
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver.set_time_step")
     def set_time_step(self):
-        """
+        """
         Sets the model the model time step
 
         If the time integrator supports automatic time step, and
@@ -457,7 +460,7 @@ 
    Source code for thetis.solver
     
    [docs]    @unfrozen
     @PETSc.Log.EventDecorator("thetis.FlowSolver.create_function_spaces")
     def create_function_spaces(self):
-        """
+        """
         Creates function spaces
 
         Function spaces are accessible via :attr:`.function_spaces`
@@ -520,14 +523,15 @@ 
    Source code for thetis.solver
     
    [docs]    @unfrozen
     @PETSc.Log.EventDecorator("thetis.FlowSolver.create_fields")
     def create_fields(self):
-        """
+        """
         Creates all fields
         """
         if not hasattr(self, 'U_2d'):
             self.create_function_spaces()
 
         if self.options.log_output and not self.options.no_exports:
-            set_log_directory(self.options.output_directory)
+            mode = "a" if self.keep_log else "w"
+            set_log_directory(self.options.output_directory, mode=mode)
 
         # mesh velocity etc fields must be in the same space as 3D coordinates
         coord_is_dg = element_continuity(self.mesh2d.coordinates.function_space().ufl_element()).horizontal == 'dg'
@@ -541,7 +545,7 @@ 
    Source code for thetis.solver
             # ----- fields
         self.fields.solution_2d = Function(self.function_spaces.V_2d)
         # correct treatment of the split 2d functions
-        uv_2d, eta2d = self.fields.solution_2d.split()
+        uv_2d, eta2d = self.fields.solution_2d.subfunctions
         self.fields.uv_2d = uv_2d
         self.fields.elev_2d = eta2d
         # elevation field seen by the 3D mode, different from elev_2d
@@ -663,7 +667,7 @@ 
    Source code for thetis.solver
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver.add_new_field")
     def add_new_field(self, function, label, name, filename, shortname=None, unit='-', preproc_func=None):
-        """
+        """
         Add a field to :attr:`fields`.
 
         :arg function: representation of the field as a :class:`Function`
@@ -680,7 +684,7 @@ 
    Source code for thetis.solver
             assert isinstance(filename, str)
         assert shortname is None or isinstance(shortname, str)
         assert isinstance(unit, str)
-        assert preproc_func is None or callable(preproc_func)
+        assert preproc_func is None or callable(preproc_func)
         assert label not in field_metadata, f"Field '{label}' already exists."
         assert ' ' not in label, "Labels cannot contain spaces"
         assert ' ' not in filename, "Filenames cannot contain spaces"
@@ -697,7 +701,7 @@ 
    Source code for thetis.solver
     
    [docs]    @unfrozen
     @PETSc.Log.EventDecorator("thetis.FlowSolver.create_equations")
     def create_equations(self):
-        """
+        """
         Creates equation instances
         """
         if 'uv_3d' not in self.fields:
@@ -705,7 +709,8 @@ 
    Source code for thetis.solver
     
         if self.options.log_output and not self.options.no_exports:
             logfile = os.path.join(create_directory(self.options.output_directory), 'log')
-            filehandler = logging.logging.FileHandler(logfile, mode='w')
+            mode = "a" if self.keep_log else "w"
+            filehandler = logging.logging.FileHandler(logfile, mode=mode)
             filehandler.setFormatter(logging.logging.Formatter('%(message)s'))
             output_logger.addHandler(filehandler)
 
@@ -885,7 +890,7 @@ 
    Source code for thetis.solver
     
    [docs]    @unfrozen
     @PETSc.Log.EventDecorator("thetis.FlowSolver.create_timestepper")
     def create_timestepper(self):
-        """
+        """
         Creates time stepper instance
         """
         if not hasattr(self, 'equations'):
@@ -909,14 +914,14 @@ 
    Source code for thetis.solver
             # compute maximal diffusivity for explicit schemes
         degree_h, degree_v = self.function_spaces.H.ufl_element().degree()
         max_diff_alpha = 1.0/60.0/max((degree_h*(degree_h + 1)), 1.0)  # FIXME depends on element type and order
-        self.fields.max_h_diff.assign(max_diff_alpha/self.dt * self.fields.h_elem_size_3d**2)
+        self.fields.max_h_diff.interpolate(max_diff_alpha/self.dt * self.fields.h_elem_size_3d**2)
         d = self.fields.max_h_diff.dat.data
         print_output('max h diff {:} - {:}'.format(d.min(), d.max()))
    
 
 
    [docs]    @unfrozen
     @PETSc.Log.EventDecorator("thetis.FlowSolver.create_exporters")
     def create_exporters(self):
-        """
+        """
         Creates file exporters
         """
         if not hasattr(self, 'timestepper'):
@@ -945,7 +950,7 @@ 
    Source code for thetis.solver
                 self.exporters['hdf5'] = e
    
 
 
    [docs]    def initialize(self):
-        """
+        """
         Creates function spaces, equations, time stepper and exporters
         """
         if not hasattr(self.function_spaces, 'U'):
@@ -961,7 +966,7 @@ 
    Source code for thetis.solver
     
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver.assign_initial_conditions")
     def assign_initial_conditions(self, elev=None, salt=None, temp=None,
                                   uv_2d=None, uv_3d=None, tke=None, psi=None):
-        """
+        """
         Assigns initial conditions
 
         :kwarg elev: Initial condition for water elevation
@@ -1017,7 +1022,7 @@ 
    Source code for thetis.solver
                 self.turbulence_model.initialize()
    
 
 
    [docs]    def add_callback(self, callback, eval_interval='export'):
-        """
+        """
         Adds callback to solver object
 
         :arg callback: :class:`.DiagnosticCallback` instance
@@ -1029,7 +1034,7 @@ 
    Source code for thetis.solver
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver.export")
     def export(self):
-        """
+        """
         Export all fields to disk
 
         Also evaluates all callbacks set to 'export' interval.
@@ -1044,26 +1049,50 @@ 
    Source code for thetis.solver
             self.fields.uv_3d -= self.fields.uv_dav_3d
    
 
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver.load_state")
-    def load_state(self, i_export, outputdir=None, t=None, iteration=None):
-        """
+    def load_state(self, i_stored, outputdir=None, t=None, iteration=None,
+                   i_export=None, legacy_mode=False):
+        """
         Loads simulation state from hdf5 outputs.
 
-        This replaces :meth:`.assign_initial_conditions` in model initilization.
+        Replaces :meth:`.assign_initial_conditions` in model initialization.
 
-        This assumes that model setup is kept the same (e.g. time step) and
-        all pronostic state variables are exported in hdf5 format.  The required
-        state variables are: elev_2d, uv_2d, uv_3d, salt_3d, temp_3d, tke_3d,
-        psi_3d
+        For example, continue simulation from export 40,
 
-        Currently hdf5 field import only works for the same number of MPI
-        processes.
+        .. code-block:: python
 
-        :arg int i_export: export index to load
+            solver_obj.load_state(40)
+
+        Continue simulation from another output directory,
+
+        .. code-block:: python
+
+            solver_obj.load_state(40, outputdir='outputs_spinup')
+
+        Continue simulation from another output directory and reset the
+        counters,
+
+        .. code-block:: python
+
+            solver_obj.load_state(40, outputdir='outputs_spinup',
+                iteration=0, t=0, i_export=0)
+
+        Model state can be loaded if all prognostic state variables have been
+        stored in hdf5 format. Required state variables are: `elev_2d`,
+        `uv_2d`, `uv_3d`, `salt_3d`, `temp_3d`, `tke_3d`, and `psi_3d`.
+
+        Simulation time and iteration can be recovered automatically provided
+        that the model setup is the same (e.g. time step and export_time).
+
+        :arg int i_stored: export index to load
         :kwarg string outputdir: (optional) directory where files are read from.
             By default ``options.output_directory``.
-        :kwarg float t: simulation time. Overrides the time stamp stored in the
-            hdf5 files.
-        :kwarg int iteration: Overrides the iteration count in the hdf5 files.
+        :kwarg float t: simulation time. Overrides the time stamp inferred from
+            model options.
+        :kwarg int iteration: Overrides the iteration count inferred from model
+            options.
+        :kwarg int i_export: Set initial export index for the present run. By
+            default, `i_stored` is used.
+        :kwarg bool legacy_mode: Load legacy `DumbCheckpoint` files.
         """
         if not self._initialized:
             self.initialize()
@@ -1079,26 +1108,27 @@ 
    Source code for thetis.solver
                                        self.fields,
                                    field_metadata,
                                    export_type='hdf5',
+                                   legacy_mode=legacy_mode,
                                    verbose=self.options.verbose > 0)
-        e.exporters['uv_2d'].load(i_export, self.fields.uv_2d)
-        e.exporters['elev_2d'].load(i_export, self.fields.elev_2d)
-        e.exporters['uv_3d'].load(i_export, self.fields.uv_3d)
+        e.exporters['uv_2d'].load(i_stored, self.fields.uv_2d)
+        e.exporters['elev_2d'].load(i_stored, self.fields.elev_2d)
+        e.exporters['uv_3d'].load(i_stored, self.fields.uv_3d)
         # NOTE remove mean from uv_3d
         self.timestepper._remove_depth_average_from_uv_3d()
         salt = temp = tke = psi = None
         if self.options.solve_salinity:
             salt = self.fields.salt_3d
-            e.exporters['salt_3d'].load(i_export, salt)
+            e.exporters['salt_3d'].load(i_stored, salt)
         if self.options.solve_temperature:
             temp = self.fields.temp_3d
-            e.exporters['temp_3d'].load(i_export, temp)
+            e.exporters['temp_3d'].load(i_stored, temp)
         if self.options.use_turbulence:
             if 'tke_3d' in self.fields:
                 tke = self.fields.tke_3d
-                e.exporters['tke_3d'].load(i_export, tke)
+                e.exporters['tke_3d'].load(i_stored, tke)
             if 'psi_3d' in self.fields:
                 psi = self.fields.psi_3d
-                e.exporters['psi_3d'].load(i_export, psi)
+                e.exporters['psi_3d'].load(i_stored, psi)
         self.assign_initial_conditions(elev=self.fields.elev_2d,
                                        uv_2d=self.fields.uv_2d,
                                        uv_3d=self.fields.uv_3d,
@@ -1107,6 +1137,8 @@ 
    Source code for thetis.solver
                                            )
 
         # time stepper bookkeeping for export time step
+        if i_export is None:
+            i_export = i_stored
         self.i_export = i_export
         self.next_export_t = self.i_export*self.options.simulation_export_time
         if iteration is None:
@@ -1126,24 +1158,53 @@ 
    Source code for thetis.solver
             for e in self.exporters.values():
             e.set_next_export_ix(self.i_export + offset)
    
 
-
    [docs]    def print_state(self, cputime):
-        """
+
    [docs]    def print_state(self, cputime, print_header=False):
+        """
         Print a summary of the model state on stdout
 
-        :arg float cputime: Measured CPU time
+        :arg float cputime: Measured CPU time in seconds
+        :kwarg print_header: Whether to print column header first
         """
+        entries = [
+            ('exp', self.i_export, '5d'),
+            ('iter', self.iteration, '5d'),
+        ]
+        # generate simulation time string
+        if self.options.simulation_initial_date is not None:
+            now = self.options.simulation_initial_date + datetime.timedelta(seconds=self.simulation_time)
+            date_str = f'{now:%Y-%m-%d}'.rjust(11)
+            time_str = f'{now:%H:%M:%S}'.rjust(9)
+            entries += [
+                ('date', date_str, '11s'),
+                ('time', time_str, '9s'),
+            ]
+        else:
+            time_str = f'{self.simulation_time:.2f}'.rjust(15)
+            entries += [
+                ('time', time_str, '15s'),
+            ]
+
         norm_h = norm(self.fields.elev_2d)
         norm_u = norm(self.fields.uv_3d)
+        # list of entries to print: (header, value, format)
+        entries += [
+            ('eta norm', norm_h, '14.4f'),
+            ('u norm', norm_u, '14.4f'),
+            ('Tcpu', cputime, '6.2f'),
+        ]
+
+        if print_header:
+            # generate header
+            header = ' '.join([e[0].rjust(len(f'{e[1]:{e[2]}}')) for e in entries])
+            print_output(header)
 
-        line = ('{iexp:5d} {i:5d} T={t:10.2f} '
-                'eta norm: {e:10.4f} u norm: {u:10.4f} {cpu:5.2f}')
-        print_output(line.format(iexp=self.i_export, i=self.iteration,
-                                 t=self.simulation_time, e=norm_h,
-                                 u=norm_u, cpu=cputime))
+        # generate line
+        line = ' '.join([f'{e[1]:{e[2]}}' for e in entries])
+        print_output(line)
         sys.stdout.flush()
    
 
     def _print_field(self, field):
-        """
+        """
         Prints min/max values of a field for debugging.
 
         :arg field: a :class:`Function` or a field string, e.g. 'salt_3d'
@@ -1159,7 +1220,7 @@ 
    Source code for thetis.solver
             print_output('    {:}: {:.4f} {:.4f}'.format(_field.name(), minval, maxval))
 
 
    [docs]    def print_state_debug(self):
-        """
+        """
         Print min/max values of prognostic/diagnostic fields for debugging.
         """
 
@@ -1183,7 +1244,7 @@ 
    Source code for thetis.solver
     
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver.iterate")
     def iterate(self, update_forcings=None, update_forcings3d=None,
                 export_func=None):
-        """
+        """
         Runs the simulation
 
         Iterates over the time loop until time ``options.simulation_end_time`` is reached.
@@ -1255,8 +1316,27 @@ 
    Source code for thetis.solver
                     for k in self.callbacks[m]:
                     self.callbacks[m][k].set_write_mode('append')
 
+        initial_simulation_time = self.simulation_time
+        internal_iteration = 0
+
+        init_date = self.options.simulation_initial_date
+        end_date = self.options.simulation_end_date
+        if (init_date is not None and end_date is not None):
+            now = init_date + datetime.timedelta(initial_simulation_time)
+            assert end_date > now, f'Simulation end date must be greater than initial time {now}'
+            print_output(
+                f'Running simulation\n'
+                f' from {now:%Y-%m-%d %H:%M:%S %Z}\n'
+                f' to   {end_date:%Y-%m-%d %H:%M:%S %Z}'
+            )
+            end_time = (end_date - now).total_seconds() + initial_simulation_time
+            if self.options.simulation_end_time is not None:
+                warning('Both simulation_end_date and simulation_end_time have been set, ignoring simulation_end_time.')
+            self.options.simulation_end_time = end_time
+        assert self.options.simulation_end_time is not None, 'simulation_end_time must be set'
+
         # initial export
-        self.print_state(0.0)
+        self.print_state(0.0, print_header=True)
         if self.export_initial_state:
             self.export()
             if export_func is not None:
@@ -1264,9 +1344,6 @@ 
    Source code for thetis.solver
                 if 'vtk' in self.exporters:
                 self.exporters['vtk'].export_bathymetry(self.fields.bathymetry_2d)
 
-        initial_simulation_time = self.simulation_time
-        internal_iteration = 0
-
         while self.simulation_time <= self.options.simulation_end_time - t_epsilon:
 
             self.timestepper.advance(self.simulation_time,
@@ -1301,8 +1378,8 @@ 
    Source code for thetis.solver
           
    
     
    
     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/solver2d.html b/_modules/thetis/solver2d.html
index be336a0..203f118 100644
--- a/_modules/thetis/solver2d.html
+++ b/_modules/thetis/solver2d.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.solver2d — Thetis 0+untagged.1888.gc4e776a documentation
+    thetis.solver2d — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -61,13 +61,13 @@ 
    Source code for thetis.solver2d
     from . import implicitexplicit
 from . import coupled_timeintegrator_2d
 from . import tracer_eq_2d
-from . import conservative_tracer_eq_2d
 from . import sediment_eq_2d
 from . import exner_eq
 import weakref
 import time as time_mod
 from mpi4py import MPI
 from . import exporter
+from .turbines import TidalTurbineFarm, DiscreteTidalTurbineFarm
 from .field_defs import field_metadata
 from .options import ModelOptions2d
 from . import callback
@@ -75,10 +75,11 @@ 
    Source code for thetis.solver2d
     from collections import OrderedDict
 import thetis.limiter as limiter
 import numpy
+import datetime
 
 
 
    [docs]class FlowSolver2d(FrozenClass):
-    """
+    """
     Main object for 2D depth averaged solver
 
     **Example**
@@ -130,8 +131,8 @@ 
    Source code for thetis.solver2d
         """
     @unfrozen
     @PETSc.Log.EventDecorator("thetis.FlowSolver2d.__init__")
-    def __init__(self, mesh2d, bathymetry_2d, options=None):
-        """
+    def __init__(self, mesh2d, bathymetry_2d, options=None, keep_log=False):
+        """
         :arg mesh2d: :class:`Mesh` object of the 2D mesh
         :arg bathymetry_2d: Bathymetry of the domain. Bathymetry stands for
             the mean water depth (positive downwards).
@@ -139,6 +140,7 @@ 
    Source code for thetis.solver2d
             :kwarg options: Model options (optional). Model options can also be
             changed directly via the :attr:`.options` class property.
         :type options: :class:`.ModelOptions2d` instance
+        :kwarg bool keep_log: append to an existing log file, or overwrite it?
         """
         self._initialized = False
         self.mesh2d = mesh2d
@@ -149,10 +151,10 @@ 
    Source code for thetis.solver2d
             self.mesh2d.boundary_len = bnd_len
 
         self.dt = None
-        """Time step"""
+        """Time step"""
 
         self.options = ModelOptions2d()
-        """
+        """
         Dictionary of all options. A :class:`.ModelOptions2d` object.
         """
         if options is not None:
@@ -165,18 +167,18 @@ 
    Source code for thetis.solver2d
             self.next_export_t = self.simulation_time + self.options.simulation_export_time
 
         self.callbacks = callback.CallbackManager()
-        """
+        """
         :class:`.CallbackManager` object that stores all callbacks
         """
 
         self.fields = FieldDict()
-        """
+        """
         :class:`.FieldDict` that holds all functions needed by the solver
         object
         """
 
         self.function_spaces = AttrDict()
-        """
+        """
         :class:`.AttrDict` that holds all function spaces needed by the
         solver object
         """
@@ -184,21 +186,22 @@ 
    Source code for thetis.solver2d
             self.fields.bathymetry_2d = bathymetry_2d
 
         self.export_initial_state = True
-        """Do export initial state. False if continuing a simulation"""
+        """Do export initial state. False if continuing a simulation"""
 
         self.sediment_model = None
-        """set up option for sediment model"""
+        """set up option for sediment model"""
 
         self.bnd_functions = {'shallow_water': {}, 'tracer': {}, 'sediment': {}}
 
         if 'tracer_2d' in field_metadata:
             field_metadata.pop('tracer_2d')
         self.solve_tracer = False
+        self.keep_log = keep_log
         self._field_preproc_funcs = {}
 
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.compute_time_step")
     def compute_time_step(self, u_scale=Constant(0.0)):
-        r"""
+        r"""
         Computes maximum explicit time step from CFL condition.
 
         .. math :: \Delta t = \frac{\Delta x}{U}
@@ -228,7 +231,7 @@ 
    Source code for thetis.solver2d
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.compute_mesh_stats")
     def compute_mesh_stats(self):
-        """
+        """
         Computes number of elements, nodes etc and prints to sdtout
         """
         nnodes = self.function_spaces.P1_2d.dim()
@@ -262,7 +265,7 @@ 
    Source code for thetis.solver2d
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.set_time_step")
     def set_time_step(self, alpha=0.05):
-        """
+        """
         Sets the model the model time step
 
         If the time integrator supports automatic time step, and
@@ -299,7 +302,7 @@ 
    Source code for thetis.solver2d
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.set_wetting_and_drying_alpha")
     def set_wetting_and_drying_alpha(self):
-        r"""
+        r"""
         Compute a wetting and drying parameter :math:`\alpha` which ensures positive water
         depth using the approximate method suggested by Karna et al. (2011).
 
@@ -355,7 +358,7 @@ 
    Source code for thetis.solver2d
     
    [docs]    @unfrozen
     @PETSc.Log.EventDecorator("thetis.FlowSolver2d.create_function_spaces")
     def create_function_spaces(self):
-        """
+        """
         Creates function spaces
 
         Function spaces are accessible via :attr:`.function_spaces`
@@ -403,7 +406,7 @@ 
    Source code for thetis.solver2d
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.add_new_field")
     def add_new_field(self, function, label, name, filename, shortname=None, unit='-', preproc_func=None):
-        """
+        """
         Add a field to :attr:`fields`.
 
         :arg function: representation of the field as a :class:`Function`
@@ -420,7 +423,7 @@ 
    Source code for thetis.solver2d
             assert isinstance(filename, str)
         assert shortname is None or isinstance(shortname, str)
         assert isinstance(unit, str)
-        assert preproc_func is None or callable(preproc_func)
+        assert preproc_func is None or callable(preproc_func)
         assert label not in field_metadata, f"Field '{label}' already exists."
         assert ' ' not in label, "Labels cannot contain spaces"
         assert ' ' not in filename, "Filenames cannot contain spaces"
@@ -435,82 +438,125 @@ 
    Source code for thetis.solver2d
                 self._field_preproc_funcs[label] = preproc_func
    
 
 
    [docs]    @unfrozen
-    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.create_equations")
+    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.create_fields")
     def create_fields(self):
-        """
+        """
         Creates field Functions
         """
         if not hasattr(self.function_spaces, 'U_2d'):
             self.create_function_spaces()
 
         if self.options.log_output and not self.options.no_exports:
-            set_log_directory(self.options.output_directory)
+            mode = "a" if self.keep_log else "w"
+            set_log_directory(self.options.output_directory, mode=mode)
 
-        self.fields.solution_2d = Function(self.function_spaces.V_2d, name='solution_2d')
-        # correct treatment of the split 2d functions
-        uv_2d, elev_2d = self.fields.solution_2d.split()
-        self.fields.uv_2d = uv_2d
-        self.fields.elev_2d = elev_2d
+        # Add general fields
         self.fields.h_elem_size_2d = Function(self.function_spaces.P1_2d)
         get_horizontal_elem_size_2d(self.fields.h_elem_size_2d)
         self.set_wetting_and_drying_alpha()
         self.depth = DepthExpression(self.fields.bathymetry_2d,
                                      use_nonlinear_equations=self.options.use_nonlinear_equations,
                                      use_wetting_and_drying=self.options.use_wetting_and_drying,
-                                     wetting_and_drying_alpha=self.options.wetting_and_drying_alpha)
    
+                                     wetting_and_drying_alpha=self.options.wetting_and_drying_alpha)
+
+        # Add fields for shallow water modelling
+        self.fields.solution_2d = Function(self.function_spaces.V_2d, name='solution_2d')
+        uv_2d, elev_2d = self.fields.solution_2d.subfunctions  # correct treatment of the split 2d functions
+        self.fields.uv_2d = uv_2d
+        self.fields.elev_2d = elev_2d
+
+        # Add tracer fields
+        self.solve_tracer = len(self.options.tracer_fields.keys()) > 0
+        for system, parent in self.options.tracer_fields.copy().items():
+            labels = system.split(',')
+            num_labels = len(labels)
+            if parent is None:
+                Q_2d = self.function_spaces.Q_2d
+                fs = Q_2d if num_labels == 1 else MixedFunctionSpace([Q_2d]*len(labels))
+                parent = Function(fs)
+            if num_labels > 1:
+                self.fields[system] = parent
+            self.options.tracer_fields[system] = parent
+            children = [parent] if num_labels == 1 else parent.subfunctions
+            for label, function in zip(labels, children):
+                tracer = self.options.tracer[label]
+                function.rename(label)
+                self.add_new_field(function,
+                                   label,
+                                   tracer.metadata['name'],
+                                   tracer.metadata['filename'],
+                                   shortname=tracer.metadata['shortname'],
+                                   unit=tracer.metadata['unit'])
+
+        # Add suspended sediment field
+        if self.options.sediment_model_options.solve_suspended_sediment:
+            self.fields.sediment_2d = Function(self.function_spaces.Q_2d, name='sediment_2d')
+
+        # Add fields for non-hydrostatic mode
+        if self.options.nh_model_options.solve_nonhydrostatic_pressure:
+            q_degree = self.options.polynomial_degree
+            if self.options.nh_model_options.q_degree is not None:
+                q_degree = self.options.nh_model_options.q_degree
+            fs_q = get_functionspace(self.mesh2d, 'CG', q_degree)
+            self.fields.q_2d = Function(fs_q, name='q_2d')  # 2D non-hydrostatic pressure at bottom
+            self.fields.w_2d = Function(self.function_spaces.H_2d, name='w_2d')  # depth-averaged vertical velocity
    
 
 
    [docs]    @unfrozen
     @PETSc.Log.EventDecorator("thetis.FlowSolver2d.create_equations")
     def create_equations(self):
-        """
+        """
         Creates equation instances
         """
         if not hasattr(self.fields, 'uv_2d'):
             self.create_fields()
-
         self.equations = AttrDict()
+
+        # tidal farms, if any
+        if len(self.options.tidal_turbine_farms) + len(self.options.discrete_tidal_turbine_farms) > 0:
+            self.tidal_farms = []
+            p = self.function_spaces.U_2d.ufl_element().degree()
+            quad_degree = 2*p + 1
+            for subdomain, farm_options in self.options.tidal_turbine_farms.items():
+                fdx = dx(subdomain, degree=quad_degree)
+                self.tidal_farms.append(TidalTurbineFarm(farm_options.turbine_density,
+                                                         fdx, farm_options))
+            for subdomain, farm_options in self.options.discrete_tidal_turbine_farms.items():
+                fdx = dx(subdomain, degree=farm_options.quadrature_degree)
+                self.tidal_farms.append(DiscreteTidalTurbineFarm(self.mesh2d, fdx, farm_options))
+        else:
+            self.tidal_farms = None
+        # Shallow water equations for hydrodynamic modelling
         self.equations.sw = shallowwater_eq.ShallowWaterEquations(
             self.fields.solution_2d.function_space(),
             self.depth,
             self.options,
+            tidal_farms=self.tidal_farms
         )
         self.equations.sw.bnd_functions = self.bnd_functions['shallow_water']
-        uv_2d, elev_2d = self.fields.solution_2d.split()
-        for label, tracer in self.options.tracer.items():
-            self.add_new_field(tracer.function or Function(self.function_spaces.Q_2d, name=label),
-                               label,
-                               tracer.metadata['name'],
-                               tracer.metadata['filename'],
-                               shortname=tracer.metadata['shortname'],
-                               unit=tracer.metadata['unit'])
-            if tracer.use_conservative_form:
-                self.equations[label] = conservative_tracer_eq_2d.ConservativeTracerEquation2D(
-                    self.function_spaces.Q_2d, self.depth, self.options, uv_2d)
-            else:
-                self.equations[label] = tracer_eq_2d.TracerEquation2D(
-                    self.function_spaces.Q_2d, self.depth, self.options, uv_2d)
-        self.solve_tracer = self.options.tracer != {}
-        if self.solve_tracer:
-            if self.options.use_limiter_for_tracers and self.options.polynomial_degree > 0:
-                self.tracer_limiter = limiter.VertexBasedP1DGLimiter(self.function_spaces.Q_2d)
-            else:
-                self.tracer_limiter = None
+        uv_2d, elev_2d = self.fields.solution_2d.subfunctions
+
+        # Passive tracer equations
+        for system, parent in self.options.tracer_fields.items():
+            self.equations[system] = tracer_eq_2d.TracerEquation2D(
+                system,
+                parent.function_space(),
+                self.depth,
+                self.options,
+                uv_2d,
+            )
 
+        # Equation for suspended sediment transport
         sediment_options = self.options.sediment_model_options
         if sediment_options.solve_suspended_sediment or sediment_options.solve_exner:
             sediment_model_class = self.options.sediment_model_options.sediment_model_class
             self.sediment_model = sediment_model_class(
                 self.options, self.mesh2d, uv_2d, elev_2d, self.depth)
         if sediment_options.solve_suspended_sediment:
-            self.fields.sediment_2d = Function(self.function_spaces.Q_2d, name='sediment_2d')
             self.equations.sediment = sediment_eq_2d.SedimentEquation2D(
                 self.function_spaces.Q_2d, self.depth, self.options, self.sediment_model,
                 conservative=sediment_options.use_sediment_conservative_form)
-            if self.options.use_limiter_for_tracers and self.options.polynomial_degree > 0:
-                self.tracer_limiter = limiter.VertexBasedP1DGLimiter(self.function_spaces.Q_2d)
-            else:
-                self.tracer_limiter = None
 
+        # Exner equation for bedload transport
         if sediment_options.solve_exner:
             if element_continuity(self.fields.bathymetry_2d.function_space().ufl_element()).horizontal in ['cg']:
                 self.equations.exner = exner_eq.ExnerEquation(
@@ -519,23 +565,24 @@ 
    Source code for thetis.solver2d
                 else:
                 raise NotImplementedError("Exner equation can currently only be implemented if the bathymetry is defined on a continuous space")
 
+        # Free surface equation for non-hydrostatic pressure
         if self.options.nh_model_options.solve_nonhydrostatic_pressure:
             print_output('Using non-hydrostatic pressure')
-            q_degree = self.options.polynomial_degree
-            if self.options.nh_model_options.q_degree is not None:
-                q_degree = self.options.nh_model_options.q_degree
-            fs_q = get_functionspace(self.mesh2d, 'CG', q_degree)
-            self.fields.q_2d = Function(fs_q, name='q_2d')  # 2D non-hydrostatic pressure at bottom
-            self.fields.w_2d = Function(self.function_spaces.H_2d, name='w_2d')  # depth-averaged vertical velocity
-            # free surface equation
             self.equations.fs = shallowwater_eq.FreeSurfaceEquation(
                 TestFunction(self.function_spaces.H_2d), self.function_spaces.H_2d, self.function_spaces.U_2d,
                 self.depth, self.options)
-            self.equations.fs.bnd_functions = self.bnd_functions['shallow_water']
    
+            self.equations.fs.bnd_functions = self.bnd_functions['shallow_water']
+
+        # Setup slope limiters
+        if self.solve_tracer or sediment_options.solve_suspended_sediment:
+            if self.options.use_limiter_for_tracers and self.options.polynomial_degree > 0:
+                self.tracer_limiter = limiter.VertexBasedP1DGLimiter(self.function_spaces.Q_2d)
+            else:
+                self.tracer_limiter = None
    
 
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.get_swe_timestepper")
     def get_swe_timestepper(self, integrator):
-        """
+        """
         Gets shallow water timestepper object with appropriate parameters
         """
         fields = {
@@ -557,7 +604,8 @@ 
    Source code for thetis.solver2d
                 self.equations.mom = shallowwater_eq.ShallowWaterMomentumEquation(
                 u_test, self.function_spaces.U_2d, self.function_spaces.H_2d,
                 self.depth,
-                options=self.options
+                options=self.options,
+                tidal_farms=self.tidal_farms
             )
             self.equations.mom.bnd_functions = bnd_conditions
             return integrator(self.equations.sw, self.equations.mom, self.fields.solution_2d,
@@ -567,38 +615,40 @@ 
    Source code for thetis.solver2d
                                   self.options.swe_timestepper_options, bnd_conditions)
    
 
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.get_tracer_timestepper")
-    def get_tracer_timestepper(self, integrator, label):
-        """
+    def get_tracer_timestepper(self, integrator, system):
+        """
         Gets tracer timestepper object with appropriate parameters
         """
-        uv, elev = self.fields.solution_2d.split()
+        uv, elev = self.fields.solution_2d.subfunctions
         fields = {
             'elev_2d': elev,
             'uv_2d': uv,
-            'diffusivity_h': self.options.tracer[label].diffusivity,
-            'source': self.options.tracer[label].source,
             'lax_friedrichs_tracer_scaling_factor': self.options.lax_friedrichs_tracer_scaling_factor,
             'tracer_advective_velocity_factor': self.options.tracer_advective_velocity_factor,
         }
+        for label in system.split(','):
+            fields[f'diffusivity_h-{label}'] = self.options.tracer[label].diffusivity
+            fields[f'source-{label}'] = self.options.tracer[label].source
         bcs = {}
-        if label in self.bnd_functions:
-            bcs = self.bnd_functions[label]
-        elif label[:-3] in self.bnd_functions:
-            bcs = self.bnd_functions[label[:-3]]
-        # TODO: Different timestepper options for different tracers
-        return integrator(self.equations[label], self.fields[label], fields, self.dt,
+        if system in self.bnd_functions:
+            bcs = self.bnd_functions[system]
+        elif system[:-3] in self.bnd_functions:
+            # TODO: Is this safe for monolithic systems?
+            bcs = self.bnd_functions[system[:-3]]
+        # TODO: Different timestepper options for different systems
+        return integrator(self.equations[system], self.fields[system], fields, self.dt,
                           self.options.tracer_timestepper_options, bcs)
    
 
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.get_sediment_timestepper")
     def get_sediment_timestepper(self, integrator):
-        """
+        """
         Gets sediment timestepper object with appropriate parameters
         """
-        uv, elev = self.fields.solution_2d.split()
+        uv, elev = self.fields.solution_2d.subfunctions
         fields = {
             'elev_2d': elev,
             'uv_2d': uv,
-            'diffusivity_h': self.options.sediment_model_options.horizontal_diffusivity,
+            'diffusivity_h-sediment_2d': self.options.sediment_model_options.horizontal_diffusivity,
             'lax_friedrichs_tracer_scaling_factor': self.options.lax_friedrichs_tracer_scaling_factor,
             'tracer_advective_velocity_factor': self.sediment_model.get_advective_velocity_correction_factor(),
         }
@@ -608,10 +658,10 @@ 
    Source code for thetis.solver2d
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.get_exner_timestepper")
     def get_exner_timestepper(self, integrator):
-        """
+        """
         Gets exner timestepper object with appropriate parameters
         """
-        uv, elev = self.fields.solution_2d.split()
+        uv, elev = self.fields.solution_2d.subfunctions
         if not self.options.sediment_model_options.solve_suspended_sediment:
             self.fields.sediment_2d = Function(elev.function_space()).interpolate(Constant(0.0))
         fields = {
@@ -627,7 +677,7 @@ 
    Source code for thetis.solver2d
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.get_fs_timestepper")
     def get_fs_timestepper(self, integrator):
-        """
+        """
         Gets free-surface correction timestepper object with appropriate parameters
         """
         fields_fs = {
@@ -641,7 +691,7 @@ 
    Source code for thetis.solver2d
     
    [docs]    @unfrozen
     @PETSc.Log.EventDecorator("thetis.FlowSolver2d.create_timestepper")
     def create_timestepper(self):
-        """
+        """
         Creates time stepper instance
         """
         if not hasattr(self, 'equations'):
@@ -694,7 +744,7 @@ 
    Source code for thetis.solver2d
     
    [docs]    @unfrozen
     @PETSc.Log.EventDecorator("thetis.FlowSolver2d.create_exporters")
     def create_exporters(self):
-        """
+        """
         Creates file exporters
         """
         if not hasattr(self, 'timestepper'):
@@ -720,7 +770,7 @@ 
    Source code for thetis.solver2d
                 self.exporters['hdf5'] = e
    
 
 
    [docs]    def initialize(self):
-        """
+        """
         Creates function spaces, equations, time stepper and exporters
         """
         if not hasattr(self.function_spaces, 'U_2d'):
@@ -735,7 +785,7 @@ 
    Source code for thetis.solver2d
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.assign_initial_conditions")
     def assign_initial_conditions(self, elev=None, uv=None, sediment=None, **tracers):
-        r"""
+        r"""
         Assigns initial conditions
 
         :kwarg elev: Initial condition for water elevation
@@ -749,7 +799,7 @@ 
    Source code for thetis.solver2d
             """
         if not self._initialized:
             self.initialize()
-        uv_2d, elev_2d = self.fields.solution_2d.split()
+        uv_2d, elev_2d = self.fields.solution_2d.subfunctions
         if elev is not None:
             elev_2d.project(elev)
         if uv is not None:
@@ -776,7 +826,7 @@ 
    Source code for thetis.solver2d
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.add_callback")
     def add_callback(self, callback, eval_interval='export'):
-        """
+        """
         Adds callback to solver object
 
         :arg callback: :class:`.DiagnosticCallback` instance
@@ -788,7 +838,7 @@ 
    Source code for thetis.solver2d
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.export")
     def export(self):
-        """
+        """
         Export all fields to disk
 
         Also evaluates all callbacks set to 'export' interval.
@@ -798,25 +848,50 @@ 
    Source code for thetis.solver2d
                 e.export()
    
 
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.load_state")
-    def load_state(self, i_export, outputdir=None, t=None, iteration=None):
-        """
+    def load_state(self, i_stored, outputdir=None, t=None, iteration=None,
+                   i_export=None, legacy_mode=False):
+        """
         Loads simulation state from hdf5 outputs.
 
-        This replaces :meth:`.assign_initial_conditions` in model initilization.
+        Replaces :meth:`.assign_initial_conditions` in model initialization.
+
+        For example, continue simulation from export 40,
+
+        .. code-block:: python
+
+            solver_obj.load_state(40)
+
+        Continue simulation from another output directory,
+
+        .. code-block:: python
 
-        This assumes that model setup is kept the same (e.g. time step) and
-        all pronostic state variables are exported in hdf5 format. The required
-        state variables are: elev_2d, uv_2d
+            solver_obj.load_state(40, outputdir='outputs_spinup')
 
-        Currently hdf5 field import only works for the same number of MPI
-        processes.
+        Continue simulation from another output directory and reset the
+        counters,
 
-        :arg int i_export: export index to load
+        .. code-block:: python
+
+            solver_obj.load_state(40, outputdir='outputs_spinup',
+                iteration=0, t=0, i_export=0)
+
+        Model state can be loaded if all prognostic state variables have been
+        stored in hdf5 format. For the hydrodynamics the required state
+        variables are: `elev_2d`, and `uv_2d`.
+
+        Simulation time and iteration can be recovered automatically provided
+        that the model setup is the same (e.g. time step and export_time).
+
+        :arg int i_stored: export index to load
         :kwarg string outputdir: (optional) directory where files are read from.
             By default ``options.output_directory``.
-        :kwarg float t: simulation time. Overrides the time stamp stored in the
-            hdf5 files.
-        :kwarg int iteration: Overrides the iteration count in the hdf5 files.
+        :kwarg float t: simulation time. Overrides the time stamp inferred from
+            model options.
+        :kwarg int iteration: Overrides the iteration count inferred from model
+            options.
+        :kwarg int i_export: Set initial export index for the present run. By
+            default, `i_stored` is used.
+        :kwarg bool legacy_mode: Load legacy `DumbCheckpoint` files.
         """
         if not self._initialized:
             self.initialize()
@@ -830,12 +905,15 @@ 
    Source code for thetis.solver2d
                                        self.fields,
                                    field_metadata,
                                    export_type='hdf5',
+                                   legacy_mode=legacy_mode,
                                    verbose=self.options.verbose > 0)
-        e.exporters['uv_2d'].load(i_export, self.fields.uv_2d)
-        e.exporters['elev_2d'].load(i_export, self.fields.elev_2d)
+        e.exporters['uv_2d'].load(i_stored, self.fields.uv_2d)
+        e.exporters['elev_2d'].load(i_stored, self.fields.elev_2d)
         self.assign_initial_conditions()
 
         # time stepper bookkeeping for export time step
+        if i_export is None:
+            i_export = i_stored
         self.i_export = i_export
         self.next_export_t = self.i_export*self.options.simulation_export_time
         if iteration is None:
@@ -855,39 +933,60 @@ 
    Source code for thetis.solver2d
             for e in self.exporters.values():
             e.set_next_export_ix(self.i_export + offset)
    
 
-
    [docs]    def print_state(self, cputime):
-        """
+
    [docs]    def print_state(self, cputime, print_header=False):
+        """
         Print a summary of the model state on stdout
 
-        :arg float cputime: Measured CPU time
+        :arg float cputime: Measured CPU time in seconds
+        :kwarg print_header: Whether to print column header first
         """
+        # list of entries to print: (header, value, format)
+        entries = [
+            ('exp', self.i_export, '5d'),
+            ('iter', self.iteration, '5d'),
+        ]
+        # generate simulation time string
+        if self.options.simulation_initial_date is not None:
+            now = self.options.simulation_initial_date + datetime.timedelta(seconds=self.simulation_time)
+            date_str = f'{now:%Y-%m-%d}'.rjust(11)
+            time_str = f'{now:%H:%M:%S}'.rjust(9)
+            entries += [
+                ('date', date_str, '11s'),
+                ('time', time_str, '9s'),
+            ]
+        else:
+            time_str = f'{self.simulation_time:.2f}'.rjust(15)
+            entries += [
+                ('time', time_str, '15s'),
+            ]
         if self.options.tracer_only:
             for label in self.options.tracer:
                 norm_q = norm(self.fields[label])
-
-                line = ('{iexp:5d} {i:5d} T={t:10.2f} '
-                        '{label:16s}: {q:10.4f} {cpu:5.2f}')
-
-                norm_label = label if len(label) < 3 or label[-3:] != '_2d' else label[:-3]
-                print_output(line.format(iexp=self.i_export, i=self.iteration,
-                                         t=self.simulation_time,
-                                         label=norm_label + ' norm',
-                                         q=norm_q, cpu=cputime))
+                e = (label, norm_q, '10.4f')
+                entries.append(e)
         else:
-            norm_h = norm(self.fields.solution_2d.split()[1])
-            norm_u = norm(self.fields.solution_2d.split()[0])
-
-            line = ('{iexp:5d} {i:5d} T={t:10.2f} '
-                    'eta norm: {e:10.4f} u norm: {u:10.4f} {cpu:5.2f}')
-            print_output(line.format(iexp=self.i_export, i=self.iteration,
-                                     t=self.simulation_time, e=norm_h,
-                                     u=norm_u, cpu=cputime))
+            norm_h = norm(self.fields.solution_2d.subfunctions[1])
+            norm_u = norm(self.fields.solution_2d.subfunctions[0])
+            entries += [
+                ('eta norm', norm_h, '14.4f'),
+                ('u norm', norm_u, '14.4f'),
+            ]
+        entries.append(('Tcpu', cputime, '6.2f'))
+
+        if print_header:
+            # generate header
+            header = ' '.join([e[0].rjust(len(f'{e[1]:{e[2]}}')) for e in entries])
+            print_output(header)
+
+        # generate line
+        line = ' '.join([f'{e[1]:{e[2]}}' for e in entries])
+        print_output(line)
         sys.stdout.flush()
    
 
 
    [docs]    @PETSc.Log.EventDecorator("thetis.FlowSolver2d.iterate")
     def iterate(self, update_forcings=None,
                 export_func=None):
-        """
+        """
         Runs the simulation
 
         Iterates over the time loop until time ``options.simulation_end_time`` is reached.
@@ -899,7 +998,6 @@ 
    Source code for thetis.solver2d
             :kwarg export_func: User-defined function (with no arguments) that will
             be called on every export.
         """
-        # TODO I think export function is obsolete as callbacks are in place
         if not self._initialized:
             self.initialize()
 
@@ -956,8 +1054,27 @@ 
    Source code for thetis.solver2d
                                                      append_to_log=True)
             self.add_callback(c, eval_interval='export')
 
+        initial_simulation_time = self.simulation_time
+        internal_iteration = 0
+
+        init_date = self.options.simulation_initial_date
+        end_date = self.options.simulation_end_date
+        if (init_date is not None and end_date is not None):
+            now = init_date + datetime.timedelta(initial_simulation_time)
+            assert end_date > now, f'Simulation end date must be greater than initial time {now}'
+            print_output(
+                f'Running simulation\n'
+                f' from {now:%Y-%m-%d %H:%M:%S %Z}\n'
+                f' to   {end_date:%Y-%m-%d %H:%M:%S %Z}'
+            )
+            end_time = (end_date - now).total_seconds() + initial_simulation_time
+            if self.options.simulation_end_time is not None:
+                warning('Both simulation_end_date and simulation_end_time have been set, ignoring simulation_end_time.')
+            self.options.simulation_end_time = end_time
+        assert self.options.simulation_end_time is not None, 'simulation_end_time must be set'
+
         # initial export
-        self.print_state(0.0)
+        self.print_state(0.0, print_header=True)
         if self.export_initial_state:
             self.export()
             if export_func is not None:
@@ -965,9 +1082,6 @@ 
    Source code for thetis.solver2d
                 if 'vtk' in self.exporters and isinstance(self.fields.bathymetry_2d, Function):
                 self.exporters['vtk'].export_bathymetry(self.fields.bathymetry_2d)
 
-        initial_simulation_time = self.simulation_time
-        internal_iteration = 0
-
         while self.simulation_time <= self.options.simulation_end_time - t_epsilon:
             self.timestepper.advance(self.simulation_time, update_forcings)
 
@@ -1000,8 +1114,8 @@ 
    Source code for thetis.solver2d
           
    
     
    
     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/stability_functions.html b/_modules/thetis/stability_functions.html
index 84e31d7..a663248 100644
--- a/_modules/thetis/stability_functions.html
+++ b/_modules/thetis/stability_functions.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.stability_functions — Thetis 0+untagged.1878.g52658ff.dirty documentation
+    thetis.stability_functions — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -107,7 +107,7 @@ 
    Source code for thetis.stability_functions
     
 
 
    [docs]def compute_normalized_frequencies(shear2, buoy2, k, eps, verbose=False):
-    r"""
+    r"""
     Computes normalized buoyancy and shear frequency squared.
 
     .. math::
@@ -132,7 +132,7 @@ 
    Source code for thetis.stability_functions
     
 
 
    [docs]class StabilityFunctionBase(ABC):
-    """
+    """
     Base class for all stability functions
     """
     @property
@@ -142,7 +142,7 @@ 
    Source code for thetis.stability_functions
     
     def __init__(self, lim_alpha_shear=True, lim_alpha_buoy=True,
                  smooth_alpha_buoy_lim=True, alpha_buoy_crit=-1.2):
-        r"""
+        r"""
         :kwarg bool lim_alpha_shear: limit maximum :math:`\alpha_M` values
             (see Umlauf and Burchard (2005) eq. 44)
         :kwarg bool lim_alpha_buoy: limit minimum (negative) :math:`\alpha_N` values
@@ -186,7 +186,7 @@ 
    Source code for thetis.stability_functions
             self.nb2 = 0.0
 
 
    [docs]    def compute_alpha_shear_steady(self, ri_st, analytical=True):
-        r"""
+        r"""
         Compute the steady-state :math:`\alpha_M`.
 
         Under steady-state conditions, the stability functions satisfy:
@@ -230,7 +230,7 @@ 
    Source code for thetis.stability_functions
             return a_shear
    
 
 
    [docs]    def compute_c3_minus(self, c1, c2, ri_st):
-        r"""
+        r"""
         Compute c3_minus parameter from c1 and c2 parameters.
 
         c3_minus is solved from equation
@@ -259,7 +259,7 @@ 
    Source code for thetis.stability_functions
             return c3_minus
    
 
 
    [docs]    def compute_cmu0(self, analytical=True):
-        r"""
+        r"""
         Compute parameter :math:`c_\mu^0`
 
         See: Umlauf and Buchard (2005) eq A.22
@@ -290,7 +290,7 @@ 
    Source code for thetis.stability_functions
             return cm0
    
 
 
    [docs]    def compute_kappa(self, sigma_psi, cmu0, n, c1, c2):
-        r"""
+        r"""
         Computes von Karman constant from the Psi Schmidt number.
 
         See: Umlauf and Burchard (2003) eq (14)
@@ -301,7 +301,7 @@ 
    Source code for thetis.stability_functions
             return cmu0 / numpy.abs(n) * numpy.sqrt(sigma_psi * (c2 - c1))
    
 
 
    [docs]    def compute_sigma_psi(self, kappa, cmu0, n, c1, c2):
-        r"""
+        r"""
         Computes the Psi Schmidt number from von Karman constant.
 
         See: Umlauf and Burchard (2003) eq (14)
@@ -312,7 +312,7 @@ 
    Source code for thetis.stability_functions
             return (n * kappa)**2 / (cmu0**2 * (c2 - c1))
    
 
 
    [docs]    def compute_length_clim(self, cmu0, ri_st):
-        r"""
+        r"""
         Computes the Galpering length scale limit.
 
         :arg cmu0: parameter :math:`c_\mu^0`
@@ -327,7 +327,7 @@ 
    Source code for thetis.stability_functions
             return clim
    
 
 
    [docs]    def get_alpha_buoy_min(self):
-        r"""
+        r"""
         Compute minimum normalized buoyancy frequency :math:`\alpha_N`
 
         See: Umlauf and Buchard (2005), Table 3
@@ -339,7 +339,7 @@ 
    Source code for thetis.stability_functions
             return an_min
    
 
 
    [docs]    def get_alpha_shear_max(self, alpha_buoy, alpha_shear):
-        r"""
+        r"""
         Compute maximum normalized shear frequency :math:`\alpha_M`
 
         from Umlauf and Buchard (2005) eq (44)
@@ -357,7 +357,7 @@ 
    Source code for thetis.stability_functions
             return as_max_n/as_max_d
    
 
 
    [docs]    def get_alpha_buoy_smooth_min(self, alpha_buoy):
-        r"""
+        r"""
         Compute smoothed minimum for normalized buoyancy frequency
 
         See: Burchard and Petersen (1999), eq (19)
@@ -367,7 +367,7 @@ 
    Source code for thetis.stability_functions
             return alpha_buoy - (alpha_buoy - self.alpha_buoy_crit)**2/(alpha_buoy + self.get_alpha_buoy_min() - 2*self.alpha_buoy_crit)
    
 
 
    [docs]    def eval_funcs(self, alpha_buoy, alpha_shear):
-        r"""
+        r"""
         Evaluate (unlimited) stability functions
 
         See: Burchard and Petersen (1999) eqns (30) and (31)
@@ -382,7 +382,7 @@ 
    Source code for thetis.stability_functions
             return c_mu, c_mu_p
    
 
 
    [docs]    def evaluate(self, shear2, buoy2, k, eps):
-        r"""
+        r"""
         Evaluate stability functions from dimensional variables.
 
         Applies limiters on :math:`\alpha_N` and :math:`\alpha_M`.
@@ -414,7 +414,7 @@ 
    Source code for thetis.stability_functions
     
 
 class GOTMStabilityFunctionBase(StabilityFunctionBase, ABC):
-    """
+    """
     Base class for stability functions defined in Umlauf and Buchard (2005)
     """
     @property
@@ -479,7 +479,7 @@ 
    Source code for thetis.stability_functions
     
     def __init__(self, lim_alpha_shear=True, lim_alpha_buoy=True,
                  smooth_alpha_buoy_lim=True, alpha_buoy_crit=-1.2):
-        r"""
+        r"""
         :kwarg bool lim_alpha_shear: limit maximum :math:`\alpha_M` values
             (see Umlauf and Burchard (2005) eq. 44)
         :kwarg bool lim_alpha_buoy: limit minimum (negative) :math:`\alpha_N` values
@@ -530,7 +530,7 @@ 
    Source code for thetis.stability_functions
     
 
 class CanutoStabilityFunctionBase(StabilityFunctionBase, ABC):
-    """
+    """
     Base class for original Canuto stability function.
     """
     @property
@@ -575,7 +575,7 @@ 
    Source code for thetis.stability_functions
     
     def __init__(self, lim_alpha_shear=True, lim_alpha_buoy=True,
                  smooth_alpha_buoy_lim=True, alpha_buoy_crit=-1.2):
-        r"""
+        r"""
         :kwarg bool lim_alpha_shear: limit maximum :math:`\alpha_M` values
             (see Umlauf and Burchard (2005) eq. 44)
         :kwarg bool lim_alpha_buoy: limit minimum (negative) :math:`\alpha_N` values
@@ -620,7 +620,7 @@ 
    Source code for thetis.stability_functions
             self.nb2 = self.cu_scalar*self.alpha_scalar*self.s6
 
     def eval_funcs_new(self, alpha_buoy, alpha_shear):
-        r"""
+        r"""
         Evaluate (unlimited) stability functions
 
         From Canuto et al (2001)
@@ -639,7 +639,7 @@ 
    Source code for thetis.stability_functions
     
 
 class ChengStabilityFunctionBase(StabilityFunctionBase):
-    """
+    """
     Base class for original Cheng stability function.
     """
     @property
@@ -684,7 +684,7 @@ 
    Source code for thetis.stability_functions
     
     def __init__(self, lim_alpha_shear=True, lim_alpha_buoy=True,
                  smooth_alpha_buoy_lim=True, alpha_buoy_crit=-1.2):
-        r"""
+        r"""
         :kwarg bool lim_alpha_shear: limit maximum :math:`\alpha_M` values
             (see Umlauf and Burchard (2005) eq. 44)
         :kwarg bool lim_alpha_buoy: limit minimum (negative) :math:`\alpha_N` values
@@ -729,7 +729,7 @@ 
    Source code for thetis.stability_functions
             self.nb2 = self.cu_scalar*self.alpha_scalar*self.s6
 
     def eval_funcs_new(self, alpha_buoy, alpha_shear):
-        r"""
+        r"""
         Evaluate (unlimited) stability functions
 
         From Canuto et al (2001)
@@ -748,7 +748,7 @@ 
    Source code for thetis.stability_functions
     
 
 
    [docs]class StabilityFunctionCanutoA(CanutoStabilityFunctionBase):
-    """
+    """
     Canuto A stability function as defined in the Canuto (2001) paper.
     """
     l1 = 0.107
@@ -763,7 +763,7 @@ 
    Source code for thetis.stability_functions
     
 
 
    [docs]class StabilityFunctionCanutoB(CanutoStabilityFunctionBase):
-    """
+    """
     Canuto B stability function as defined in the Canuto (2001) paper.
     """
     l1 = 0.127
@@ -778,7 +778,7 @@ 
    Source code for thetis.stability_functions
     
 
 
    [docs]class StabilityFunctionCheng(ChengStabilityFunctionBase):
-    """
+    """
     Cheng stability function as defined in the Cheng et al (2002) paper.
     """
     l1 = 0.107
@@ -793,7 +793,7 @@ 
    Source code for thetis.stability_functions
     
 
 
    [docs]class GOTMStabilityFunctionCanutoA(GOTMStabilityFunctionBase):
-    """
+    """
     Canuto et al. (2001) version A stability functions
 
     Parameters are from Umlauf and Buchard (2005), Table 1
@@ -814,7 +814,7 @@ 
    Source code for thetis.stability_functions
     
 
 
    [docs]class GOTMStabilityFunctionCanutoB(GOTMStabilityFunctionBase):
-    """
+    """
     Canuto et al. (2001) version B stability functions
 
     Parameters are from Umlauf and Buchard (2005), Table 1
@@ -835,7 +835,7 @@ 
    Source code for thetis.stability_functions
     
 
 
    [docs]class GOTMStabilityFunctionKanthaClayson(GOTMStabilityFunctionBase):
-    """
+    """
     Kantha and Clayson (1994) quasi-equilibrium stability functions
 
     Parameters are from Umlauf and Buchard (2005), Table 1
@@ -856,7 +856,7 @@ 
    Source code for thetis.stability_functions
     
 
 
    [docs]class GOTMStabilityFunctionCheng(GOTMStabilityFunctionBase):
-    """
+    """
     Cheng et al. (2002) quasi-equilibrium stability functions
 
     Parameters are from Umlauf and Buchard (2005), Table 1
@@ -884,8 +884,8 @@ 
    Source code for thetis.stability_functions
           
    
     
    
     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/timeintegrator.html b/_modules/thetis/timeintegrator.html
index 1f9dd67..ebf8e22 100644
--- a/_modules/thetis/timeintegrator.html
+++ b/_modules/thetis/timeintegrator.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.timeintegrator — Thetis 0+untagged.1888.gc4e776a documentation
+    thetis.timeintegrator — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -55,25 +55,24 @@ 
    Source code for thetis.timeintegrator
     Generic time integration schemes to advance equations in time.
 """
 from .utility import *
-from abc import ABCMeta, abstractmethod
+from abc import ABC, abstractmethod
 import numpy
+from pyop2.profiling import timed_region, timed_stage
 
 CFL_UNCONDITIONALLY_STABLE = numpy.inf
 # CFL coefficient for unconditionally stable methods
 
 
-
    [docs]class TimeIntegratorBase(object):
-    """
+
    [docs]class TimeIntegratorBase(ABC):
+    """
     Abstract class that defines the API for all time integrators
 
     Both :class:`TimeIntegrator` and :class:`CoupledTimeIntegrator` inherit
     from this class.
     """
-    __metaclass__ = ABCMeta
-
 
    [docs]    @abstractmethod
     def advance(self, t, update_forcings=None):
-        """
+        """
         Advances equations for one time step
 
         :arg t: simulation time
@@ -85,7 +84,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @abstractmethod
     def initialize(self, init_solution):
-        """
+        """
         Initialize the time integrator
 
         :arg init_solution: initial solution
@@ -94,11 +93,11 @@ 
    Source code for thetis.timeintegrator
     
 
 
    [docs]class TimeIntegrator(TimeIntegratorBase):
-    """
+    """
     Base class for all time integrator objects that march a single equation
     """
     def __init__(self, equation, solution, fields, dt, options):
-        """
+        """
         :arg equation: the equation to solve
         :type equation: :class:`Equation` object
         :arg solution: :class:`Function` where solution will be stored
@@ -122,12 +121,12 @@ 
    Source code for thetis.timeintegrator
             self.solver_parameters = options.solver_parameters
 
 
    [docs]    def set_dt(self, dt):
-        """Update time step"""
+        """Update time step"""
         self.dt = dt
         self.dt_const.assign(dt)
    
 
 
    [docs]    def advance_picard(self, t, update_forcings=None, update_lagged=True, update_fields=True):
-        """
+        """
         Advances equations for one time step within a Picard iteration
 
         :arg t: simulation time
@@ -137,16 +136,42 @@ 
    Source code for thetis.timeintegrator
             :kwarg update_lagged: should the old solution be updated?
         :kwarg update_fields: should the fields be updated?
         """
-        raise NotImplementedError(f"Picard iterations are not supported for {self} time integrators.")
    
+        raise NotImplementedError(f"Picard iterations are not supported for {self} time integrators.")
    
+
+
    [docs]    def create_fields_old(self):
+        """
+        Create a copy of fields dictionary to store beginning of timestep values
+        """
+        # create functions to hold the values of previous time step
+        self.fields_old = {}
+        for k in sorted(self.fields):
+            if self.fields[k] is not None:
+                if isinstance(self.fields[k], Function):
+                    self.fields_old[k] = Function(
+                        self.fields[k].function_space())
+                elif isinstance(self.fields[k], Constant):
+                    # Although Constants may be changed by the user, we just
+                    # use the same value here, as assigning to domain-less Constants
+                    # causes issues in the adjoint. Constants with a domain, created by Constaint(.., domain=...),
+                    # are now just Functions on the Real space, so are dealt with under isinstance(..., Function)
+                    self.fields_old[k] = self.fields[k]
    
+
+
    [docs]    def update_fields_old(self):
+        """
+        Update the values of fields_old with those of fields
+        """
+        for k in sorted(self.fields):
+            if isinstance(self.fields[k], Function):
+                self.fields_old[k].assign(self.fields[k])
    
 
 
 
    [docs]class ForwardEuler(TimeIntegrator):
-    """Standard forward Euler time integration scheme."""
+    """Standard forward Euler time integration scheme."""
     cfl_coeff = 1.0
 
     @PETSc.Log.EventDecorator("thetis.ForwardEuler.__init__")
     def __init__(self, equation, solution, fields, dt, options, bnd_conditions):
-        """
+        """
         :arg equation: the equation to solve
         :type equation: :class:`Equation` object
         :arg solution: :class:`Function` where solution will be stored
@@ -159,15 +184,7 @@ 
    Source code for thetis.timeintegrator
             super(ForwardEuler, self).__init__(equation, solution, fields, dt, options)
         self.solution_old = Function(self.equation.function_space)
 
-        # create functions to hold the values of previous time step
-        self.fields_old = {}
-        for k in sorted(self.fields):
-            if self.fields[k] is not None:
-                if isinstance(self.fields[k], Function):
-                    self.fields_old[k] = Function(
-                        self.fields[k].function_space())
-                elif isinstance(self.fields[k], Constant):
-                    self.fields_old[k] = Constant(self.fields[k])
+        self.create_fields_old()
 
         u_old = self.solution_old
         u_tri = self.equation.trial
@@ -187,31 +204,27 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.ForwardEuler.initialize")
     def initialize(self, solution):
-        """Assigns initial conditions to all required fields."""
+        """Assigns initial conditions to all required fields."""
         self.solution_old.assign(solution)
-        # assign values to old functions
-        for k in sorted(self.fields_old):
-            self.fields_old[k].assign(self.fields[k])
    
+        self.update_fields_old()
    
 
 
    [docs]    @PETSc.Log.EventDecorator("thetis.ForwardEuler.advance")
     def advance(self, t, update_forcings=None):
-        """Advances equations for one time step."""
+        """Advances equations for one time step."""
         if update_forcings is not None:
             update_forcings(t + self.dt)
         self.solution_old.assign(self.solution)
         self.solver.solve()
-        # shift time
-        for k in sorted(self.fields_old):
-            self.fields_old[k].assign(self.fields[k])
    
+        self.update_fields_old()
    
 
 
 
    [docs]class CrankNicolson(TimeIntegrator):
-    """Standard Crank-Nicolson time integration scheme."""
+    """Standard Crank-Nicolson time integration scheme."""
     cfl_coeff = CFL_UNCONDITIONALLY_STABLE
 
     @PETSc.Log.EventDecorator("thetis.CrankNicolson.__init__")
     def __init__(self, equation, solution, fields, dt, options, bnd_conditions):
-        """
+        """
         :arg equation: the equation to solve
         :type equation: :class:`Equation` object
         :arg solution: :class:`Function` where solution will be stored
@@ -229,16 +242,7 @@ 
    Source code for thetis.timeintegrator
             else:
             self.solver_parameters.setdefault('snes_type', 'newtonls')
         self.solution_old = Function(self.equation.function_space, name='solution_old')
-        # create functions to hold the values of previous time step
-        # TODO is this necessary? is self.fields sufficient?
-        self.fields_old = {}
-        for k in sorted(self.fields):
-            if self.fields[k] is not None:
-                if isinstance(self.fields[k], Function):
-                    self.fields_old[k] = Function(
-                        self.fields[k].function_space(), name=self.fields[k].name()+'_old')
-                elif isinstance(self.fields[k], Constant):
-                    self.fields_old[k] = Constant(self.fields[k])
+        self.create_fields_old()
 
         u = self.solution
         u_old = self.solution_old
@@ -264,7 +268,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.CrankNicolson.update_solver")
     def update_solver(self):
-        """Create solver objects"""
+        """Create solver objects"""
         # Ensure LU assembles monolithic matrices
         if self.solver_parameters.get('pc_type') == 'lu':
             self.solver_parameters['mat_type'] = 'aij'
@@ -276,39 +280,33 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.CrankNicolson.initialize")
     def initialize(self, solution):
-        """Assigns initial conditions to all required fields."""
+        """Assigns initial conditions to all required fields."""
         self.solution_old.assign(solution)
-        # assign values to old functions
-        for k in sorted(self.fields_old):
-            self.fields_old[k].assign(self.fields[k])
    
+        self.update_fields_old()
    
 
 
    [docs]    @PETSc.Log.EventDecorator("thetis.CrankNicolson.advance")
     def advance(self, t, update_forcings=None):
-        """Advances equations for one time step."""
+        """Advances equations for one time step."""
         if update_forcings is not None:
             update_forcings(t + self.dt)
         self.solution_old.assign(self.solution)
         self.solver.solve()
-        # shift time
-        for k in sorted(self.fields_old):
-            self.fields_old[k].assign(self.fields[k])
    
+        self.update_fields_old()
    
 
 
    [docs]    @PETSc.Log.EventDecorator("thetis.CrankNicolson.advance_picard")
     def advance_picard(self, t, update_forcings=None, update_lagged=True, update_fields=True):
-        """Advances equations for one time step in a Picard iteration."""
+        """Advances equations for one time step in a Picard iteration."""
         if update_forcings is not None:
             update_forcings(t + self.dt)
         if update_lagged:
             self.solution_old.assign(self.solution)
         self.solver.solve()
         if update_fields:
-            # shift time
-            for k in sorted(self.fields_old):
-                self.fields_old[k].assign(self.fields[k])
    
+            self.update_fields_old()
    
 
 
 
    [docs]class SteadyState(TimeIntegrator):
-    """
+    """
     Time integrator that solves the steady state equations, leaving out the
     mass terms
     """
@@ -316,7 +314,7 @@ 
    Source code for thetis.timeintegrator
     
     @PETSc.Log.EventDecorator("thetis.SteadyState.__init__")
     def __init__(self, equation, solution, fields, dt, options, bnd_conditions):
-        """
+        """
         :arg equation: the equation to solve
         :type equation: :class:`Equation` object
         :arg solution: :class:`Function` where solution will be stored
@@ -333,7 +331,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.SteadyState.update_solver")
     def update_solver(self):
-        """Create solver objects"""
+        """Create solver objects"""
         # Ensure LU assembles monolithic matrices
         if self.solver_parameters.get('pc_type') == 'lu':
             self.solver_parameters['mat_type'] = 'aij'
@@ -345,20 +343,20 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.SteadyState.initialize")
     def initialize(self, solution):
-        """Assigns initial conditions to all required fields."""
+        """Assigns initial conditions to all required fields."""
         # nothing to do here as the initial condition is passed in via solution
         return
    
 
 
    [docs]    @PETSc.Log.EventDecorator("thetis.SteadyState.advance")
     def advance(self, t, update_forcings=None):
-        """Advances equations for one time step."""
+        """Advances equations for one time step."""
         if update_forcings is not None:
             update_forcings(t + self.dt)
         self.solver.solve()
    
 
 
 
    [docs]class PressureProjectionPicard(TimeIntegrator):
-    """
+    """
     Pressure projection scheme with Picard iteration for shallow water
     equations
 
@@ -368,7 +366,7 @@ 
    Source code for thetis.timeintegrator
         # TODO add more documentation
     @PETSc.Log.EventDecorator("thetis.PressureProjectionPicard.__init__")
     def __init__(self, equation, equation_mom, solution, fields, dt, options, bnd_conditions):
-        """
+        """
         :arg equation: free surface equation
         :type equation: :class:`Equation` object
         :arg equation_mom: momentum equation
@@ -407,8 +405,8 @@ 
    Source code for thetis.timeintegrator
                 self.solution_lagged = Function(self.equation.function_space)
         else:
             self.solution_lagged = self.solution_old
-        uv_lagged, eta_lagged = self.solution_lagged.split()
-        uv_old, eta_old = self.solution_old.split()
+        uv_lagged, eta_lagged = self.solution_lagged.subfunctions
+        uv_old, eta_old = self.solution_old.subfunctions
 
         if (solver_parameters['ksp_type'] == 'preonly'
                 and 'fieldsplit_H_2d' in solver_parameters
@@ -424,15 +422,8 @@ 
    Source code for thetis.timeintegrator
                 raise Exception("The timestepper PressureProjectionPicard is only recommended in combination with the "
                             "dg-cg element_family. If you want to use it in combination with dg-dg or rt-dg you need to adjust the solver_parameters_pressure option.")
 
-        # create functions to hold the values of previous time step
-        self.fields_old = {}
-        for k in sorted(self.fields):
-            if self.fields[k] is not None:
-                if isinstance(self.fields[k], Function):
-                    self.fields_old[k] = Function(
-                        self.fields[k].function_space())
-                elif isinstance(self.fields[k], Constant):
-                    self.fields_old[k] = Constant(self.fields[k])
+        self.create_fields_old()
+
         # for the mom. eqn. the 'eta' field is just one of the 'other' fields
         fields_mom = self.fields.copy()
         fields_mom_old = self.fields_old.copy()
@@ -485,7 +476,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.PressureProjectionPicard.update_solver")
     def update_solver(self):
-        """Create solver objects"""
+        """Create solver objects"""
         prob = NonlinearVariationalProblem(self.F_mom, self.uv_star)
         self.solver_mom = NonlinearVariationalSolver(prob,
                                                      solver_parameters=self.solver_parameters_mom,
@@ -503,16 +494,14 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.PressureProjectionPicard.initialize")
     def initialize(self, solution):
-        """Assigns initial conditions to all required fields."""
+        """Assigns initial conditions to all required fields."""
         self.solution_old.assign(solution)
         self.solution_lagged.assign(solution)
-        # assign values to old functions
-        for k in sorted(self.fields_old):
-            self.fields_old[k].assign(self.fields[k])
    
+        self.update_fields_old()
    
 
 
    [docs]    @PETSc.Log.EventDecorator("thetis.PressureProjectionPicard.advance")
     def advance(self, t, update_forcings=None):
-        """Advances equations for one time step."""
+        """Advances equations for one time step."""
         if update_forcings is not None:
             update_forcings(t + self.dt)
         self.solution_old.assign(self.solution)
@@ -525,13 +514,11 @@ 
    Source code for thetis.timeintegrator
                 with timed_stage("Pressure solve"):
                 self.solver.solve()
 
-        # shift time
-        for k in sorted(self.fields_old):
-            self.fields_old[k].assign(self.fields[k])
    
+        self.update_fields_old()
    
 
 
 
    [docs]class LeapFrogAM3(TimeIntegrator):
-    """
+    """
     Leap-Frog Adams-Moulton 3 ALE time integrator
 
     Defined in (2.27)-(2.30) in [1]; (2.21)-(2.22) in [2]
@@ -549,7 +536,7 @@ 
    Source code for thetis.timeintegrator
     
     @PETSc.Log.EventDecorator("thetis.LeapFrogAM3.__init__")
     def __init__(self, equation, solution, fields, dt, options, bnd_conditions, terms_to_add='all'):
-        """
+        """
         :arg equation: equation to solve
         :type equation: :class:`Equation` object
         :arg solution: :class:`Function` where solution will be stored
@@ -591,7 +578,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.LeapFrogAM3.initialize")
     def initialize(self, solution):
-        """Assigns initial conditions to all required fields."""
+        """Assigns initial conditions to all required fields."""
         self.mass_matrix = assemble(self.a)
         self.solution.assign(solution)
         self.solution_old.assign(solution)
@@ -600,7 +587,7 @@ 
    Source code for thetis.timeintegrator
             # TODO: Linear solver is not annotated and does not accept ad_block_tag
 
     def _solve_system(self):
-        """
+        """
         Solves system mass_matrix*solution = rhs_func
 
         If the function space is fully discontinuous, the mass matrix is
@@ -610,7 +597,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.LeapFrogAM3.predice")
     def predict(self):
-        r"""
+        r"""
         Prediction step from :math:`t_{n-1/2}` to :math:`t_{n+1/2}`
 
         Let :math:`M_n` denote the mass matrix at time :math:`t_{n}`.
@@ -640,7 +627,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.LeapFrogAM3.correct")
     def correct(self):
-        r"""
+        r"""
         Correction step from :math:`t_{n}` to :math:`t_{n+1}`
 
         Let :math:`M_n` denote the mass matrix at time :math:`t_{n}`.
@@ -663,7 +650,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.LeapFrogAM3.advance")
     def advance(self, t, update_forcings=None):
-        """Advances equations for one time step."""
+        """Advances equations for one time step."""
         if self._nontrivial:
             if update_forcings is not None:
                 update_forcings(t + self.dt)
@@ -673,7 +660,7 @@ 
    Source code for thetis.timeintegrator
     
 
 
    [docs]class SSPRK22ALE(TimeIntegrator):
-    r"""
+    r"""
     SSPRK(2,2) ALE time integrator for 3D fields
 
     The scheme is
@@ -689,7 +676,7 @@ 
    Source code for thetis.timeintegrator
     
     @PETSc.Log.EventDecorator("thetis.SSPRK22ALE.__init__")
     def __init__(self, equation, solution, fields, dt, options, bnd_conditions, terms_to_add='all'):
-        """
+        """
         :arg equation: equation to solve
         :type equation: :class:`Equation` object
         :arg solution: :class:`Function` where solution will be stored
@@ -726,7 +713,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.SSPRK22ALE.initialize")
     def initialize(self, solution):
-        """Assigns initial conditions to all required fields."""
+        """Assigns initial conditions to all required fields."""
         self.solution.assign(solution)
 
         mass_matrix = assemble(self.a)
@@ -737,7 +724,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.SSPRK22ALE.stage_one_prep")
     def stage_one_prep(self):
-        """
+        """
         Preprocess first stage: compute all forms on the old geometry
         """
         if self._nontrivial:
@@ -753,7 +740,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.SSPRK22ALE.stage_one_solve")
     def stage_one_solve(self):
-        r"""
+        r"""
         First stage: solve :math:`u^{(1)}` given previous solution :math:`u^n`.
 
         This is a forward Euler ALE step between domains :math:`\Omega^n` and :math:`\Omega^{(1)}`:
@@ -772,7 +759,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.SSPRK22ALE.stage_two_prep")
     def stage_two_prep(self):
-        """
+        """
         Preprocess 2nd stage: compute all forms on the old geometry
         """
         if self._nontrivial:
@@ -785,7 +772,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.SSPRK22ALE.stage_two_solve")
     def stage_two_solve(self):
-        r"""
+        r"""
         2nd stage: solve :math:`u^{n+1}` given previous solutions :math:`u^n, u^{(1)}`.
 
         This is an ALE step:
@@ -806,7 +793,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.SSPRK22ALE.solve_stage")
     def solve_stage(self, i_stage):
-        """Solves i-th stage"""
+        """Solves i-th stage"""
         if i_stage == 0:
             self.stage_one_solve()
         else:
@@ -814,7 +801,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.SSPRK22ALE.prepare_stage")
     def prepare_stage(self, i_stage, t, update_forcings=None):
-        """
+        """
         Preprocess stage i_stage.
 
         This must be called prior to updating mesh geometry.
@@ -828,7 +815,7 @@ 
    Source code for thetis.timeintegrator
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.SSPRK22ALE.advance")
     def advance(self, t, update_forcings=None):
-        """Advances equations for one time step."""
+        """Advances equations for one time step."""
         if not self._initialized:
             self.initialize(self.solution)
         for i_stage in range(self.n_stages):
@@ -844,8 +831,8 @@ 
    Source code for thetis.timeintegrator
           
    
     
    
     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/timezone.html b/_modules/thetis/timezone.html
index 8965614..766e46c 100644
--- a/_modules/thetis/timezone.html
+++ b/_modules/thetis/timezone.html
@@ -5,14 +5,13 @@
   
     
     
-    thetis.timezone — Thetis 0+untagged.1636.g63e82af documentation
-    
-    
-    
+    thetis.timezone — Thetis 0+untagged.2009.gd7af522 documentation
+    
+    
+    
     
     
     
-    
     
     
 
   
    
@@ -63,11 +61,11 @@ 
    Source code for thetis.timezone
     
 
 
    [docs]class FixedTimeZone(pytz._FixedOffset):
-    """
+    """
     Class that represents a fixed time zone defined by UTC offset in hours.
     """
     def __init__(self, offset, name):
-        """
+        """
         arg int offset: timezone UTC offset in hours
         arg str name: timezone name
         """
@@ -84,14 +82,14 @@ 
    Source code for thetis.timezone
     
 
 
    [docs]def datetime_to_epoch(t):
-    """
+    """
     Convert python datetime object to epoch time stamp.
     """
     return (t - epoch).total_seconds()
    
 
 
 
    [docs]def epoch_to_datetime(t):
-    """
+    """
     Convert python datetime object to epoch time stamp.
     """
     return epoch + datetime.timedelta(seconds=t)
    
@@ -105,8 +103,8 @@ 
    Source code for thetis.timezone
           
    
     
    
     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/tracer_eq.html b/_modules/thetis/tracer_eq.html
index caeb245..15022df 100644
--- a/_modules/thetis/tracer_eq.html
+++ b/_modules/thetis/tracer_eq.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.tracer_eq — Thetis 0+untagged.1878.g52658ff.dirty documentation
+    thetis.tracer_eq — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -83,7 +83,7 @@ 
    Source code for thetis.tracer_eq
     
 
 
    [docs]class TracerTerm(Term):
-    """
+    """
     Generic tracer term that provides commonly used members and mapping for
     boundary functions.
     """
@@ -92,7 +92,7 @@ 
    Source code for thetis.tracer_eq
                      use_symmetric_surf_bnd=True, use_lax_friedrichs=True,
                  sipg_factor=Constant(1.0),
                  sipg_factor_vertical=Constant(1.0)):
-        """
+        """
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :kwarg bathymetry: bathymetry of the domain
         :type bathymetry: 3D :class:`Function` or :class:`Constant`
@@ -128,7 +128,7 @@ 
    Source code for thetis.tracer_eq
             self.ds_bottom = ds_bottom(degree=self.quad_degree)
 
 
    [docs]    def get_bnd_functions(self, c_in, uv_in, elev_in, bnd_id, bnd_conditions):
-        """
+        """
         Returns external values of tracer and uv for all supported
         boundary conditions.
 
@@ -169,7 +169,7 @@ 
    Source code for thetis.tracer_eq
     
 
 
    [docs]class HorizontalAdvectionTerm(TracerTerm):
-    r"""
+    r"""
     Horizontal advection term :math:`\nabla_h \cdot (\textbf{u} T)`
 
     The weak formulation reads
@@ -243,7 +243,7 @@ 
    Source code for thetis.tracer_eq
     
 
 
    [docs]class VerticalAdvectionTerm(TracerTerm):
-    r"""
+    r"""
     Vertical advection term :math:`\partial (w T)/(\partial z)`
 
     The weak form reads
@@ -292,7 +292,7 @@ 
    Source code for thetis.tracer_eq
     
 
 
    [docs]class HorizontalDiffusionTerm(TracerTerm):
-    r"""
+    r"""
     Horizontal diffusion term :math:`-\nabla_h \cdot (\mu_h \nabla_h T)`
 
     Using the symmetric interior penalty method the weak form becomes
@@ -356,7 +356,7 @@ 
    Source code for thetis.tracer_eq
     
 
 
    [docs]class VerticalDiffusionTerm(TracerTerm):
-    r"""
+    r"""
     Vertical diffusion term :math:`-\frac{\partial}{\partial z} \Big(\mu \frac{T}{\partial z}\Big)`
 
     Using the symmetric interior penalty method the weak form becomes
@@ -412,7 +412,7 @@ 
    Source code for thetis.tracer_eq
     
 
 
    [docs]class SourceTerm(TracerTerm):
-    r"""
+    r"""
     Generic source term
 
     The weak form reads
@@ -431,7 +431,7 @@ 
    Source code for thetis.tracer_eq
     
 
 
    [docs]class TracerEquation(Equation):
-    """
+    """
     3D tracer advection-diffusion equation :eq:`tracer_eq` in conservative form
     """
     def __init__(self, function_space,
@@ -439,7 +439,7 @@ 
    Source code for thetis.tracer_eq
                      use_symmetric_surf_bnd=True, use_lax_friedrichs=True,
                  sipg_factor=Constant(10.0),
                  sipg_factor_vertical=Constant(10.0)):
-        """
+        """
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :kwarg bathymetry: bathymetry of the domain
         :type bathymetry: 3D :class:`Function` or :class:`Constant`
@@ -472,8 +472,8 @@ 
    Source code for thetis.tracer_eq
           
    
     
    
     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/tracer_eq_2d.html b/_modules/thetis/tracer_eq_2d.html
index fffc2ef..9e78b6d 100644
--- a/_modules/thetis/tracer_eq_2d.html
+++ b/_modules/thetis/tracer_eq_2d.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.tracer_eq_2d — Thetis 0+untagged.1888.gc4e776a documentation
+    thetis.tracer_eq_2d — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -62,6 +62,18 @@ 
    Source code for thetis.tracer_eq_2d
         = \nabla_h \cdot (\mu_h \nabla_h T)
     :label: tracer_eq_2d
 
+where :math:`\nabla_h` denotes horizontal gradient, :math:`\textbf{u}` are the horizontal
+velocities, and
+:math:`\mu_h` denotes horizontal diffusivity.
+
+The advection-diffusion equation of depth-integrated tracer :math:`q=HT` in conservative form reads
+
+.. math::
+    \frac{\partial q}{\partial t}
+    + \nabla_h \cdot (\textbf{u} q)
+    = \nabla_h \cdot (\mu_h \nabla_h q)
+    :label: cons_tracer_eq_2d
+
 where :math:`\nabla_h` denotes horizontal gradient, :math:`\textbf{u}` are the horizontal
 velocities, and
 :math:`\mu_h` denotes horizontal diffusivity.
@@ -75,23 +87,34 @@ 
    Source code for thetis.tracer_eq_2d
         'HorizontalAdvectionTerm',
     'HorizontalDiffusionTerm',
     'SourceTerm',
+    'ConservativeTracerTerm',
+    'ConservativeHorizontalAdvectionTerm',
+    'ConservativeHorizontalDiffusionTerm',
+    'ConservativeSourceTerm',
 ]
 
 
 
    [docs]class TracerTerm(Term):
-    """
+    """
     Generic tracer term that provides commonly used members and mapping for
     boundary functions.
     """
-    def __init__(self, function_space, depth, options, test_function=None):
-        """
+    def __init__(self, index, label, function_space, depth, options,
+                 test_function=None, trial_function=None):
+        """
+        :arg index: index for this tracer component within the vector
+        :arg label: label for this tracer component
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :arg depth: :class:`DepthExpression` containing depth info
         :arg options: :class`ModelOptions2d` containing parameters
-        :kwarg test_function: custom :class:`TestFunction`.
+        :kwarg test_function: custom :class:`TestFunction`
+        :kwarg trial_function: custom :class:`TrialFunction`
         """
         super(TracerTerm, self).__init__(function_space,
-                                         test_function=test_function)
+                                         test_function=test_function,
+                                         trial_function=trial_function)
+        self.index = index
+        self.label = label
         self.depth = depth
         self.options = options
         self.cellsize = CellSize(self.mesh)
@@ -106,7 +129,7 @@ 
    Source code for thetis.tracer_eq_2d
             self.ds = ds(degree=self.quad_degree)
 
 
    [docs]    def get_bnd_functions(self, c_in, uv_in, elev_in, bnd_id, bnd_conditions):
-        """
+        """
         Returns external values of tracer and uv for all supported
         boundary conditions.
 
@@ -142,11 +165,17 @@ 
    Source code for thetis.tracer_eq_2d
             else:
             uv_ext = uv_in
 
-        return c_ext, uv_ext, elev_ext
    
+        return c_ext, uv_ext, elev_ext
    
+
+
    [docs]    def component(self, f):
+        """
+        Return the component of a function corresponding to this tracer.
+        """
+        return f if len(f.dof_dset.dim) == 1 else split(f)[self.index]
    
 
 
 
    [docs]class HorizontalAdvectionTerm(TracerTerm):
-    r"""
+    r"""
     Advection of tracer term, :math:`\bar{\textbf{u}} \cdot \nabla T`
 
     The weak form is
@@ -173,14 +202,15 @@ 
    Source code for thetis.tracer_eq_2d
                 return 0
         elev = fields_old['elev_2d']
         self.corr_factor = fields_old.get('tracer_advective_velocity_factor')
+        c = self.component(solution)
 
         uv = self.corr_factor * fields_old['uv_2d']
         # FIXME is this an option?
         lax_friedrichs_factor = fields_old.get('lax_friedrichs_tracer_scaling_factor')
 
         f = 0
-        f += -(Dx(uv[0] * self.test, 0) * solution
-               + Dx(uv[1] * self.test, 1) * solution) * self.dx
+        f += -(Dx(uv[0] * self.test, 0) * c
+               + Dx(uv[1] * self.test, 1) * c) * self.dx
 
         if self.horizontal_dg:
             # add interface term
@@ -188,19 +218,19 @@ 
    Source code for thetis.tracer_eq_2d
                 un_av = (uv_av[0]*self.normal('-')[0]
                      + uv_av[1]*self.normal('-')[1])
             s = 0.5*(sign(un_av) + 1.0)
-            c_up = solution('-')*s + solution('+')*(1-s)
+            c_up = c('-')*s + c('+')*(1-s)
 
             f += c_up*(jump(self.test, uv[0] * self.normal[0])
                        + jump(self.test, uv[1] * self.normal[1])) * self.dS
             # Lax-Friedrichs stabilization
             if self.options.use_lax_friedrichs_tracer:
                 gamma = 0.5*abs(un_av)*lax_friedrichs_factor
-                f += gamma*dot(jump(self.test), jump(solution))*self.dS
+                f += gamma*dot(jump(self.test), jump(c))*self.dS
 
         for bnd_marker in self.boundary_markers:
             funcs = bnd_conditions.get(bnd_marker)
             ds_bnd = ds(int(bnd_marker), degree=self.quad_degree)
-            c_in = solution
+            c_in = c
             if funcs is not None:
                 c_ext, uv_ext, eta_ext = self.get_bnd_functions(c_in, uv, elev, bnd_marker, bnd_conditions)
                 uv_av = 0.5*(uv + uv_ext)
@@ -217,7 +247,7 @@ 
    Source code for thetis.tracer_eq_2d
     
 
 
    [docs]class HorizontalDiffusionTerm(TracerTerm):
-    r"""
+    r"""
     Horizontal diffusion term :math:`-\nabla_h \cdot (\mu_h \nabla_h T)`
 
     Using the symmetric interior penalty method the weak form becomes
@@ -247,13 +277,14 @@ 
    Source code for thetis.tracer_eq_2d
         https://dial.uclouvain.be/pr/boreal/object/boreal:128254/
     """
 
    [docs]    def residual(self, solution, solution_old, fields, fields_old, bnd_conditions):
-        if fields_old.get('diffusivity_h') is None:
+        diffusivity_h = fields_old.get(f'diffusivity_h-{self.label}')
+        if diffusivity_h is None:
             return 0
-        diffusivity_h = fields_old['diffusivity_h']
+        c = self.component(solution)
         diff_tensor = as_matrix([[diffusivity_h, 0, ],
                                  [0, diffusivity_h, ]])
         grad_test = grad(self.test)
-        diff_flux = dot(diff_tensor, grad(solution))
+        diff_flux = dot(diff_tensor, grad(c))
         sipg_factor = self.options.sipg_factor_tracer
 
         f = 0
@@ -272,16 +303,16 @@ 
    Source code for thetis.tracer_eq_2d
                 ds_interior = self.dS
             f += sigma_max * inner(
                 jump(self.test, self.normal),
-                dot(avg(diff_tensor), jump(solution, self.normal)))*ds_interior
+                dot(avg(diff_tensor), jump(c, self.normal)))*ds_interior
             f += -inner(avg(dot(diff_tensor, grad(self.test))),
-                        jump(solution, self.normal))*ds_interior
+                        jump(c, self.normal))*ds_interior
             f += -inner(jump(self.test, self.normal),
-                        avg(dot(diff_tensor, grad(solution))))*ds_interior
+                        avg(dot(diff_tensor, grad(c))))*ds_interior
 
         for bnd_marker in self.boundary_markers:
             funcs = bnd_conditions.get(bnd_marker)
             ds_bnd = ds(int(bnd_marker), degree=self.quad_degree)
-            c_in = solution
+            c_in = c
             elev = fields_old['elev_2d']
             self.corr_factor = fields_old.get('tracer_advective_velocity_factor')
             uv = self.corr_factor * fields_old['uv_2d']
@@ -301,7 +332,7 @@ 
    Source code for thetis.tracer_eq_2d
     
 
 
    [docs]class SourceTerm(TracerTerm):
-    r"""
+    r"""
     Generic source term
 
     The weak form reads
@@ -314,46 +345,213 @@ 
    Source code for thetis.tracer_eq_2d
         """
 
    [docs]    def residual(self, solution, solution_old, fields, fields_old, bnd_conditions):
         f = 0
-        source = fields_old.get('source')
+        source = fields_old.get(f'source-{self.label}')
         if source is not None:
             f += -inner(source, self.test)*self.dx
         return -f
    
 
 
+
    [docs]class ConservativeTracerTerm(TracerTerm):
+    """
+    Generic depth-integrated tracer term that provides commonly used members and mapping for
+    boundary functions.
+    """
+    def __init__(self, index, label, function_space, depth, options,
+                 test_function=None, trial_function=None):
+        """
+        :arg index: index for this tracer component within the vector
+        :arg label: label for this tracer component
+        :arg function_space: :class:`FunctionSpace` where the solution belongs
+        :arg depth: :class: `DepthExpression` containing depth info
+        :arg options: :class`ModelOptions2d` containing parameters
+        :kwarg test_function: custom :class:`TestFunction`.
+        :kwarg trial_function: custom :class:`TrialFunction`.
+        """
+        super(ConservativeTracerTerm, self).__init__(index, label, function_space, depth, options,
+                                                     test_function=test_function,
+                                                     trial_function=trial_function)
    
+
+    # TODO: at the moment this is the same as TracerTerm, but we probably want to overload its
+    # get_bnd_functions method
+
+
+
    [docs]class ConservativeHorizontalAdvectionTerm(ConservativeTracerTerm):
+    r"""
+    Advection of tracer term, :math:`\nabla \cdot \bar{\textbf{u}} \nabla q`
+
+    The weak form is
+
+    .. math::
+        \int_\Omega \boldsymbol{\psi} \nabla\cdot \bar{\textbf{u}} \nabla q  dx
+        = - \int_\Omega \left(\nabla_h \boldsymbol{\psi})\right) \cdot \bar{\textbf{u}} \cdot q dx
+        + \int_\Gamma \text{avg}(q\bar{\textbf{u}}\cdot\textbf{n}) \cdot \text{jump}(\boldsymbol{\psi}) dS
+
+    where the right hand side has been integrated by parts;
+    :math:`\textbf{n}` is the unit normal of
+    the element interfaces, and :math:`\text{jump}` and :math:`\text{avg}` denote the
+    jump and average operators across the interface.
+    """
+
    [docs]    def residual(self, solution, solution_old, fields, fields_old, bnd_conditions=None):
+        if fields_old.get('uv_2d') is None:
+            return 0
+        elev = fields_old['elev_2d']
+        self.corr_factor = fields_old.get('tracer_advective_velocity_factor')
+        c = self.component(solution)
+
+        uv = self.corr_factor * fields_old['uv_2d']
+        uv_p1 = fields_old.get('uv_p1')
+        uv_mag = fields_old.get('uv_mag')
+
+        lax_friedrichs_factor = fields_old.get('lax_friedrichs_tracer_scaling_factor')
+
+        f = 0
+        f += -(Dx(self.test, 0) * uv[0] * c
+               + Dx(self.test, 1) * uv[1] * c) * self.dx
+
+        if self.horizontal_dg:
+            # add interface term
+            uv_av = avg(uv)
+            un_av = (uv_av[0]*self.normal('-')[0]
+                     + uv_av[1]*self.normal('-')[1])
+            s = 0.5*(sign(un_av) + 1.0)
+            flux_up = c('-')*uv('-')*s + c('+')*uv('+')*(1-s)
+
+            f += (flux_up[0] * jump(self.test, self.normal[0])
+                  + flux_up[1] * jump(self.test, self.normal[1])) * self.dS
+            # Lax-Friedrichs stabilization
+            if self.options.use_lax_friedrichs_tracer:
+                if uv_p1 is not None:
+                    gamma = 0.5*abs((avg(uv_p1)[0]*self.normal('-')[0]
+                                     + avg(uv_p1)[1]*self.normal('-')[1]))*lax_friedrichs_factor
+                elif uv_mag is not None:
+                    gamma = 0.5*avg(uv_mag)*lax_friedrichs_factor
+                else:
+                    gamma = 0.5*abs(un_av)*lax_friedrichs_factor
+                f += gamma*dot(jump(self.test), jump(c))*self.dS
+        if bnd_conditions is not None:
+            for bnd_marker in self.boundary_markers:
+                funcs = bnd_conditions.get(bnd_marker)
+                ds_bnd = ds(int(bnd_marker), degree=self.quad_degree)
+                c_in = c
+                if funcs is not None:
+                    c_ext, uv_ext, eta_ext = self.get_bnd_functions(c_in, uv, elev, bnd_marker, bnd_conditions)
+                    uv_av = 0.5*(uv + uv_ext)
+                    un_av = self.normal[0]*uv_av[0] + self.normal[1]*uv_av[1]
+                    s = 0.5*(sign(un_av) + 1.0)
+                    flux_up = c_in*uv*s + c_ext*uv_ext*(1-s)
+                    f += (flux_up[0]*self.normal[0]
+                          + flux_up[1]*self.normal[1])*self.test*ds_bnd
+                else:
+                    f += c_in * (uv[0]*self.normal[0]
+                                 + uv[1]*self.normal[1])*self.test*ds_bnd
+
+        return -f
    
+
+
+
    [docs]class ConservativeHorizontalDiffusionTerm(ConservativeTracerTerm, HorizontalDiffusionTerm):
+    r"""
+    Horizontal diffusion term :math:`-\nabla_h \cdot (\mu_h \nabla_h q)`
+
+    Using the symmetric interior penalty method the weak form becomes
+
+    .. math::
+        -\int_\Omega \nabla_h \cdot (\mu_h \nabla_h q) \phi dx
+        =& \int_\Omega \mu_h (\nabla_h \phi) \cdot (\nabla_h q) dx \\
+        &- \int_{\mathcal{I}_h\cup\mathcal{I}_v} \text{jump}(\phi \textbf{n}_h)
+        \cdot \text{avg}(\mu_h \nabla_h q) dS
+        - \int_{\mathcal{I}_h\cup\mathcal{I}_v} \text{jump}(q \textbf{n}_h)
+        \cdot \text{avg}(\mu_h  \nabla \phi) dS \\
+        &+ \int_{\mathcal{I}_h\cup\mathcal{I}_v} \sigma \text{avg}(\mu_h) \text{jump}(q \textbf{n}_h) \cdot
+            \text{jump}(\phi \textbf{n}_h) dS \\
+        &- \int_\Gamma \mu_h (\nabla_h \phi) \cdot \textbf{n}_h ds
+
+    where :math:`\sigma` is a penalty parameter, see Hillewaert (2013).
+
+    Hillewaert, Koen (2013). Development of the discontinuous Galerkin method
+    for high-resolution, large scale CFD and acoustics in industrial
+    geometries. PhD Thesis. Université catholique de Louvain.
+    https://dial.uclouvain.be/pr/boreal/object/boreal:128254/
+    """
    
+    # TODO: at the moment the same as HorizontalDiffusionTerm
+    # do we need additional H-derivative term?
+    # would also become different if ConservativeTracerTerm gets different bc options
+
+
+
    [docs]class ConservativeSourceTerm(ConservativeTracerTerm):
+    r"""
+    Generic source term
+
+    The weak form reads
+
+    .. math::
+        F_s = \int_\Omega \sigma \phi dx
+
+    where :math:`\sigma` is a user defined scalar :class:`Function`.
+
+    """
+
    [docs]    def residual(self, solution, solution_old, fields, fields_old, bnd_conditions=None):
+        f = 0
+        source = fields_old.get(f'source-{self.label}')
+        if source is not None:
+            H = self.depth.get_total_depth(fields_old['elev_2d'])
+            f += -inner(H*source, self.test)*self.dx
+        return -f
    
+
+
 
    [docs]class TracerEquation2D(Equation):
-    """
-    2D tracer advection-diffusion equation :eq:`tracer_eq` in conservative form
+    """
+    Several 2D tracer advection-diffusion equations :eq:`tracer_eq` in
+    conservative or non-conservative form, solved as a mixed system.
     """
-    def __init__(self, function_space, depth, options, velocity):
-        """
+    def __init__(self, system, function_space, depth, options, velocity):
+        """
+        :arg system: comma separated tracer fields under consideration
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :arg depth: :class: `DepthExpression` containing depth info
         :arg options: :class`ModelOptions2d` containing parameters
         :arg velocity: velocity field associated with the shallow water model
         """
-        super(TracerEquation2D, self).__init__(function_space)
-
-        # Apply SUPG stabilisation
-        kwargs = {}
-        if options.use_supg_tracer:
-            unorm = options.horizontal_velocity_scale
-            if unorm.values()[0] > 0:
-                self.cellsize = anisotropic_cell_size(function_space.mesh())
-                tau = 0.5*self.cellsize/unorm
-                D = options.horizontal_diffusivity_scale
-                if D.values()[0] > 0:
-                    Pe = 0.5*unorm*self.cellsize/D
-                    tau = min_value(tau, Pe/3)
-                self.test = self.test + tau*dot(velocity, grad(self.test))
-                kwargs['test_function'] = self.test
-
-        args = (function_space, depth, options)
-        self.add_terms(*args, **kwargs)
-
-
    [docs]    def add_terms(self, *args, **kwargs):
-        self.add_term(HorizontalAdvectionTerm(*args, **kwargs), 'explicit')
-        self.add_term(HorizontalDiffusionTerm(*args, **kwargs), 'explicit')
-        self.add_term(SourceTerm(*args, **kwargs), 'source')
    
+        super().__init__(function_space)
+        self.depth = depth
+        self.options = options
+        self.velocity = velocity
+        test_functions = TestFunctions(function_space)
+        trial_functions = TrialFunctions(function_space)
+        for i, label in enumerate(system.split(',')):
+            args = (i, label, function_space[i], depth, options)
+            kwargs = {'trial_function': trial_functions[i]}
+            if options.use_supg_tracer:
+                # TODO: allow SUPG for only some fields
+                kwargs['test_function'] = self.apply_supg(test_functions[i])
+            else:
+                kwargs['test_function'] = test_functions[i]
+            if options.tracer[label].use_conservative_form:
+                self.add_conservative_terms(*args, **kwargs)
+            else:
+                self.add_nonconservative_terms(*args, **kwargs)
+
+
    [docs]    def add_nonconservative_terms(self, *args, **kwargs):
+        self.add_term(HorizontalAdvectionTerm(*args, **kwargs), 'explicit', args[1])
+        self.add_term(HorizontalDiffusionTerm(*args, **kwargs), 'explicit', args[1])
+        self.add_term(SourceTerm(*args, **kwargs), 'source', args[1])
    
+
+
    [docs]    def add_conservative_terms(self, *args, **kwargs):
+        self.add_term(ConservativeHorizontalAdvectionTerm(*args, **kwargs), 'explicit', args[1])
+        self.add_term(ConservativeHorizontalDiffusionTerm(*args, **kwargs), 'explicit', args[1])
+        self.add_term(ConservativeSourceTerm(*args, **kwargs), 'source', args[1])
    
+
+
    [docs]    def apply_supg(self, test_function):
+        """
+        Apply SUPG stabilization in the sense of modifying the test function.
+        """
+        unorm = self.options.horizontal_velocity_scale
+        self.cellsize = anisotropic_cell_size(self.function_space.mesh())
+        tau = 0.5*self.cellsize/unorm
+        D = self.options.horizontal_diffusivity_scale
+        Pe = 0.5*unorm*self.cellsize/D
+        tau = conditional(D > 0, min_value(tau, Pe/3), tau)
+        inc = tau*dot(self.velocity, grad(self.test))
+        return test_function + conditional(unorm > 0, inc, 0)
    
 
    
 
             
    
@@ -364,8 +562,8 @@ 
    Source code for thetis.tracer_eq_2d
           
    
     
    
     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/turbines.html b/_modules/thetis/turbines.html
index 54966e9..cb5a82b 100644
--- a/_modules/thetis/turbines.html
+++ b/_modules/thetis/turbines.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.turbines — Thetis 0+untagged.1878.g52658ff.dirty documentation
+    thetis.turbines — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -59,68 +59,158 @@ 
    Source code for thetis.turbines
     from .log import *
 from .callback import DiagnosticCallback
 from .optimisation import DiagnosticOptimisationCallback
+import pyadjoint
 import numpy
 
 
-
    [docs]class TurbineFarm(object):
-    """
-    Evaluates power output, costs and profit of tidal turbine farm
-
-    Cost is simply the number of turbines evaluated as the integral of the
-    turbine density.
-
-    Power output is calculated as the amount of energy
-    taken out of the flow by the turbine drag term. This is in general
-    an over-estimate of the 'usefully extractable' energy. Furthermore no
-    density is included, so that assuming a density of 1000 kg/m^3 and
-    all further quantities in SI, the power is measured in kW.
-
-    Profit is calculated as:
-
-      Profit = Power - break_even_wattage * Cost
-
-    With the above assumptions, break_even_wattage should be specified in
-    kW and can be interpreted as the average power per turbine required
-    to 'break even', i.e. Profit=0.
-
-    Power output and profit are time-integrated (simple first order) and
-    time-averages are available as average_power and average_profit.
+
    [docs]class TidalTurbine:
+    def __init__(self, options, upwind_correction=False):
+        self.diameter = options.diameter
+        self.C_support = options.C_support
+        self.A_support = options.A_support
+        self.upwind_correction = upwind_correction
+
+    def _thrust_area(self, uv):
+        C_T = self.thrust_coefficient(uv)
+        A_T = pi * self.diameter**2 / 4
+        fric = C_T * A_T
+        if self.C_support:
+            fric += self.C_support * self.A_support
+        return fric
+
+
    [docs]    def velocity_correction(self, uv, depth):
+        fric = self._thrust_area(uv)
+        if self.upwind_correction:
+            return 0.5*(1+sqrt(1-fric/(self.diameter*depth)))
+        else:
+            return 1
    
+
+
    [docs]    def friction_coefficient(self, uv, depth):
+        thrust_area = self._thrust_area(uv)
+        alpha = self.velocity_correction(uv, depth)
+        return thrust_area/2./alpha**2
    
+
+
    [docs]    def power(self, uv, depth):
+        # ratio of discrete to upstream velocity (NOTE: should include support drag!)
+        alpha = self.velocity_correction(uv, depth)
+        A_T = pi * self.diameter**2 / 4
+        uv3 = dot(uv, uv)**1.5 / alpha**3  # upwind cubed velocity
+        C_T = self.thrust_coefficient(uv3**(1/3))
+        # this assumes the velocity through the turbine does not change due to the support (is this correct?)
+        return 0.25*C_T*A_T*(1+sqrt(1-C_T))*uv3
    
+
+
+
    [docs]class ConstantThrustTurbine(TidalTurbine):
+    def __init__(self, options, upwind_correction=False):
+        super().__init__(options, upwind_correction=upwind_correction)
+        self.C_T = options.thrust_coefficient
+
+
    [docs]    def thrust_coefficient(self, uv):
+        return self.C_T
    
+
+
+
    [docs]def linearly_interpolate_table(x_list, y_list, y_final, x):
+    """Return UFL expression that linearly interpolates between y-values in x-points
+
+    :param x_list: (1D) x-points
+    :param y_list: y-values in those points
+    :param y_final: value for x>x_list[-1]
+    :param x: point to interpolate (assumed x>x_list[0])
+    """
+    # below x1, interpolate between x0 and x1:
+    below_x1 = ((x_list[1]-x)*y_list[0] + (x-x_list[0])*y_list[1])/(x_list[1]-x_list[0])
+    # above x1, interpolate from rest of the table, or take final value:
+    if len(x_list) > 2:
+        above_x1 = linearly_interpolate_table(x_list[1:], y_list[1:], y_final, x)
+    else:
+        above_x1 = y_final
+
+    return conditional(x < x_list[1], below_x1, above_x1)
    
+
+
+
    [docs]class TabulatedThrustTurbine(TidalTurbine):
+    def __init__(self, options, upwind_correction=False):
+        super().__init__(options, upwind_correction=upwind_correction)
+        self.C_T = options.thrust_coefficients
+        self.speeds = options.thrust_speeds
+        if not len(self.C_T) == len(self.speeds):
+            raise ValueError("In tabulated thrust curve the number of thrust coefficients and speed values should be the same.")
+
+
    [docs]    def thrust_coefficient(self, uv):
+        umag = dot(uv, uv)**0.5
+        return conditional(umag < self.speeds[0], 0, linearly_interpolate_table(self.speeds, self.C_T, 0, umag))
    
+
+
+
    [docs]class TidalTurbineFarm:
+    def __init__(self, turbine_density, dx, options):
+        """
+        :arg turbine_density: turbine distribution density field or expression
+        :arg dx: measure to integrate power output, n/o turbines
+        :arg options: a :class:`TidalTurbineFarmOptions` options dictionary
+        """
+        upwind_correction = getattr(options, 'upwind_correction', False)
+        if options.turbine_type == 'constant':
+            self.turbine = ConstantThrustTurbine(options.turbine_options, upwind_correction=upwind_correction)
+        elif options.turbine_type == 'table':
+            self.turbine = TabulatedThrustTurbine(options.turbine_options, upwind_correction=upwind_correction)
+        self.dx = dx
+        self.turbine_density = turbine_density
+        self.break_even_wattage = options.break_even_wattage
+
+
    [docs]    def number_of_turbines(self):
+        return assemble(self.turbine_density * self.dx)
    
+
+
    [docs]    def power_output(self, uv, depth):
+        return assemble(self.turbine.power(uv, depth) * self.turbine_density * self.dx)
    
+
+
    [docs]    def friction_coefficient(self, uv, depth):
+        return self.turbine.friction_coefficient(uv, depth)
    
+
+
+
    [docs]class DiscreteTidalTurbineFarm(TidalTurbineFarm):
+    """
+    Class that can be used for the addition of turbines in the turbine density field
     """
-    def __init__(self, farm_options, subdomain_id, u, v, dt):
-        """
-        :arg farm_options: a :class:`TidalTurbineFarmOptions` object that define the farm and the turbines used
-        :arg int subdomain_id: the farm is restricted to this subdomain
-        :arg u,v: the depth-averaged velocity field
-        :arg float dt: used for time-integration."""
-        turbine_density = farm_options.turbine_density
-        C_T = farm_options.turbine_options.thrust_coefficient
-        A_T = pi*(farm_options.turbine_options.diameter/2.)**2
-        C_D = C_T*A_T/2.*turbine_density
-        self.power_integral = C_D * (u*u + v*v)**1.5 * dx(subdomain_id)
-        # cost integral is n/o turbines = \int turbine_density
-        self.cost = assemble(turbine_density * dx(subdomain_id))
-        self.break_even_wattage = farm_options.break_even_wattage
-        self.dt = dt
-
-        # time-integrated quantities:
-        self.integrated_power = 0.
-        self.average_power = 0.
-        self.average_profit = 0.
-        self.time_period = 0.
 
-
    [docs]    @PETSc.Log.EventDecorator("thetis.TurbineFarm.evaluate_timestep")
-    def evaluate_timestep(self):
-        """Perform time integration and return current power and time-averaged power and profit."""
-        self.time_period = self.time_period + self.dt
-        current_power = assemble(self.power_integral)
-        self.integrated_power = self.integrated_power + current_power * self.dt
-        self.average_power = self.integrated_power / self.time_period
-        self.average_profit = self.average_power - self.break_even_wattage * self.cost
-        return current_power, self.average_power, self.average_profit
    
+    def __init__(self, mesh, dx, options):
+        """
+        :arg mesh: mesh domain
+        :arg dx: measure to integrate power output, n/o turbines
+        :arg options: a :class:`TidalTurbineFarmOptions` options dictionary
+        """
+
+        # Preliminaries
+        self.mesh = mesh
+        # this sets self.turbine_expr=0
+        super().__init__(0, dx, options)
+
+        # Adding turbine distribution in the domain
+        self.add_turbines(options.turbine_coordinates)
+
+
    [docs]    def add_turbines(self, coordinates):
+        """
+        :param coords: Array with turbine coordinates to be positioned
+        :param function: turbine density function to be adapted
+        :param mesh: computational mesh domain
+        :param radius: radius where the bump will be applied
+        :return: updated turbine density field
+        """
+        x = SpatialCoordinate(self.mesh)
+
+        radius = self.turbine.diameter * 0.5
+        for coord in coordinates:
+            dx0 = (x[0] - coord[0])/radius
+            dx1 = (x[1] - coord[1])/radius
+            psi_x = conditional(lt(abs(dx0), 1), exp(1-1/(1-dx0**2)), 0)
+            psi_y = conditional(lt(abs(dx1), 1), exp(1-1/(1-dx1**2)), 0)
+            bump = psi_x * psi_y
+
+            unit_bump_integral = 1.45661  # integral of bump function for radius=1 (copied from OpenTidalFarm who used Wolfram)
+            self.turbine_density = self.turbine_density + bump/(radius**2 * unit_bump_integral)
    
 
 
 
    [docs]class TurbineFunctionalCallback(DiagnosticCallback):
-    """
+    """
     :class:`.DiagnosticCallback` that evaluates the performance of each tidal turbine farm."""
 
     name = 'turbine'  # this name will be used in the hdf5 file
@@ -128,47 +218,51 @@ 
    Source code for thetis.turbines
     
     @PETSc.Log.EventDecorator("thetis.TurbineFunctionalCallback.__init__")
     def __init__(self, solver_obj, **kwargs):
-        """
+        """
         :arg solver_obj: a :class:`.FlowSolver2d` object containing the tidal_turbine_farms
         :arg kwargs: see :class:`DiagnosticCallback`"""
-        nfarms = len(solver_obj.options.tidal_turbine_farms)
+        if not hasattr(solver_obj, 'tidal_farms'):
+            solver_obj.create_equations()
+        self.farms = solver_obj.tidal_farms
+        nfarms = len(self.farms)
         super().__init__(solver_obj, array_dim=nfarms, **kwargs)
 
-        solver_obj.create_equations()
-        # TODO: was u, eta = split(solution)
-        u, v, eta = solver_obj.fields.solution_2d
-        dt = solver_obj.options.timestep
+        self.uv, eta = split(solver_obj.fields.solution_2d)
+        self.dt = solver_obj.options.timestep
+        self.depth = solver_obj.fields.bathymetry_2d
 
-        self.farms = [TurbineFarm(farm_options, subdomain_id, u, v, dt) for subdomain_id, farm_options in solver_obj.options.tidal_turbine_farms.items()]
-        """The sum of the number of turbines in all farms"""
-        self.cost = sum(farm.cost for farm in self.farms)
+        self.cost = [farm.number_of_turbines() for farm in self.farms]
         if self.append_to_log:
-            print_output('Number of turbines = {}'.format(self.cost))
-
-    def __call__(self):
-        return numpy.transpose([farm.evaluate_timestep() for farm in self.farms])
+            print_output('Number of turbines = {}'.format(sum(self.cost)))
+        self.break_even_wattage = [farm.break_even_wattage for farm in self.farms]
 
-
    [docs]    def message_str(self, current_power, average_power, average_profit):
-        return 'Current power, average power and profit for each farm: {}, {}, {}'.format(current_power, average_power, average_profit)
    
+        # time-integrated quantities:
+        self.integrated_power = [0] * nfarms
+        self.average_power = [0] * nfarms
+        self.average_profit = [0] * nfarms
+        self.time_period = 0.
 
-    @property
-    def average_profit(self):
-        """The sum of the time-averaged profit output of all farms"""
-        return sum(farm.average_profit for farm in self.farms)
+    def _evaluate_timestep(self):
+        """Perform time integration and return current power and time-averaged power and profit."""
+        self.time_period = self.time_period + self.dt
+        current_power = []
+        for i, farm in enumerate(self.farms):
+            power = farm.power_output(self.uv, self.depth)
+            current_power.append(power)
+            self.integrated_power[i] += power * self.dt
+            self.average_power[i] = self.integrated_power[i] / self.time_period
+            self.average_profit[i] = self.average_power[i] - self.break_even_wattage[i] * self.cost[i]
+        return current_power, self.average_power, self.average_profit
 
-    @property
-    def average_power(self):
-        """The sum of the time-averaged power output of all farms"""
-        return sum(farm.average_power for farm in self.farms)
+    def __call__(self):
+        return self._evaluate_timestep()
 
-    @property
-    def integrated_power(self):
-        """The sum of the time-integrated power output of all farms"""
-        return sum(farm.integrated_power for farm in self.farms)
    
+
    [docs]    def message_str(self, current_power, average_power, average_profit):
+        return 'Current power, average power and profit for each farm: {}, {}, {}'.format(current_power, average_power, average_profit)
    
 
 
 
    [docs]class TurbineOptimisationCallback(DiagnosticOptimisationCallback):
-    """
+    """
     :class:`DiagnosticOptimisationCallback` that evaluates the performance of each tidal turbine farm during an optimisation.
 
     See the :py:mod:`optimisation` module for more info about the use of OptimisationCallbacks."""
@@ -176,7 +270,7 @@ 
    Source code for thetis.turbines
         variable_names = ['cost', 'average_power', 'average_profit']
 
     def __init__(self, solver_obj, turbine_functional_callback, **kwargs):
-        """
+        """
         :arg solver_obj: a :class:`.FlowSolver2d` object
         :arg turbine_functional_callback: a :class:`.TurbineFunctionalCallback` used in the forward model
         :args kwargs: see :class:`.DiagnosticOptimisationCallback`"""
@@ -184,13 +278,89 @@ 
    Source code for thetis.turbines
             super().__init__(solver_obj, **kwargs)
 
 
    [docs]    def compute_values(self, *args):
-        costs = [farm.cost.block_variable.saved_output for farm in self.tfc.farms]
-        powers = [farm.average_power.block_variable.saved_output for farm in self.tfc.farms]
-        profits = [farm.average_profit.block_variable.saved_output for farm in self.tfc.farms]
+        costs = [x.block_variable.saved_output for x in self.tfc.cost]
+        powers = [x.block_variable.saved_output for x in self.tfc.average_power]
+        profits = [x.block_variable.saved_output for x in self.tfc.average_profit]
         return costs, powers, profits
    
 
 
    [docs]    def message_str(self, cost, average_power, average_profit):
         return 'Costs, average power and profit for each farm: {}, {}, {}'.format(cost, average_power, average_profit)
    
+
+
+
    [docs]class MinimumDistanceConstraints(pyadjoint.InequalityConstraint):
+    """This class implements minimum distance constraints between turbines.
+
+    .. note:: This class subclasses ``pyadjoint.InequalityConstraint``. The
+        following methods must be implemented:
+
+        * ``length(self)``
+        * ``function(self, m)``
+        * ``jacobian(self, m)``
+    """
+    def __init__(self, turbine_positions, minimum_distance):
+        """Create MinimumDistanceConstraints
+
+        :param turbine_positions: list of [x,y] where x and y are either float or Constant
+        :param minimum_distance: The minimum distance allowed between turbines.
+        """
+        self._turbines = [float(xi) for xy in turbine_positions for xi in xy]
+        self._minimum_distance = minimum_distance
+        self._nturbines = len(turbine_positions)
+
+
    [docs]    def length(self):
+        """Returns the number of constraints ``len(function(m))``."""
+        return int(self._nturbines*(self._nturbines-1)/2)
    
+
+
    [docs]    def function(self, m):
+        """Return an object which must be positive for the point to be feasible.
+
+        :param m: The serialized paramaterisation of the turbines.
+        :type m: numpy.ndarray.
+        :returns: numpy.ndarray -- each entry must be positive for the positions to be
+            feasible.
+        """
+        print_output("Calculating minimum distance constraints.")
+        inequality_constraints = []
+        for i in range(self._nturbines):
+            for j in range(self._nturbines):
+                if i <= j:
+                    continue
+                inequality_constraints.append((m[2*i]-m[2*j])**2 + (m[2*i+1]-m[2*j+1])**2 - self._minimum_distance**2)
+
+        inequality_constraints = numpy.array(inequality_constraints)
+        if any(inequality_constraints <= 0):
+            print_output(
+                "Minimum distance inequality constraints (should all "
+                "be > 0): %s" % inequality_constraints)
+        return inequality_constraints
    
+
+
    [docs]    def jacobian(self, m):
+        """Returns the gradient of the constraint function.
+
+        Return a list of vector-like objects representing the gradient of the
+        constraint function with respect to the parameter m.
+
+        :param m: The serialized paramaterisation of the turbines.
+        :type m: numpy.ndarray.
+        :returns: numpy.ndarray -- the gradient of the constraint function with
+            respect to each input parameter m.
+        """
+        print_output("Calculating gradient of equality constraint")
+
+        grad_h = numpy.zeros((self.length(), self._nturbines*2))
+        row = 0
+        for i in range(self._nturbines):
+            for j in range(self._nturbines):
+                if i <= j:
+                    continue
+
+                grad_h[row, 2*i] = 2*(m[2*i] - m[2*j])
+                grad_h[row, 2*j] = -2*(m[2*i] - m[2*j])
+                grad_h[row, 2*i+1] = 2*(m[2*i+1] - m[2*j+1])
+                grad_h[row, 2*j+1] = -2*(m[2*i+1] - m[2*j+1])
+                row += 1
+
+        return grad_h
    
 
    
 
             
    
@@ -201,8 +371,8 @@ 
    Source code for thetis.turbines
           
    
     
    
     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/turbulence.html b/_modules/thetis/turbulence.html
index d06471b..4144806 100644
--- a/_modules/thetis/turbulence.html
+++ b/_modules/thetis/turbulence.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.turbulence — Thetis 0+untagged.1878.g52658ff.dirty documentation
+    thetis.turbulence — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -198,30 +198,32 @@ 
    Source code for thetis.turbulence
     from .log import *
 from .options import GLSModelOptions, PacanowskiPhilanderModelOptions
 import numpy
+from abc import ABC, abstractmethod
+from pyop2.profiling import timed_stage
 
 
 
    [docs]def set_func_min_val(f, minval):
-    """
+    """
     Sets a minimum value to a :class:`Function`
     """
     f.dat.data[f.dat.data < minval] = minval
    
 
 
 
    [docs]def set_func_max_val(f, maxval):
-    """
+    """
     Sets a minimum value to a :class:`Function`
     """
     f.dat.data[f.dat.data > maxval] = maxval
    
 
 
 
    [docs]class VerticalGradSolver(object):
-    """
+    """
     Computes vertical gradient in the weak sense.
 
     """
     @PETSc.Log.EventDecorator("thetis.VerticalGradSolver.__init__")
     def __init__(self, source, solution, solver_parameters=None):
-        """
+        """
         :arg source: A :class:`Function` or expression to differentiate.
         :arg solution: A :class:`Function` where the solution will be stored.
             Must be in P0 space.
@@ -255,12 +257,12 @@ 
    Source code for thetis.turbulence
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.VerticalGradSolver.solve")
     def solve(self):
-        """Computes the gradient"""
+        """Computes the gradient"""
         self.weak_grad_solver.solve()
    
 
 
 
    [docs]class ShearFrequencySolver(object):
-    r"""
+    r"""
     Computes vertical shear frequency squared form the given horizontal
     velocity field.
 
@@ -270,7 +272,7 @@ 
    Source code for thetis.turbulence
         """
     @PETSc.Log.EventDecorator("thetis.ShearFrequencySolver.__init__")
     def __init__(self, uv, m2, mu, mv, mu_tmp, relaxation=1.0, minval=1e-12):
-        """
+        """
         :arg uv: horizontal velocity field
         :type uv: :class:`Function`
         :arg m2: :math:`M^2` field
@@ -298,7 +300,7 @@ 
    Source code for thetis.turbulence
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.ShearFrequencySolver.solve")
     def solve(self, init_solve=False):
-        """
+        """
         Computes buoyancy frequency
 
         :kwarg bool init_solve: Set to True if solving for the first time, skips
@@ -312,13 +314,13 @@ 
    Source code for thetis.turbulence
                     gamma = self.relaxation if not init_solve else 1.0
                 mu_comp[i_comp].assign(gamma*self.mu_tmp
                                        + (1.0 - gamma)*mu_comp[i_comp])
-                self.m2 += mu_comp[i_comp]*mu_comp[i_comp]
+                self.m2.interpolate(self.m2 + mu_comp[i_comp]*mu_comp[i_comp])
             # crop small/negative values
             set_func_min_val(self.m2, self.minval)
    
 
 
 
    [docs]class BuoyFrequencySolver(object):
-    r"""
+    r"""
     Computes buoyancy frequency squared form the given horizontal
     velocity field.
 
@@ -327,7 +329,7 @@ 
    Source code for thetis.turbulence
         """
     @PETSc.Log.EventDecorator("thetis.BuoyFrequencySolver.__init__")
     def __init__(self, rho, n2, n2_tmp, relaxation=1.0, minval=1e-12):
-        """
+        """
         :arg rho: water density field
         :type rho: :class:`Function`
         :arg n2: :math:`N^2` field
@@ -356,7 +358,7 @@ 
    Source code for thetis.turbulence
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.BuoyFrequencySolver.solve")
     def solve(self, init_solve=False):
-        """
+        """
         Computes buoyancy frequency
 
         :kwarg bool init_solve: Set to True if solving for the first time, skips
@@ -370,17 +372,17 @@ 
    Source code for thetis.turbulence
                                    + (1.0 - gamma)*self.n2)
    
 
 
-
    [docs]class TurbulenceModel(object):
-    """Base class for all vertical turbulence models"""
+
    [docs]class TurbulenceModel(ABC):
+    """Base class for all vertical turbulence models"""
 
 
    [docs]    @abstractmethod
     def initialize(self):
-        """Initialize all turbulence fields"""
+        """Initialize all turbulence fields"""
         pass
    
 
 
    [docs]    @abstractmethod
     def preprocess(self, init_solve=False):
-        """
+        """
         Computes all diagnostic variables that depend on the mean flow model
         variables.
 
@@ -390,7 +392,7 @@ 
    Source code for thetis.turbulence
     
 
    [docs]    @abstractmethod
     def postprocess(self):
-        """
+        """
         Updates all diagnostic variables that depend on the turbulence state
         variables.
 
@@ -400,7 +402,7 @@ 
    Source code for thetis.turbulence
     
 
 
    [docs]class GenericLengthScaleModel(TurbulenceModel):
-    """
+    """
     Generic Length Scale turbulence closure model implementation
     """
     @PETSc.Log.EventDecorator("thetis.GenericLengthScaleModel.__init__")
@@ -408,7 +410,7 @@ 
    Source code for thetis.turbulence
                      l_field, epsilon_field,
                  eddy_diffusivity, eddy_viscosity,
                  n2, m2, options=None):
-        """
+        """
         :arg solver: FlowSolver object
         :arg k_field: turbulent kinetic energy (TKE) field
         :type k_field: :class:`Function`
@@ -512,7 +514,7 @@ 
    Source code for thetis.turbulence
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.GenericLengthScaleModel.initialize")
     def initialize(self, l_init=0.01):
-        """
+        """
         Initialize turbulent fields.
 
         k is set to k_min value. epsilon and psi are set to a value that
@@ -536,7 +538,7 @@ 
    Source code for thetis.turbulence
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.GenericLengthScaleModel.preprocess")
     def preprocess(self, init_solve=False):
-        """
+        """
         Computes all diagnostic variables that depend on the mean flow model
         variables.
 
@@ -559,7 +561,7 @@ 
    Source code for thetis.turbulence
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.GenericLengthScaleModel.postprocess")
     def postprocess(self):
-        r"""
+        r"""
         Updates all diagnostic variables that depend on the turbulence state
         variables :math:`k,\psi`.
 
@@ -592,7 +594,7 @@ 
    Source code for thetis.turbulence
                 set_func_min_val(self.psi, o.psi_min)
 
             # udpate epsilon
-            self.epsilon.assign(cmu0**(3.0 + p/n)*self.k**(3.0/2.0 + m/n)*self.psi**(-1.0/n))
+            self.epsilon.interpolate(cmu0**(3.0 + p/n)*self.k**(3.0/2.0 + m/n)*self.psi**(-1.0/n))
             if o.limit_eps:
                 # impose Galperin limit on eps
                 eps_min = cmu0**3.0/(numpy.sqrt(2)*galp_clim)*numpy.sqrt(n2_pos)*k_arr
@@ -601,7 +603,7 @@ 
    Source code for thetis.turbulence
                 set_func_min_val(self.epsilon, o.eps_min)
 
             # update L
-            self.l.assign(cmu0**3.0 * self.k**(3.0/2.0) / self.epsilon)
+            self.l.interpolate(cmu0**3.0 * self.k**(3.0/2.0) / self.epsilon)
             if o.limit_len_min:
                 set_func_min_val(self.l, o.len_min)
             if o.limit_len:
@@ -628,7 +630,7 @@ 
    Source code for thetis.turbulence
                 set_func_min_val(self.diffusivity, o.diff_min)
    
 
 
    [docs]    def print_debug(self):
-        """
+        """
         Print diagnostic field values for debugging.
         """
         fmt = '{:8s} {:10.3e} {:10.3e}'
@@ -648,7 +650,7 @@ 
    Source code for thetis.turbulence
     
 
 
    [docs]class TKESourceTerm(TracerTerm):
-    r"""
+    r"""
     Production and destruction terms of the TKE equation :eq:`turb_tke_eq`
 
     .. math::
@@ -665,7 +667,7 @@ 
    Source code for thetis.turbulence
         """
     def __init__(self, function_space, gls_model,
                  bathymetry=None, v_elem_size=None, h_elem_size=None):
-        """
+        """
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :arg gls_model: :class:`.GenericLengthScaleModel` object
         :kwarg bathymetry: bathymetry of the domain
@@ -703,7 +705,7 @@ 
    Source code for thetis.turbulence
     
 
 
    [docs]class PsiSourceTerm(TracerTerm):
-    r"""
+    r"""
     Production and destruction terms of the Psi equation :eq:`turb_psi_eq`
 
     .. math::
@@ -730,7 +732,7 @@ 
    Source code for thetis.turbulence
         """
     def __init__(self, function_space, gls_model,
                  bathymetry=None, v_elem_size=None, h_elem_size=None):
-        """
+        """
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :arg gls_model: :class:`.GenericLengthScaleModel` object
         :kwarg bathymetry: bathymetry of the domain
@@ -810,14 +812,14 @@ 
    Source code for thetis.turbulence
     
 
 
    [docs]class GLSVerticalDiffusionTerm(VerticalDiffusionTerm):
-    """
+    """
     Vertical diffusion term where the diffusivity is replaced by
     viscosity/Schmidt number.
     """
     def __init__(self, function_space, schmidt_nb,
                  bathymetry=None, v_elem_size=None, h_elem_size=None,
                  sipg_factor=Constant(1.0)):
-        """
+        """
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :arg schmidt_nb: the Schmidt number of TKE or Psi
         :kwarg bathymetry: bathymetry of the domain
@@ -842,14 +844,14 @@ 
    Source code for thetis.turbulence
     
 
 
    [docs]class TKEEquation(Equation):
-    """
+    """
     Turbulent kinetic energy equation :eq:`turb_tke_eq` without advection terms.
 
     Advection of TKE is implemented using the standard tracer equation.
     """
     def __init__(self, function_space, gls_model,
                  bathymetry=None, v_elem_size=None, h_elem_size=None):
-        """
+        """
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :arg gls_model: :class:`.GenericLengthScaleModel` object
         :kwarg bathymetry: bathymetry of the domain
@@ -872,14 +874,14 @@ 
    Source code for thetis.turbulence
     
 
 
    [docs]class PsiEquation(Equation):
-    r"""
+    r"""
     Generic length scale equation :eq:`turb_psi_eq` without advection terms.
 
     Advection of :math:`\psi` is implemented using the standard tracer equation.
     """
     def __init__(self, function_space, gls_model,
                  bathymetry=None, v_elem_size=None, h_elem_size=None):
-        """
+        """
         :arg function_space: :class:`FunctionSpace` where the solution belongs
         :arg gls_model: :class:`.GenericLengthScaleModel` object
         :kwarg bathymetry: bathymetry of the domain
@@ -902,7 +904,7 @@ 
    Source code for thetis.turbulence
     
 
 
    [docs]class PacanowskiPhilanderModel(TurbulenceModel):
-    r"""
+    r"""
     Gradient Richardson number based model by Pacanowski and Philander (1981).
 
     Given the gradient Richardson number :math:`Ri` the eddy viscosity and
@@ -924,7 +926,7 @@ 
    Source code for thetis.turbulence
         def __init__(self, solver, uv_field, rho_field,
                  eddy_diffusivity, eddy_viscosity,
                  n2, m2, options=None):
-        """
+        """
         :arg solver: FlowSolver object
         :arg uv_field: horizontal velocity field
         :type uv_field: :class:`Function`
@@ -974,7 +976,7 @@ 
    Source code for thetis.turbulence
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.PacanowskiPhilanderModel.initialize")
     def initialize(self):
-        """Initializes fields"""
+        """Initializes fields"""
         self.n2.assign(1e-12)
         self.n2_pos.assign(1e-12)
         self.preprocess(init_solve=True)
@@ -982,7 +984,7 @@ 
    Source code for thetis.turbulence
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.PacanowskiPhilanderModel.preprocess")
     def preprocess(self, init_solve=False):
-        """
+        """
         Computes all diagnostic variables that depend on the mean flow model
         variables.
 
@@ -1000,7 +1002,7 @@ 
    Source code for thetis.turbulence
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.PacanowskiPhilanderModel.postprocess")
     def postprocess(self):
-        """
+        """
         Updates all diagnostic variables that depend on the turbulence state
         variables.
 
@@ -1020,8 +1022,8 @@ 
    Source code for thetis.turbulence
           
    
     
    
     
   
 
\ No newline at end of file
diff --git a/_modules/thetis/utility.html b/_modules/thetis/utility.html
index effda12..a08bcb7 100644
--- a/_modules/thetis/utility.html
+++ b/_modules/thetis/utility.html
@@ -5,7 +5,7 @@
   
     
     
-    thetis.utility — Thetis 0+untagged.1888.gb6666e3 documentation
+    thetis.utility — Thetis 0+untagged.2009.gd7af522 documentation
     
     
     
@@ -55,16 +55,15 @@ 
    Source code for thetis.utility
     Utility functions and classes for 2D and 3D ocean models
 """
 import os
+import glob
 import sys
-from abc import ABCMeta, abstractmethod  # NOQA
 from collections import OrderedDict, namedtuple  # NOQA
 
-import coffee.base as ast  # NOQA
 import ufl  # NOQA
 from firedrake import *
 from firedrake.petsc import PETSc
 from mpi4py import MPI  # NOQA
-from pyop2.profiling import timed_function, timed_region, timed_stage  # NOQA
+from pyop2.profiling import timed_stage  # NOQA
 import numpy
 from functools import wraps
 from pyadjoint.tape import no_annotations
@@ -78,7 +77,7 @@ 
    Source code for thetis.utility
     
 
 
    [docs]class FrozenClass(object):
-    """
+    """
     A class where creating a new attribute will raise an exception if
     :attr:`_isfrozen` is ``True``.
 
@@ -95,7 +94,7 @@ 
    Source code for thetis.utility
     
 
 
    [docs]def unfrozen(method):
-    """
+    """
     Decorator to temporarily unfreeze an object
     whilst one of its methods is being called.
     """
@@ -112,11 +111,11 @@ 
    Source code for thetis.utility
     
 
 
    [docs]class SumFunction(object):
-    """
+    """
     Helper class to keep track of sum of Coefficients.
     """
     def __init__(self):
-        """
+        """
         Initialize empty sum.
 
         get operation returns Constant(0)
@@ -124,7 +123,7 @@ 
    Source code for thetis.utility
             self.coeff_list = []
 
 
    [docs]    def add(self, coeff):
-        """
+        """
         Adds a coefficient to self
         """
         if coeff is None:
@@ -132,7 +131,7 @@ 
    Source code for thetis.utility
             self.coeff_list.append(coeff)
    
 
 
    [docs]    def get_sum(self):
-        """
+        """
         Returns a sum of all added Coefficients
         """
         if len(self.coeff_list) == 0:
@@ -141,7 +140,7 @@ 
    Source code for thetis.utility
     
 
 
    [docs]class AttrDict(dict):
-    """
+    """
     Dictionary that provides both self['key'] and self.key access to members.
 
     http://stackoverflow.com/questions/4984647/accessing-dict-keys-like-an-attribute-in-python
@@ -154,7 +153,7 @@ 
    Source code for thetis.utility
     
 
 
    [docs]class FieldDict(AttrDict):
-    """
+    """
     AttrDict that checks that all added fields have proper meta data.
 
     Values can be either Function or Constant objects.
@@ -175,7 +174,7 @@ 
    Source code for thetis.utility
                     raise Exception(msg)
 
     def _set_functionname(self, key, value):
-        """Set function.name to key to ensure consistent naming"""
+        """Set function.name to key to ensure consistent naming"""
         if isinstance(value, Function):
             value.rename(name=key)
 
@@ -229,18 +228,18 @@ 
    Source code for thetis.utility
     
 
 
    [docs]def get_extruded_base_element(ufl_element):
-    """
+    """
     Return UFL TensorProductElement of an extruded UFL element.
 
     In case of a non-extruded mesh, returns the element itself.
     """
-    if isinstance(ufl_element, ufl.HDivElement):
+    if isinstance(ufl_element, firedrake.HDivElement):
         ufl_element = ufl_element._element
-    if isinstance(ufl_element, ufl.MixedElement):
-        ufl_element = ufl_element.sub_elements()[0]
-    if isinstance(ufl_element, ufl.VectorElement):
-        ufl_element = ufl_element.sub_elements()[0]  # take the first component
-    if isinstance(ufl_element, ufl.EnrichedElement):
+    if isinstance(ufl_element, firedrake.MixedElement):
+        ufl_element = ufl_element.sub_elements[0]
+    if isinstance(ufl_element, firedrake.VectorElement):
+        ufl_element = ufl_element.sub_elements[0]  # take the first component
+    if isinstance(ufl_element, firedrake.EnrichedElement):
         ufl_element = ufl_element._elements[0]
     return ufl_element
    
 
@@ -254,7 +253,7 @@ 
    Source code for thetis.utility
     
 
 
    [docs]def element_continuity(ufl_element):
-    """Return an :class:`ElementContinuity` instance with the
+    """Return an :class:`ElementContinuity` instance with the
     continuity of a given element.
 
     :arg ufl_element: The UFL element to determine the continuity
@@ -274,11 +273,11 @@ 
    Source code for thetis.utility
         }
 
     base_element = get_extruded_base_element(ufl_element)
-    if isinstance(elem, ufl.HDivElement):
+    if isinstance(elem, firedrake.HDivElement):
         horiz_type = 'hdiv'
         vert_type = 'hdiv'
-    elif isinstance(base_element, ufl.TensorProductElement):
-        a, b = base_element.sub_elements()
+    elif isinstance(base_element, firedrake.TensorProductElement):
+        a, b = base_element.sub_elements
         horiz_type = elem_types[a.family()]
         vert_type = elem_types[b.family()]
     else:
@@ -288,7 +287,7 @@ 
    Source code for thetis.utility
     
 
 
    [docs]def create_directory(path, comm=COMM_WORLD):
-    """
+    """
     Create a directory on disk
 
     Raises IOError if a file with the same name already exists.
@@ -303,9 +302,51 @@ 
    Source code for thetis.utility
         return path
    
 
 
+
    [docs]def read_mesh_from_checkpoint(filename_or_outputdir, mesh_name=None):
+    """
+    Read mesh from a hdf5 checkpoint file.
+    :arg filename_or_outputdir: file name of the hdf5 checkpoint file,
+       or the output directory used in the previous run that create the
+       checkpoint. In the latter case the mesh will be read from the
+       first file that's found in outputdir/hdf5/*.h5
+
+    When loading fields from a checkpoint file in a run, the mesh used in that
+    run needs to be read from that same checkpoint file.
+
+    When multiple checkpoint files are used as model inputs which have
+    been created in separate runs/scripts, make sure that only one of those
+    scripts created the original mesh (by reading a .msh file or using a mesh
+    creation utility like RectangleMesh) and all other script read their mesh
+    from the checkpoint created by that first script.  For example, a
+    preprocessing script might read in a .msh file and interpolate and smoothen
+    the bathymetry on that mesh and write out the result in a checkpoint. A
+    first Thetis run should read its mesh from that checkpoint, and may then
+    save a series of hdf5 files.  A second Thetis run can then read in the mesh
+    and bathymetry from the checkpoint of the preprocessing script, and call
+    load_state() to pick up from the previous run.
+    """
+    if os.path.isdir(filename_or_outputdir):
+        # grab the first outputdir/hdf5/*.h5 file we can find
+        path = os.path.join(filename_or_outputdir, 'hdf5', '*.h5')
+        try:
+            filename = next(glob.iglob(path))
+        except StopIteration:
+            raise IOError(f'No checkpoint files found in {path}')
+    else:
+        # assume we're being pointed to the hdf5 file directly
+        filename = filename_or_outputdir
+
+    with CheckpointFile(filename, 'r') as f:
+        if mesh_name is None:
+            mesh_name = 'firedrake_default'
+        mesh = f.load_mesh(mesh_name)
+
+    return mesh
    
+
+
 
    [docs]@PETSc.Log.EventDecorator("thetis.get_facet_mask")
 def get_facet_mask(function_space, facet='bottom'):
-    """
+    """
     Returns the top/bottom nodes of extruded 3D elements.
 
     :arg function_space: Firedrake :class:`FunctionSpace` object
@@ -323,7 +364,7 @@ 
    Source code for thetis.utility
         assert isinstance(elem, TensorProductElement), \
         f'function space must be defined on an extruded 3D mesh: {elem}'
     # figure out number of nodes in sub elements
-    h_elt, v_elt = elem.sub_elements()
+    h_elt, v_elt = elem.sub_elements
     nb_nodes_h = create_finat_element(h_elt).space_dimension()
     nb_nodes_v = create_finat_element(v_elt).space_dimension()
     # compute top/bottom facet indices
@@ -337,7 +378,7 @@ 
    Source code for thetis.utility
     
    [docs]@PETSc.Log.EventDecorator("thetis.extrude_mesh_sigma")
 def extrude_mesh_sigma(mesh2d, n_layers, bathymetry_2d, z_stretch_fact=1.0,
                        min_depth=None):
-    """
+    """
     Extrudes a 2d surface mesh with bathymetry data defined in a 2d field.
 
     Generates a uniform terrain following mesh.
@@ -408,14 +449,14 @@ 
    Source code for thetis.utility
     
 
    [docs]@PETSc.Log.EventDecorator("thetis.comp_volume_2d")
 def comp_volume_2d(eta, bath):
-    """Computes volume of the 2D domain as an integral of the elevation field"""
+    """Computes volume of the 2D domain as an integral of the elevation field"""
     val = assemble((eta+bath)*dx)
     return val
    
 
 
 
    [docs]@PETSc.Log.EventDecorator("thetis.comp_volume_3d")
 def comp_volume_3d(mesh):
-    """Computes volume of the 3D domain as an integral"""
+    """Computes volume of the 3D domain as an integral"""
     one = Constant(1.0, domain=mesh.coordinates.ufl_domain())
     val = assemble(one*dx)
     return val
    
@@ -423,7 +464,7 @@ 
    Source code for thetis.utility
     
 
    [docs]@PETSc.Log.EventDecorator("thetis.comp_tracer_mass_2d")
 def comp_tracer_mass_2d(scalar_func, total_depth):
-    """
+    """
     Computes total tracer mass in the 2D domain
     :arg scalar_func: depth-averaged scalar :class:`Function` to integrate
     :arg total_depth: scalar UFL expression (e.g. from get_total_depth())
@@ -434,7 +475,7 @@ 
    Source code for thetis.utility
     
 
    [docs]@PETSc.Log.EventDecorator("thetis.comp_tracer_mass_3d")
 def comp_tracer_mass_3d(scalar_func):
-    """
+    """
     Computes total tracer mass in the 3D domain
 
     :arg scalar_func: scalar :class:`Function` to integrate
@@ -445,7 +486,7 @@ 
    Source code for thetis.utility
     
 
    [docs]@PETSc.Log.EventDecorator("thetis.get_zcoord_from_mesh")
 def get_zcoord_from_mesh(zcoord, solver_parameters=None):
-    """
+    """
     Evaluates z coordinates from the 3D mesh
 
     :arg zcoord: scalar :class:`Function` where coordinates will be stored
@@ -466,7 +507,7 @@ 
    Source code for thetis.utility
     
 
    [docs]@PETSc.Log.EventDecorator("thetis.compute_baroclinic_head")
 def compute_baroclinic_head(solver):
-    r"""
+    r"""
     Computes the baroclinic head :math:`r` from the density field
 
     .. math::
@@ -483,7 +524,7 @@ 
    Source code for thetis.utility
     
 
    [docs]@PETSc.Log.EventDecorator("thetis.extend_function_to_3d")
 def extend_function_to_3d(func, mesh_extruded):
-    """
+    """
     Returns a 3D view of a 2D :class:`Function` on the extruded domain.
 
     The 3D function resides in V x R function space, where V is the function
@@ -496,7 +537,7 @@ 
    Source code for thetis.utility
         family = ufl_elem.family()
     degree = ufl_elem.degree()
     name = func.name()
-    if isinstance(ufl_elem, ufl.VectorElement):
+    if isinstance(ufl_elem, firedrake.VectorElement):
         # vector function space
         fs_extended = get_functionspace(mesh_extruded, family, degree, 'R', 0,
                                         dim=2, vector=True)
@@ -508,7 +549,7 @@ 
    Source code for thetis.utility
     
 
 
    [docs]class ExtrudedFunction(Function):
-    """
+    """
     A 2D :class:`Function` that provides a 3D view on the extruded domain.
 
     The 3D function can be accessed as `ExtrudedFunction.view_3d`.
@@ -517,7 +558,7 @@ 
    Source code for thetis.utility
         function.
     """
     def __init__(self, *args, mesh_3d=None, **kwargs):
-        """
+        """
         Create a 2D :class:`Function` with a 3D view on extruded mesh.
 
         :arg mesh_3d: Extruded 3D mesh where the function will be extended to.
@@ -530,7 +571,7 @@ 
    Source code for thetis.utility
     
 
 
    [docs]class SubdomainProjector(object):
-    """
+    """
     Projector that projects the restriction of an expression to the specified subdomain.
     """
     def __init__(self, v, v_out, subdomain_id, solver_parameters=None, constant_jacobian=True):
@@ -562,7 +603,7 @@ 
    Source code for thetis.utility
     
 
    [docs]    @PETSc.Log.EventDecorator("thetis.SubdomainProjector.project")
     def project(self):
-        """
+        """
         Apply the projection.
         """
         self.solver.solve()
    
@@ -570,7 +611,7 @@ 
    Source code for thetis.utility
     
 
    [docs]@PETSc.Log.EventDecorator("thetis.compute_elem_height")
 def compute_elem_height(zcoord, output):
-    """
+    """
     Computes the element height on an extruded mesh.
 
     :arg zcoord: field that contains the z coordinates of the mesh
@@ -609,7 +650,7 @@ 
    Source code for thetis.utility
     
    [docs]@no_annotations
 @PETSc.Log.EventDecorator("thetis.get_horizontal_elem_size_2d")
 def get_horizontal_elem_size_2d(sol2d):
-    """
+    """
     Computes horizontal element size from the 2D mesh
 
     :arg sol2d: 2D :class:`Function` where result is stored
@@ -631,7 +672,7 @@ 
    Source code for thetis.utility
     
 
    [docs]@PETSc.Log.EventDecorator("thetis.get_facet_areas")
 def get_facet_areas(mesh):
-    """
+    """
     Compute area of each facet of `mesh`. The facet areas are stored as a HDiv trace field.
 
     NOTES:
@@ -654,7 +695,7 @@ 
    Source code for thetis.utility
     
 
    [docs]@PETSc.Log.EventDecorator("thetis.get_minimum_angles_2d")
 def get_minimum_angles_2d(mesh2d):
-    """
+    """
     Compute the minimum angle in each element of a triangular mesh, `mesh2d`, using the
     cosine rule. The minimum angles are outputted as a P0 field.
     """
@@ -665,8 +706,12 @@ 
    Source code for thetis.utility
             raise NotImplementedError("Minimum angle only currently implemented for triangles.")
     edge_lengths = get_facet_areas(mesh2d)
     min_angles = Function(FunctionSpace(mesh2d, "DG", 0))
-    par_loop("""for (int i=0; i<angle.dofs; i++) {
+    edge_cell_node_map = edge_lengths.function_space().cell_node_map()
+    min_angle_cell_node_map = min_angles.function_space().cell_node_map()
 
+    kernel = op2.Kernel("""
+        void minimum_angle_kernel(double *edges, double *angle) {
+            for (int i=0; i<%(nodes)d; i++) {
                   double min_edge = edges[0];
                   int min_index = 0;
 
@@ -689,13 +734,17 @@ 
    Source code for thetis.utility
                         }
                   }
                   angle[0] = acos(numerator/denominator);
-                }""", dx, {'edges': (edge_lengths, READ), 'angle': (min_angles, RW)})
+            }
+        }""" % {"nodes": edge_cell_node_map.arity}, "minimum_angle_kernel")
+    op2.par_loop(kernel, mesh2d.cell_set,
+                 edge_lengths.dat(op2.READ, edge_cell_node_map),
+                 min_angles.dat(op2.RW, min_angle_cell_node_map))
     return min_angles
    
 
 
 
    [docs]@PETSc.Log.EventDecorator("thetis.get_cell_widths_2d")
 def get_cell_widths_2d(mesh2d):
-    """
+    """
     Compute widths of mesh elements in each coordinate direction as the maximum distance
     between components of vertex coordinates.
     """
@@ -705,16 +754,24 @@ 
    Source code for thetis.utility
         except AssertionError:
         raise NotImplementedError("Cell widths only currently implemented for triangles.")
     cell_widths = Function(VectorFunctionSpace(mesh2d, "DG", 0)).assign(numpy.finfo(0.0).min)
-    par_loop("""for (int i=0; i<coords.dofs; i++) {
-                  widths[0] = fmax(widths[0], fabs(coords[2*i] - coords[(2*i+2)%6]));
-                  widths[1] = fmax(widths[1], fabs(coords[2*i+1] - coords[(2*i+3)%6]));
-                }""", dx, {'coords': (mesh2d.coordinates, READ), 'widths': (cell_widths, RW)})
+    coords_cell_node_map = mesh2d.coordinates.function_space().cell_node_map()
+    widths_cell_node_map = cell_widths.function_space().cell_node_map()
+    kernel = op2.Kernel("""
+        void cell_width_kernel(double *coords, double *widths) {
+            for (int i=0; i<%(nodes)d; i++) {
+                  widths[0] = fmax(widths[0], fabs(coords[2*i] - coords[(2*i+2)%%6]));
+                  widths[1] = fmax(widths[1], fabs(coords[2*i+1] - coords[(2*i+3)%%6]));
+            }
+        }""" % {"nodes": coords_cell_node_map.arity}, "cell_width_kernel")
+    op2.par_loop(kernel, mesh2d.cell_set,
+                 mesh2d.coordinates.dat(op2.READ, coords_cell_node_map),
+                 cell_widths.dat(op2.MAX, widths_cell_node_map))
     return cell_widths
    
 
 
 
    [docs]@PETSc.Log.EventDecorator("thetis.anisotropic_cell_size")
 def anisotropic_cell_size(mesh):
-    """
+    """
     Measure of cell size for anisotropic meshes, as described in
     Micheletti et al. (2003).
 
@@ -725,52 +782,23 @@ 
    Source code for thetis.utility
         the advection-diffusion and the Stokes problems. SIAM Journal
     on Numerical Analysis 41.3: 1131-1162.
     """
-    try:
-        from firedrake.slate.slac.compiler import PETSC_ARCH
-    except ImportError:
-        PETSC_ARCH = os.path.join(os.environ.get('PETSC_DIR'), os.environ.get('PETSC_ARCH'))
-    include_dir = ["%s/include/eigen3" % PETSC_ARCH]
-
-    # Compute cell Jacobian
-    P0_ten = TensorFunctionSpace(mesh, "DG", 0)
-    J = Function(P0_ten, name="Cell Jacobian")
-    J.interpolate(Jacobian(mesh))
 
-    # Compute minimum eigenvalue
-    P0 = FunctionSpace(mesh, "DG", 0)
-    min_evalue = Function(P0, name="Minimum eigenvalue")
-    kernel_str = """
-#include <Eigen/Dense>
+    JTJ = dot(Jacobian(mesh).T, Jacobian(mesh))
 
-using namespace Eigen;
+    # based on https://github.com/michalhabera/dolfiny/blob/master/dolfiny/invariants.py
+    I1 = JTJ[0, 0] + JTJ[1, 1]
+    D = (JTJ[0, 0] - JTJ[1, 1])**2 + 4 * JTJ[0, 1] * JTJ[1, 0]
+    D += numpy.finfo(float).eps
 
-void eigmin(double minEval[1], const double * J_) {
-
-  // Map input onto an Eigen object
-  Map<Matrix<double, 2, 2, RowMajor> > J((double *)J_);
-
-  // Compute J^T * J
-  Matrix<double, 2, 2, RowMajor> A = J.transpose()*J;
-
-  // Solve eigenvalue problem
-  SelfAdjointEigenSolver<Matrix<double, 2, 2, RowMajor>> eigensolver(A);
-  Vector2d D = eigensolver.eigenvalues();
-
-  // Take the square root
-  double lambda1 = sqrt(fabs(D(0)));
-  double lambda2 = sqrt(fabs(D(1)));
-
-  // Select minimum eigenvalue in modulus
-  minEval[0] = fmin(lambda1, lambda2);
-}
-"""
-    kernel = op2.Kernel(kernel_str, 'eigmin', cpp=True, include_dirs=include_dir)
-    op2.par_loop(kernel, P0_ten.node_set, min_evalue.dat(op2.RW), J.dat(op2.READ))
-    return min_evalue
    
+    # Compute smallest singular value of J
+    P0 = FunctionSpace(mesh, "DG", 0)
+    cell_size = Function(P0, name="Smallest singular value")
+    cell_size.interpolate(sqrt((I1-sqrt(D))/2.))
+    return cell_size
    
 
 
 
    [docs]def beta_plane_coriolis_params(latitude):
-    r"""
+    r"""
     Computes beta plane parameters :math:`f_0,\beta` based on latitude
 
     :arg float latitude: latitude in degrees
@@ -792,7 +820,7 @@ 
    Source code for thetis.utility
     
 
 
    [docs]def beta_plane_coriolis_function(latitude, out_function, y_offset=0.0):
-    """
+    """
     Interpolates beta plane Coriolis function to a field
 
     :arg float latitude: latitude in degrees
@@ -807,7 +835,7 @@ 
    Source code for thetis.utility
     
 
 
    [docs]def tensor_jump(v, n):
-    r"""
+    r"""
     Jump term for vector functions based on the tensor product
 
     .. math::
@@ -821,7 +849,7 @@ 
    Source code for thetis.utility
     
 
 
    [docs]def compute_boundary_length(mesh2d):
-    """
+    """
     Computes the length of the boundary segments in given 2d mesh
     """
     p1 = get_functionspace(mesh2d, 'CG', 1)
@@ -836,7 +864,7 @@ 
    Source code for thetis.utility