Fix Diagnostics of Diagnostics

examples/compressible/skamarock_klemp_nonhydrostatic.py

@@ -27,10 +27,10 @@
     tmax = dt
     dumpfreq = 1
 else:
-    nlayers = 10
-    columns = 150
-    tmax = 3600.
-    dumpfreq = int(tmax / (2*dt))
+    nlayers = 5
+    columns = 30
+    tmax = 10*dt
+    dumpfreq = 1
 
 # ---------------------------------------------------------------------------- #
 # Set up model objects
@@ -44,22 +44,22 @@
 # Equation
 Tsurf = 300.
 parameters = CompressibleParameters()
-eqns = CompressibleEulerEquations(domain, parameters)
+eqns = CompressibleEulerEquations(domain, parameters, u_transport_option='vector_advection_form')
+print(f'Number of DOFs = {eqns.X.function_space().dim()}')
 
 # I/O
 points_x = np.linspace(0., L, 100)
 points_z = [H/2.]
 points = np.array([p for p in itertools.product(points_x, points_z)])
-dirname = 'skamarock_klemp_nonlinear'
+dirname = 'testing'
 
 # Dumping point data using legacy PointDataOutput is not supported in parallel
 if COMM_WORLD.size == 1:
     output = OutputParameters(
         dirname=dirname,
         dumpfreq=dumpfreq,
-        pddumpfreq=dumpfreq,
-        dumplist=['u'],
-        point_data=[('theta_perturbation', points)],
+        dump_nc=True,
+        dump_vtus=False
     )
 else:
     logger.warning(
@@ -69,13 +69,11 @@
     output = OutputParameters(
         dirname=dirname,
         dumpfreq=dumpfreq,
-        pddumpfreq=dumpfreq,
-        dumplist=['u'],
-    )
+        dump_nc=True,
+        dump_vtus=False
+          )
 
-diagnostic_fields = [CourantNumber(), Gradient("u"), Perturbation('theta'),
-                     Gradient("theta_perturbation"), Perturbation('rho'),
-                     RichardsonNumber("theta", parameters.g/Tsurf), Gradient("theta")]
+diagnostic_fields = [CourantNumber(), Pressure(eqns), SteadyStateError('Pressure_Vt')]
 io = IO(domain, output, diagnostic_fields=diagnostic_fields)
 
 # Transport schemes

gusto/diagnostics.py

@@ -1033,12 +1033,33 @@ def setup(self, domain, state_fields):
             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
+
+            self.init_field_set = False
+
+            #HACK Defo not the best way to do this perhaps we want a flag for checkpoint or prognostic
+            if not self.name in ['rho', 'u', 'theta']:
+                self.init_field_set = True
+
+            # Attach state fields to self so that we can pick it up in compute
+            self.state_fields = state_fields
             field2.assign(field1)
 
+        else:
+            self.init_field_set = True
+
         super().setup(domain, state_fields)
 
+    def compute(self):
+        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."""
@@ -1200,7 +1221,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"))
 
 
@@ -1243,6 +1267,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"))