From 7dbfcb15b174741183b7cfcae8e07676d80b964c Mon Sep 17 00:00:00 2001 From: jshipton Date: Tue, 1 Sep 2026 13:31:48 +0100 Subject: [PATCH 1/2] routine to generate initial fields for SWE which have zero divergence and zero time-change to divergence, given vorticity --- gusto/initialisation/__init__.py | 3 +- gusto/initialisation/sw_balance.py | 46 ++++++++++++ integration-tests/balance/test_sw_balance.py | 77 ++++++++++++++++++++ 3 files changed, 125 insertions(+), 1 deletion(-) create mode 100644 gusto/initialisation/sw_balance.py create mode 100644 integration-tests/balance/test_sw_balance.py diff --git a/gusto/initialisation/__init__.py b/gusto/initialisation/__init__.py index 6aa7f52ea..4fd48bcc9 100644 --- a/gusto/initialisation/__init__.py +++ b/gusto/initialisation/__init__.py @@ -1,2 +1,3 @@ from gusto.initialisation.hydrostatic_initialisation import * # noqa -from gusto.initialisation.numerical_integrator import * # noqa \ No newline at end of file +from gusto.initialisation.numerical_integrator import * # noqa +from gusto.initialisation.sw_balance import * # noqa diff --git a/gusto/initialisation/sw_balance.py b/gusto/initialisation/sw_balance.py new file mode 100644 index 000000000..6238a4642 --- /dev/null +++ b/gusto/initialisation/sw_balance.py @@ -0,0 +1,46 @@ +from firedrake import TestFunction, TrialFunction, Function, \ + dot, grad, dx, VectorSpaceBasis, solve, TestFunctions, TrialFunctions, \ + inner, div, Constant, assemble + +def nondivergent_velocity(equation, zeta0, u0, D0): + """ + + Args: + """ + + domain = equation.domain + Vcg = domain.spaces("H1") + + # compute initial streamfunction + 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 + u0.project(domain.perp(grad(psi))) + + # solve for the depth 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) diff --git a/integration-tests/balance/test_sw_balance.py b/integration-tests/balance/test_sw_balance.py new file mode 100644 index 000000000..cad9e811d --- /dev/null +++ b/integration-tests/balance/test_sw_balance.py @@ -0,0 +1,77 @@ +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 + # ------------------------------------------------------------------------ # + + g = parameters.g + Omega = parameters.Omega + + 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_velocity(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-16 + assert divu.dat.data.max() < tol and abs(divu.dat.data.min()) < tol From 4f3628442524112cd762095a49e997e1362b6238 Mon Sep 17 00:00:00 2001 From: jshipton Date: Tue, 1 Sep 2026 19:43:52 +0100 Subject: [PATCH 2/2] lint, docs, sensible tol for test --- gusto/initialisation/sw_balance.py | 21 +++++++++++++++----- integration-tests/balance/test_sw_balance.py | 10 +++------- 2 files changed, 19 insertions(+), 12 deletions(-) diff --git a/gusto/initialisation/sw_balance.py b/gusto/initialisation/sw_balance.py index 6238a4642..07bc3df80 100644 --- a/gusto/initialisation/sw_balance.py +++ b/gusto/initialisation/sw_balance.py @@ -2,16 +2,23 @@ dot, grad, dx, VectorSpaceBasis, solve, TestFunctions, TrialFunctions, \ inner, div, Constant, assemble -def nondivergent_velocity(equation, zeta0, u0, D0): + +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 + # compute initial streamfunction from vorticity by solving Poisson equation v = TestFunction(Vcg) p = TrialFunction(Vcg) psi = Function(Vcg) @@ -21,10 +28,11 @@ def nondivergent_velocity(equation, zeta0, u0, D0): solve(a == L, psi, nullspace=nullspace, solver_parameters={'ksp_type': 'cg', 'pc_type': 'none'}) - # compute initial velocity + # compute initial velocity from streamfunction u0.project(domain.perp(grad(psi))) - # solve for the depth so that we don't generate any divergence initially + # 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 @@ -33,7 +41,10 @@ def nondivergent_velocity(equation, zeta0, u0, D0): 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 + 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 diff --git a/integration-tests/balance/test_sw_balance.py b/integration-tests/balance/test_sw_balance.py index cad9e811d..6e80ca90d 100644 --- a/integration-tests/balance/test_sw_balance.py +++ b/integration-tests/balance/test_sw_balance.py @@ -33,10 +33,6 @@ def setup_balance(dirname): # ------------------------------------------------------------------------ # # Initial conditions # ------------------------------------------------------------------------ # - - g = parameters.g - Omega = parameters.Omega - stepper = model.stepper u0 = stepper.fields("u") D0 = stepper.fields("D") @@ -53,7 +49,7 @@ def setup_balance(dirname): # calculate corresponding velocity and depth such that initial # conditions are nondivergent and div(u_t)=0 - nondivergent_velocity(model.equation, zeta0, u0, D0) + nondivergent_flow(model.equation, zeta0, u0, D0) Dbar = Function(D0.function_space()).assign(mean_depth) stepper.set_reference_profiles([('D', Dbar)]) @@ -69,9 +65,9 @@ def run_balance(dirname): def test_nondivergent_sw(tmpdir): - + dirname = str(tmpdir) hdiv_space, u = run_balance(dirname) divu = Function(hdiv_space).project(div(u)) - tol = 1e-16 + tol = 1e-6 assert divu.dat.data.max() < tol and abs(divu.dat.data.min()) < tol