routine to generate balanced initial fields for SWE

gusto/initialisation/__init__.py

@@ -1,2 +1,3 @@
 from gusto.initialisation.hydrostatic_initialisation import *   # noqa
-from gusto.initialisation.numerical_integrator import *         # noqa
+from gusto.initialisation.numerical_integrator import *         # noqa
+from gusto.initialisation.sw_balance import *                   # noqa

gusto/initialisation/sw_balance.py

@@ -0,0 +1,57 @@
+from firedrake import TestFunction, TrialFunction, Function, \
+    dot, grad, dx, VectorSpaceBasis, solve, TestFunctions, TrialFunctions, \
+    inner, div, Constant, assemble
+
+
+def nondivergent_flow(equation, zeta0, u0, D0):
+    """
+    Returns u0 and D0, balanced velocity and depth fields, given a
+    vorticity field zeta0. Balance is defined as
+
+    Args:
+        equation (:class:`PrognosticEquation`): the model's equation object.
+        zeta0 (:class:`ufl.Expr`): the input vorticity field.
+        u0 (:class:`Function`): the velocity to be returned.
+        D0 (:class:`Function`): the depth to be returned.
+    """
+
+    domain = equation.domain
+    Vcg = domain.spaces("H1")
+
+    # compute initial streamfunction from vorticity by solving Poisson equation
+    v = TestFunction(Vcg)
+    p = TrialFunction(Vcg)
+    psi = Function(Vcg)
+    a = -dot(grad(v), grad(p)) * dx
+    L = v * zeta0 * dx
+    nullspace = VectorSpaceBasis(constant=True)
+    solve(a == L, psi, nullspace=nullspace,
+          solver_parameters={'ksp_type': 'cg', 'pc_type': 'none'})
+
+    # compute initial velocity from streamfunction
+    u0.project(domain.perp(grad(psi)))
+
+    # solve mixed Poisson problem for (v, depth) with v=u_t and
+    # div(v)=0 so that we don't generate any divergence initially
+    VHdiv = domain.spaces("HDiv")
+    Vdg = domain.spaces("L2")
+    W = VHdiv * Vdg
+    v, h = TrialFunctions(W)
+    p, q = TestFunctions(W)
+    g = equation.parameters.g
+    f = equation.prescribed_fields("coriolis")
+    a = inner(p, v) * dx - g * div(p) * h * dx + q * div(v) * dx
+    L = (
+        -(f + zeta0) * inner(p, domain.perp(u0)) * dx
+        + 0.5 * div(p) * dot(u0, u0) * dx
+    )
+    w = Function(W)
+    solve(a == L, w, nullspace=nullspace)
+    _, D = w.subfunctions
+    D0.assign(D)
+
+    # adjust depth to have initial mean of H as set in the parameters
+    C = Function(Vdg).assign(Constant(1.0))
+    area = assemble(C*dx)
+    Dmean = assemble(D*dx)/area
+    D0.assign(D0 - Dmean + equation.parameters.H)

integration-tests/balance/test_sw_balance.py

@@ -0,0 +1,73 @@
+from gusto import *
+from firedrake import SpatialCoordinate, conditional, Function
+
+
+def setup_balance(dirname):
+    # ------------------------------------------------------------------------ #
+    # Parameters for test case
+    # ------------------------------------------------------------------------ #
+
+    radius = 6371220.                 # planetary radius (m)
+    mean_depth = 222.                 # reference depth (m)
+    dt = 3600.                        # timestep (s)
+    tmax = 10 * dt                    # final time (s)
+
+    # ------------------------------------------------------------------------ #
+    # Set up model objects
+    # ------------------------------------------------------------------------ #
+
+    # Domain
+    mesh = GeneralIcosahedralSphereMesh(radius, 12, degree=2)
+
+    # Equation
+    parameters = ShallowWaterParameters(mesh, H=mean_depth)
+    eqns = ShallowWaterEquations
+
+    # I/O
+    output = OutputParameters(dirname=dirname, dumpfreq=10)
+
+    # model
+    model = SIQNModel(mesh, dt, parameters, eqns, family='BDM')
+    model.setup(output)
+
+    # ------------------------------------------------------------------------ #
+    # Initial conditions
+    # ------------------------------------------------------------------------ #
+    stepper = model.stepper
+    u0 = stepper.fields("u")
+    D0 = stepper.fields("D")
+
+    # set initial vorticity to be nonzero in a latitude band
+    Vcg = model.domain.spaces("H1")
+    phi_c = pi/18
+    phi_w = 4.5*pi/180
+    zeta_s = 3e-5
+    x, y, z = SpatialCoordinate(mesh)
+    _, phi, _ = lonlatr_from_xyz(x, y, z)
+    zeta_expr = conditional(abs(phi-phi_c) > phi_w/2, 0, zeta_s)
+    zeta0 = Function(Vcg).interpolate(zeta_expr)
+
+    # calculate corresponding velocity and depth such that initial
+    # conditions are nondivergent and div(u_t)=0
+    nondivergent_flow(model.equation, zeta0, u0, D0)
+
+    Dbar = Function(D0.function_space()).assign(mean_depth)
+    stepper.set_reference_profiles([('D', Dbar)])
+
+    return stepper, tmax, model.domain.spaces("L2")
+
+
+def run_balance(dirname):
+
+    stepper, tmax, hdiv_space = setup_balance(dirname)
+    stepper.run(t=0, tmax=tmax)
+    return hdiv_space, stepper.fields("u")
+
+
+def test_nondivergent_sw(tmpdir):
+
+    dirname = str(tmpdir)
+    hdiv_space, u = run_balance(dirname)
+    divu = Function(hdiv_space).project(div(u))
+    tol = 1e-6
+    assert divu.dat.data.max() < tol and abs(divu.dat.data.min()) < tol