I think this fixes the tests but might break mesh movement - check si…

…gn of ke term!!

gusto/domain.py

@@ -25,7 +25,8 @@ class Domain(object):
     the same then this can be specified through the "degree" argument.
     """
     def __init__(self, mesh, dt, family, degree=None,
-                 horizontal_degree=None, vertical_degree=None):
+                 horizontal_degree=None, vertical_degree=None,
+                 move_mesh=False):
         """
         Args:
             mesh (:class:`Mesh`): the model's mesh.
@@ -59,6 +60,9 @@ def __init__(self, mesh, dt, family, degree=None,
         else:
             raise TypeError(f'dt must be a Constant, float or int, not {type(dt)}')
 
+        # store whether we are moving the mesh
+        self.move_mesh = move_mesh
+
         # -------------------------------------------------------------------- #
         # Build compatible function spaces
         # -------------------------------------------------------------------- #

gusto/equations.py

@@ -669,20 +669,29 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None,
         # -------------------------------------------------------------------- #
         # Transport Terms
         # -------------------------------------------------------------------- #
+
+
+        # Mesh movement requires the circulation form, and an
+        # additional modification
+        if domain.move_mesh:
+            assert u_transport_option == "circulation_form"
+
         # Velocity transport term -- depends on formulation
         if u_transport_option == "vector_invariant_form":
             u_adv = prognostic(vector_invariant_form(domain, w, u, u), 'u')
         elif u_transport_option == "vector_advection_form":
             u_adv = prognostic(advection_form(w, u, u), 'u')
         elif u_transport_option == "circulation_form":
-            ke_form = prognostic(kinetic_energy_form(domain, w, u), "u")
-            ke_form = transport.remove(ke_form)
-            ke_form = ke_form.label_map(
-                lambda t: t.has_label(transporting_velocity),
-                lambda t: Term(ufl.replace(
-                    t.form, {t.get(transporting_velocity): u}), t.labels))
-            ke_form = transporting_velocity.remove(ke_form)
-            u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") - ke_form
+            ke_form = prognostic(kinetic_energy_form(w, u, u), "u")
+            if domain.move_mesh:
+                ke_form = transport.remove(ke_form)
+                ke_form = ke_form.label_map(
+                    lambda t: t.has_label(transporting_velocity),
+                    lambda t: Term(ufl.replace(
+                        t.form, {t.get(transporting_velocity): u}), t.labels))
+                ke_form = transporting_velocity.remove(ke_form)
+                ke_form *= -1
+            u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), "u") + ke_form
         else:
             raise ValueError("Invalid u_transport_option: %s" % u_transport_option)
 

gusto/timeloop.py

@@ -678,6 +678,9 @@ def __init__(self, equation_set, io, transport_schemes,
         Vu = equation_set.domain.spaces("HDiv")
         self.uadv = Function(Vu)
 
+        assert isinstance(equation_set, ShallowWaterEquations), \
+            'Mesh movement only works with the shallow water equation set'
+
         super().__init__(equation_set=equation_set,
                          io=io,
                          transport_schemes=transport_schemes,
@@ -708,6 +711,10 @@ def __init__(self, equation_set, io, transport_schemes,
         self.v1 = Function(mesh.coordinates.function_space())
         self.v1_V1 = Function(Vu)
 
+        # Set up diagnostics - also called in the run method, is it ok
+        # to do this twice (I suspect not)
+        self.io.setup_diagnostics(self.fields)
+
         mesh_generator.monitor.setup(self.fields)
         mesh_generator.setup()