Moving meshes final

examples/shallow_water/mm_ot_sw_galewsky_jet.py

@@ -0,0 +1,223 @@
+from gusto import *
+from firedrake import IcosahedralSphereMesh, \
+    Constant, ge, le, exp, cos, \
+    conditional, interpolate, SpatialCoordinate, VectorFunctionSpace, \
+    Function, assemble, dx, pi, CellNormal
+
+import numpy as np
+import sys
+
+day = 24.*60.*60.
+dt = 240.
+
+if '--running-tests' in sys.argv:
+    ref_level = 3
+    dt = 480.
+    tmax = 1440.
+    cubed_sphere = False
+else:
+    # setup resolution and timestepping parameters
+    cubed_sphere = True
+    if cubed_sphere:
+        ncells = 24
+        dt = 200.
+    else:
+        ref_level = 4
+        dt = 240.
+    tmax = 6*day
+
+# setup shallow water parameters
+R = 6371220.
+H = 10000.
+parameters = ShallowWaterParameters(H=H)
+
+perturb = True
+if perturb:
+    dirname = "mm_ot_sw_galewsky_jet_perturbed_dt%s" % dt
+else:
+    dirname = "mm_ot_sw_galewsky_jet_unperturbed"
+
+# Domain
+if cubed_sphere:
+    dirname += "_cs%s" % ncells
+    mesh = GeneralCubedSphereMesh(radius=R,
+                                  num_cells_per_edge_of_panel=ncells,
+                                  degree=2)
+    family = "RTCF"
+else:
+    dirname += "_is%s" % ref_level
+    mesh = IcosahedralSphereMesh(radius=R,
+                                 refinement_level=ref_level, degree=2)
+    family = "BDM"
+domain = Domain(mesh, dt, family, 1, move_mesh=True)
+
+# Equation
+Omega = parameters.Omega
+x = SpatialCoordinate(domain.mesh)
+fexpr = 2*Omega*x[2]/R
+eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr,
+                             u_transport_option="vector_invariant_form")
+
+# I/O
+output = OutputParameters(dirname=dirname,
+                          dumplist_latlon=['D', 'PotentialVorticity',
+                                           'RelativeVorticity'])
+
+pv = PotentialVorticity()
+diagnostic_fields = [pv, RelativeVorticity()]
+io = IO(domain, output, diagnostic_fields=diagnostic_fields)
+
+# Transport schemes
+transported_fields = [TrapeziumRule(domain, "u"),
+                      SSPRK3(domain, "D")]
+
+transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")]
+
+
+# Mesh movement
+def update_pv():
+    pv()
+
+
+def reinterpolate_coriolis():
+    domain.k = interpolate(x/R, domain.mesh.coordinates.function_space())
+    domain.outward_normals.interpolate(CellNormal(domain.mesh))
+    eqns.prescribed_fields("coriolis").interpolate(fexpr)
+
+
+monitor = MonitorFunction("PotentialVorticity", adapt_to="gradient")
+mesh_generator = OptimalTransportMeshGenerator(domain.mesh,
+                                               monitor,
+                                               pre_meshgen_callback=update_pv,
+                                               post_meshgen_callback=reinterpolate_coriolis)
+
+# Time stepper
+stepper = MeshMovement(eqns, io, transported_fields,
+                       spatial_methods=transport_methods,
+                       mesh_generator=mesh_generator)
+
+# initial conditions
+u0 = stepper.fields("u")
+D0 = stepper.fields("D")
+
+# get lat lon coordinates
+theta, lamda = latlon_coords(mesh)
+
+# expressions for meridional and zonal velocity
+u_max = 80.0
+theta0 = pi/7.
+theta1 = pi/2. - theta0
+en = np.exp(-4./((theta1-theta0)**2))
+u_zonal_expr = (u_max/en)*exp(1/((theta - theta0)*(theta - theta1)))
+u_zonal = conditional(ge(theta, theta0), conditional(le(theta, theta1), u_zonal_expr, 0.), 0.)
+u_merid = 0.0
+
+# get cartesian components of velocity
+uexpr = sphere_to_cartesian(mesh, u_zonal, u_merid)
+u0.project(uexpr, form_compiler_parameters={'quadrature_degree': 12})
+
+Rc = Constant(R)
+g = Constant(parameters.g)
+
+
+def D_integrand(th):
+    # Initial D field is calculated by integrating D_integrand w.r.t. theta
+    # Assumes the input is between theta0 and theta1.
+    # Note that this function operates on vectorized input.
+    from scipy import exp, sin, tan
+    f = 2.0*parameters.Omega*sin(th)
+    u_zon = (80.0/en)*exp(1.0/((th - theta0)*(th - theta1)))
+    return u_zon*(f + tan(th)*u_zon/R)
+
+
+def Dval(X):
+    # Function to return value of D at X
+    from scipy import integrate
+
+    # Preallocate output array
+    val = np.zeros(len(X))
+
+    angles = np.zeros(len(X))
+
+    # Minimize work by only calculating integrals for points with
+    # theta between theta_0 and theta_1.
+    # For theta <= theta_0, the integral is 0
+    # For theta >= theta_1, the integral is constant.
+
+    # Precalculate this constant:
+    poledepth, _ = integrate.fixed_quad(D_integrand, theta0, theta1, n=64)
+    poledepth *= -R/parameters.g
+
+    angles[:] = np.arcsin(X[:, 2]/R)
+
+    for ii in range(len(X)):
+        if angles[ii] <= theta0:
+            val[ii] = 0.0
+        elif angles[ii] >= theta1:
+            val[ii] = poledepth
+        else:
+            # Fixed quadrature with 64 points gives absolute errors below 1e-13
+            # for a quantity of order 1e-3.
+            v, _ = integrate.fixed_quad(D_integrand, theta0, angles[ii], n=64)
+            val[ii] = -(R/parameters.g)*v
+
+    return val
+
+
+# Get coordinates to pass to Dval function
+W = VectorFunctionSpace(mesh, D0.ufl_element())
+X = interpolate(mesh.coordinates, W)
+D0.dat.data[:] = Dval(X.dat.data_ro)
+
+# Adjust mean value of initial D
+C = Function(D0.function_space()).assign(Constant(1.0))
+area = assemble(C*dx)
+Dmean = assemble(D0*dx)/area
+D0 -= Dmean
+D0 += Constant(parameters.H)
+
+# optional perturbation
+if perturb:
+    alpha = Constant(1/3.)
+    beta = Constant(1/15.)
+    Dhat = Constant(120.)
+    theta2 = Constant(pi/4.)
+    g = Constant(parameters.g)
+    D_pert = Function(D0.function_space()).interpolate(Dhat*cos(theta)*exp(-(lamda/alpha)**2)*exp(-((theta2 - theta)/beta)**2))
+    D0 += D_pert
+
+
+def initialise_fn():
+    u0 = stepper.fields("u")
+    D0 = stepper.fields("D")
+
+    u0.project(uexpr, form_compiler_parameters={'quadrature_degree': 12})
+
+    X = interpolate(domain.mesh.coordinates, W)
+    D0.dat.data[:] = Dval(X.dat.data_ro)
+    area = assemble(C*dx)
+    Dmean = assemble(D0*dx)/area
+    D0 -= Dmean
+    D0 += parameters.H
+    if perturb:
+        theta, lamda = latlon_coords(domain.mesh)
+        D_pert.interpolate(Dhat*cos(theta)*exp(-(lamda/alpha)**2)*exp(-((theta2 - theta)/beta)**2))
+        D0 += D_pert
+    domain.k = interpolate(x/R, domain.mesh.coordinates.function_space())
+    domain.outward_normals.interpolate(CellNormal(domain.mesh))
+    eqns.prescribed_fields("coriolis").interpolate(fexpr)
+    pv()
+
+
+pv.setup(domain, stepper.fields)
+mesh_generator.get_first_mesh(initialise_fn)
+
+domain.k = interpolate(x/R, domain.mesh.coordinates.function_space())
+domain.outward_normals.interpolate(CellNormal(domain.mesh))
+eqns.prescribed_fields("coriolis").interpolate(fexpr)
+pv()
+
+Dbar = Function(D0.function_space()).assign(H)
+stepper.set_reference_profiles([('D', Dbar)])
+
+stepper.run(t=0, tmax=tmax)

gusto/__init__.py

@@ -32,6 +32,7 @@ def perp(self, o, a):
 from gusto.limiters import *                         # noqa
 from gusto.linear_solvers import *                   # noqa
 from gusto.meshes import *                           # noqa
+from gusto.moving_mesh import *                      # noqa
 from gusto.numerical_integrator import *             # noqa
 from gusto.physics import *                          # noqa
 from gusto.preconditioners import *                  # noqa

gusto/diagnostics.py

@@ -1452,7 +1452,7 @@ def setup(self, domain, state_fields, vorticity_type=None):
                 f = state_fields("coriolis")
                 L += gamma*f*dx
 
-            problem = LinearVariationalProblem(a, L, self.field)
+            problem = LinearVariationalProblem(a, L, self.field, constant_jacobian=False)
             self.evaluator = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"})
 
 

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

@@ -10,7 +10,8 @@
 from gusto.fml import (Term, all_terms, keep, drop, Label, subject, name,
                        replace_subject, replace_trial_function)
 from gusto.labels import (time_derivative, transport, prognostic, hydrostatic,
-                          linearisation, pressure_gradient, coriolis)
+                          linearisation, pressure_gradient, coriolis,
+                          transporting_velocity)
 from gusto.thermodynamics import exner_pressure
 from gusto.common_forms import (advection_form, continuity_form,
                                 vector_invariant_form, kinetic_energy_form,
@@ -236,6 +237,7 @@ def __init__(self, field_names, domain, space_names,
 
         # Make the full mixed function space
         W = MixedFunctionSpace(self.spaces)
+        self.W = W
 
         # Can now call the underlying PrognosticEquation
         full_field_name = "_".join(self.field_names)
@@ -668,14 +670,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 == "vector_invariant_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')
+            if domain.move_mesh:
+                ke_form = prognostic(kinetic_energy_form(w, u, 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 -= ke_form
         elif u_transport_option == "circulation_form":
-            ke_form = prognostic(kinetic_energy_form(w, u, u), 'u')
-            u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), 'u') + ke_form
+            ke_form = prognostic(kinetic_energy_form(w, u, u), "u")
+            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/forcing.py

@@ -90,11 +90,13 @@ def __init__(self, equation, alpha):
 
         # now we can set up the explicit and implicit problems
         explicit_forcing_problem = LinearVariationalProblem(
-            a.form, L_explicit.form, self.xF, bcs=bcs
+            a.form, L_explicit.form, self.xF, bcs=bcs,
+            constant_jacobian=False
         )
 
         implicit_forcing_problem = LinearVariationalProblem(
-            a.form, L_implicit.form, self.xF, bcs=bcs
+            a.form, L_implicit.form, self.xF, bcs=bcs,
+            constant_jacobian=False
         )
 
         self.solvers = {}

gusto/linear_solvers.py

@@ -599,7 +599,8 @@ def __init__(self, equation, alpha):
         bcs = [DirichletBC(W.sub(0), bc.function_arg, bc.sub_domain) for bc in equation.bcs['u']]
         problem = LinearVariationalProblem(aeqn.form,
                                            action(Leqn.form, self.xrhs),
-                                           self.dy, bcs=bcs)
+                                           self.dy, bcs=bcs,
+                                           constant_jacobian=False)
 
         self.solver = LinearVariationalSolver(problem,
                                               solver_parameters=self.solver_parameters,

gusto/moving_mesh/__init__.py

@@ -0,0 +1,3 @@
+from gusto.moving_mesh.monge_ampere import *         # noqa
+from gusto.moving_mesh.monitor import *              # noqa
+from gusto.moving_mesh.utility_functions import *    # noqa