Log the Courant number

gusto/configuration.py

@@ -93,6 +93,7 @@ class OutputParameters(Configuration):
     checkpoint_pickup_filename = None
     chkptfreq = 1
     dirname = None
+    log_courant = True
     #: TODO: Should the output fields be interpolated or projected to
     #: a linear space?  Default is interpolation.
     project_fields = False

gusto/diagnostics.py

@@ -156,6 +156,7 @@ def __init__(self, space=None, method='interpolate', required_fields=()):
         self.space = space
         self.method = method
         self.expr = None
+        self.to_dump = True
 
         # Property to allow graceful failures if solve method not valid
         if not hasattr(self, "solve_implemented"):
@@ -195,7 +196,7 @@ def setup(self, domain, state_fields, space=None):
                 f'Diagnostics {self.name} is using a function space which does not have a name'
             domain.spaces(space.name, V=space)
 
-            self.field = state_fields(self.name, space=space, dump=True, pick_up=False)
+            self.field = state_fields(self.name, space=space, dump=self.to_dump, pick_up=False)
 
             if self.method != 'solve':
                 assert self.expr is not None, \
@@ -233,6 +234,52 @@ class CourantNumber(DiagnosticField):
     """Dimensionless Courant number diagnostic field."""
     name = "CourantNumber"
 
+    def __init__(self, velocity='u', name=None, to_dump=True, space=None,
+                 method='interpolate', required_fields=()):
+        """
+        Args:
+            velocity (str or :class:`ufl.Expr`, optional): the velocity field to
+                take the Courant number of. Can be a string referring to an
+                existing field, or an expression. If it is an expression, the
+                name argument is required. Defaults to 'u'.
+            name (str, optional): the name to append to "CourantNumber" to form
+                the name of this diagnostic. This argument must be provided if
+                the velocity is an expression (rather than a string). Defaults
+                to None.
+            to_dump (bool, optional): whether this diagnostic should be dumped.
+                Defaults to True.
+            space (:class:`FunctionSpace`, optional): the function space to
+                evaluate the diagnostic field in. Defaults to None, in which
+                case a default space will be chosen for this diagnostic.
+            method (str, optional): a string specifying the method of evaluation
+                for this diagnostic. Valid options are 'interpolate', 'project',
+                'assign' and 'solve'. Defaults to 'interpolate'.
+            required_fields (tuple, optional): tuple of names of the fields that
+                are required for the computation of this diagnostic field.
+                Defaults to ().
+        """
+        # Work out whether to take Courant number from field or expression
+        if type(velocity) is str:
+            # Default name should just be CourantNumber
+            if velocity == 'u':
+                self.name = 'CourantNumber'
+            elif name is None:
+                self.name = 'CourantNumber_'+velocity
+            else:
+                self.name = 'CourantNumber_'+name
+        else:
+            if name is None:
+                raise ValueError('CourantNumber diagnostic: if provided '
+                                 + 'velocity is an expression then the name '
+                                 + 'argument must be provided')
+            self.name = 'CourantNumber_'+name
+
+        self.velocity = velocity
+        super().__init__(space=space, method=method, required_fields=required_fields)
+
+        # Done after super init to ensure that it is not always set to True
+        self.to_dump = to_dump
+
     def setup(self, domain, state_fields):
         """
         Sets up the :class:`Function` for the diagnostic field.
@@ -247,7 +294,11 @@ def setup(self, domain, state_fields):
         test = TestFunction(V)
         self.area = Function(V)
         assemble(test*dx, tensor=self.area)
-        u = state_fields("u")
+
+        if type(self.velocity) is str:
+            u = state_fields(self.velocity)
+        else:
+            u = self.velocity
 
         self.expr = sqrt(dot(u, u))/sqrt(self.area)*domain.dt
 

gusto/io.py

@@ -5,7 +5,7 @@
 from netCDF4 import Dataset
 import sys
 import time
-from gusto.diagnostics import Diagnostics
+from gusto.diagnostics import Diagnostics, CourantNumber
 from gusto.meshes import get_flat_latlon_mesh
 from firedrake import (Function, functionspaceimpl, File,
                        DumbCheckpoint, FILE_CREATE, FILE_READ, CheckpointFile)
@@ -238,6 +238,60 @@ def log_parameters(self, equation):
             logger.info("Physical parameters that take non-default values:")
             logger.info(", ".join("%s: %s" % (k, float(v)) for (k, v) in vars(equation.parameters).items()))
 
+    def setup_log_courant(self, state_fields, name='u', expression=None):
+        """
+        Sets up Courant number diagnostics to be logged.
+
+        Args:
+            state_fields (:class:`StateFields`): the model's field container.
+            name (str, optional): the name of the field to log the Courant
+                number of. Defaults to 'u'.
+            expression (:class:`ufl.Expr`, optional): expression of velocity
+                field to take Courant number of. Defaults to None, in which case
+
+        """
+
+        if self.output.log_courant:
+            diagnostic_names = [diagnostic.name for diagnostic in self.diagnostic_fields]
+            courant_name = None if name == 'u' else name
+
+            # Set up diagnostic if it hasn't already been
+            if courant_name not in diagnostic_names and 'u' in state_fields._field_names:
+                if expression is None:
+                    diagnostic = CourantNumber(to_dump=False)
+                elif expression is not None:
+                    diagnostic = CourantNumber(velocity=expression, name=courant_name, to_dump=False)
+
+                self.diagnostic_fields.append(diagnostic)
+                diagnostic.setup(self.domain, state_fields)
+                self.diagnostics.register(diagnostic.name)
+
+    def log_courant(self, state_fields, name='u', message=None):
+        """
+        Logs the maximum Courant number value.
+
+        Args:
+            state_fields (:class:`StateFields`): the model's field container.
+            name (str, optional): the name of the field to log the Courant
+                number of. Defaults to 'u'.
+            message (str, optional): an extra message to be logged. Defaults to
+                None.
+        """
+
+        if self.output.log_courant and 'u' in state_fields._field_names:
+            diagnostic_names = [diagnostic.name for diagnostic in self.diagnostic_fields]
+            courant_name = 'CourantNumber' if name == 'u' else 'CourantNumber_'+name
+            courant_idx = diagnostic_names.index(courant_name)
+            courant_diagnostic = self.diagnostic_fields[courant_idx]
+            courant_diagnostic.compute()
+            courant_field = state_fields(courant_name)
+            courant_max = self.diagnostics.max(courant_field)
+
+            if message is None:
+                logger.info(f'Max Courant: {courant_max:.2e}')
+            else:
+                logger.info(f'Max Courant {message}: {courant_max:.2e}')
+
     def setup_diagnostics(self, state_fields):
         """
         Prepares the I/O for computing the model's global diagnostics and

gusto/timeloop.py

@@ -152,6 +152,10 @@ def run(self, t, tmax, pick_up=False):
 
         # Set up diagnostics, which may set up some fields necessary to pick up
         self.io.setup_diagnostics(self.fields)
+        self.io.setup_log_courant(self.fields)
+        if self.transporting_velocity != "prognostic":
+            self.io.setup_log_courant(self.fields, name='transporting_velocity',
+                                      expression=self.transporting_velocity)
 
         if pick_up:
             # Pick up fields, and return other info to be picked up
@@ -171,6 +175,8 @@ def run(self, t, tmax, pick_up=False):
 
             self.x.update()
 
+            self.io.log_courant(self.fields)
+
             self.timestep()
 
             self.t.assign(self.t + self.dt)
@@ -551,6 +557,8 @@ def timestep(self):
         for k in range(self.maxk):
 
             with timed_stage("Transport"):
+                self.io.log_courant(self.fields, 'transporting_velocity',
+                                    f'transporting velocity, outer iteration {k}')
                 for name, scheme in self.active_transport:
                     # transports a field from xstar and puts result in xp
                     scheme.apply(xp(name), xstar(name))