Fix Diagnostics of Diagnostics

examples/compressible/dcmip_3_1_meanflow_quads.py

@@ -82,7 +82,9 @@
     dirname=dirname,
     dumpfreq=dumpfreq,
 )
-diagnostic_fields = [Perturbation('theta'), Perturbation('rho'), CompressibleKineticEnergy()]
+diagnostic_fields = [Perturbation('theta'), Perturbation('rho'),
+                     CompressibleKineticEnergy(), PotentialEnergy(eqns)]
+
 io = IO(domain, output, diagnostic_fields=diagnostic_fields)
 
 # Transport schemes

examples/compressible/skamarock_klemp_nonhydrostatic.py

@@ -29,7 +29,7 @@
 else:
     nlayers = 10
     columns = 150
-    tmax = 3600.
+    tmax = 3600
     dumpfreq = int(tmax / (2*dt))
 
 # ---------------------------------------------------------------------------- #
@@ -73,9 +73,9 @@
         dumplist=['u'],
     )
 
-diagnostic_fields = [CourantNumber(), Gradient("u"), Perturbation('theta'),
-                     Gradient("theta_perturbation"), Perturbation('rho'),
-                     RichardsonNumber("theta", parameters.g/Tsurf), Gradient("theta")]
+diagnostic_fields = [CourantNumber(), Gradient('u'), Perturbation('theta'),
+                     Gradient('theta_perturbation'), Perturbation('rho'),
+                     RichardsonNumber('theta', parameters.g/Tsurf), Gradient('theta')]
 io = IO(domain, output, diagnostic_fields=diagnostic_fields)
 
 # Transport schemes

examples/compressible/unsaturated_bubble.py

@@ -59,7 +59,7 @@
     checkpoint=False
 )
 diagnostic_fields = [RelativeHumidity(eqns), Perturbation('theta'),
-                     Perturbation('water_vapour'), Perturbation('rho')]
+                     Perturbation('water_vapour'), Perturbation('rho'), Perturbation('RelativeHumidity')]
 io = IO(domain, output, diagnostic_fields=diagnostic_fields)
 
 # Transport schemes -- specify options for using recovery wrapper

examples/shallow_water/williamson_2.py

@@ -62,7 +62,8 @@
         dump_nc=True,
     )
 
-    diagnostic_fields = [RelativeVorticity(), PotentialVorticity(),
+    diagnostic_fields = [RelativeVorticity(), SteadyStateError('RelativeVorticity'),
+                         PotentialVorticity(),
                          ShallowWaterKineticEnergy(),
                          ShallowWaterPotentialEnergy(parameters),
                          ShallowWaterPotentialEnstrophy(),

gusto/diagnostics.py

@@ -1078,15 +1078,39 @@ def setup(self, domain, state_fields):
         """
         # Check if initial field already exists -- otherwise needs creating
         if not hasattr(state_fields, self.field_name2):
+            field1 = state_fields(self.field_name1)
+            field2 = state_fields(self.field_name2, space=field1.function_space(),
+                                  pick_up=True, dump=False)
+            # Attach state fields to self so that we can pick it up in compute
+            self.state_fields = state_fields
+            # The initial value for fields may not have already been set yet so we
+            # postpone setting it until the compute method is called
+            self.init_field_set = False
+        else:
             field1 = state_fields(self.field_name1)
             field2 = state_fields(self.field_name2, space=field1.function_space(),
                                   pick_up=True, dump=False)
             # By default set this new field to the current value
             # This may be overwritten if picking up from a checkpoint
             field2.assign(field1)
+            self.state_fields = state_fields
+            self.init_field_set = True
 
         super().setup(domain, state_fields)
 
+    def compute(self):
+        # The first time the compute method is called we set the initial field.
+        # We do not want to do this if picking up from a checkpoint
+        if not self.init_field_set:
+            # Set initial field
+            full_field = self.state_fields(self.field_name1)
+            init_field = self.state_fields(self.field_name2)
+            init_field.assign(full_field)
+
+            self.init_field_set = True
+
+        super().compute()
+
     @property
     def name(self):
         """Gives the name of this diagnostic field."""
@@ -1248,7 +1272,10 @@ def setup(self, domain, state_fields):
             state_fields (:class:`StateFields`): the model's field container.
         """
         x = SpatialCoordinate(domain.mesh)
-        self.expr = self.rho_averaged * (1 + self.r_t) * self.parameters.g * dot(x, domain.k)
+        self._setup_thermodynamics(domain, state_fields)
+        z = Function(self.rho_averaged.function_space())
+        z.interpolate(dot(x, domain.k))
+        self.expr = self.rho_averaged * (1 + self.r_t) * self.parameters.g * z
         super().setup(domain, state_fields, space=domain.spaces("DG"))
 
 
@@ -1291,6 +1318,7 @@ def setup(self, domain, state_fields):
             state_fields (:class:`StateFields`): the model's field container.
         """
         u = state_fields('u')
+        self._setup_thermodynamics(domain, state_fields)
         self.expr = 0.5 * self.rho_averaged * (1 + self.r_t) * dot(u, u)
         super().setup(domain, state_fields, space=domain.spaces("DG"))
 

gusto/io.py

@@ -641,6 +641,11 @@ def pick_up_from_checkpoint(self, state_fields):
         else:
             raise ValueError("Must set checkpoint True if picking up")
 
+        # Prevent any steady-state diagnostics overwriting their original fields
+        for diagnostic_field in self.diagnostic_fields:
+            if hasattr(diagnostic_field, "init_field_set"):
+                diagnostic_field.init_field_set = True
+
         return t, reference_profiles, step, initial_steps
 
     def dump(self, state_fields, t, step, initial_steps=None):

gusto/timeloop.py

@@ -227,16 +227,19 @@ def set_reference_profiles(self, reference_profiles):
             else:
                 raise ValueError(f'When initialising reference profile {field_name}'
                                  + ' the passed profile must be a Function')
+            # if field name is not prognostic we need to add it
             ref.interpolate(profile)
-
             # Assign profile to X_ref belonging to equation
             if isinstance(self.equation, PrognosticEquationSet):
-                assert field_name in self.equation.field_names, \
-                    f'Cannot set reference profile as field {field_name} not found'
-                idx = self.equation.field_names.index(field_name)
-                X_ref = self.equation.X_ref.subfunctions[idx]
-                X_ref.assign(ref)
-
+                if field_name in self.equation.field_names:
+                    idx = self.equation.field_names.index(field_name)
+                    X_ref = self.equation.X_ref.subfunctions[idx]
+                    X_ref.assign(ref)
+                else:
+                    # reference profile of a diagnostic
+                    # warn user in case they made a typo
+                    logger.warning(f'Setting reference profile for diagnostic {field_name}')
+                    # Don't need to do anything else as value in field container has already been set
         self.reference_profiles_initialised = True