diff --git a/examples/shallow_water/moist_convective_williamson2.py b/examples/shallow_water/moist_convective_williamson2.py new file mode 100644 index 000000000..22ca6b78e --- /dev/null +++ b/examples/shallow_water/moist_convective_williamson2.py @@ -0,0 +1,163 @@ +""" +A moist convective version of the Williamson 2 shallow water test (steady state +geostrophically-balanced flow). The saturation function depends on height, +with a constant background buoyancy/temperature field. +Vapour is initialised very close to saturation and small overshoots will +generate clouds. +""" +from gusto import * +from firedrake import (IcosahedralSphereMesh, SpatialCoordinate, sin, cos, exp) +import sys + +# ----------------------------------------------------------------- # +# Test case parameters +# ----------------------------------------------------------------- # + +dt = 120 + +if '--running-tests' in sys.argv: + tmax = dt + dumpfreq = 1 +else: + day = 24*60*60 + tmax = 5*day + ndumps = 5 + dumpfreq = int(tmax / (ndumps*dt)) + +R = 6371220. +u_max = 20 +phi_0 = 3e4 +epsilon = 1/300 +theta_0 = epsilon*phi_0**2 +g = 9.80616 +H = phi_0/g +xi = 0 +q0 = 200 +beta1 = 110 +alpha = 16 +gamma_v = 0.98 +qprecip = 1e-4 +gamma_r = 1e-3 + +# ----------------------------------------------------------------- # +# Set up model objects +# ----------------------------------------------------------------- # + +# Domain +mesh = IcosahedralSphereMesh(radius=R, refinement_level=3, degree=2) +degree = 1 +domain = Domain(mesh, dt, 'BDM', degree) +x = SpatialCoordinate(mesh) + +# Equations +parameters = ShallowWaterParameters(H=H, g=g) +Omega = parameters.Omega +fexpr = 2*Omega*x[2]/R + +tracers = [WaterVapour(space='DG'), CloudWater(space='DG'), Rain(space='DG')] + +eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, + u_transport_option='vector_advection_form', + active_tracers=tracers) + +# IO +dirname = "moist_convective_williamson2" +output = OutputParameters(dirname=dirname, + dumpfreq=dumpfreq, + dumplist_latlon=['D', 'D_error'], + dump_nc=True, + dump_vtus=True) + +diagnostic_fields = [CourantNumber(), RelativeVorticity(), + PotentialVorticity(), + ShallowWaterKineticEnergy(), + ShallowWaterPotentialEnergy(parameters), + ShallowWaterPotentialEnstrophy(), + SteadyStateError('u'), SteadyStateError('D'), + SteadyStateError('water_vapour'), + SteadyStateError('cloud_water')] + +io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + +# define saturation function +def sat_func(x_in): + h = x_in.split()[1] + lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2]) + numerator = theta_0 + sigma*((cos(phi))**2) * ((w + sigma)*(cos(phi))**2 + 2*(phi_0 - w - sigma)) + denominator = phi_0**2 + (w + sigma)**2*(sin(phi))**4 - 2*phi_0*(w + sigma)*(sin(phi))**2 + theta = numerator/denominator + return q0/(g*h) * exp(20*(theta)) + + +transport_methods = [DGUpwind(eqns, field_name) for field_name in eqns.field_names] + +limiter = DG1Limiter(domain.spaces('DG')) + +transported_fields = [TrapeziumRule(domain, "u"), + SSPRK3(domain, "D"), + SSPRK3(domain, "water_vapour", limiter=limiter), + SSPRK3(domain, "cloud_water", limiter=limiter), + SSPRK3(domain, "rain", limiter=limiter) + ] + +linear_solver = MoistConvectiveSWSolver(eqns) + +sat_adj = SWSaturationAdjustment(eqns, sat_func, + time_varying_saturation=True, + convective_feedback=True, beta1=beta1, + gamma_v=gamma_v, time_varying_gamma_v=False, + parameters=parameters) + +inst_rain = InstantRain(eqns, qprecip, vapour_name="cloud_water", + rain_name="rain", gamma_r=gamma_r) + +physics_schemes = [(sat_adj, ForwardEuler(domain)), + (inst_rain, ForwardEuler(domain))] + +stepper = SemiImplicitQuasiNewton(eqns, io, + transport_schemes=transported_fields, + spatial_methods=transport_methods, + linear_solver=linear_solver, + physics_schemes=physics_schemes) + +# ----------------------------------------------------------------- # +# Initial conditions +# ----------------------------------------------------------------- # + +u0 = stepper.fields("u") +D0 = stepper.fields("D") +v0 = stepper.fields("water_vapour") + +lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2]) + +uexpr = xyz_vector_from_lonlatr(u_max*cos(phi), 0, 0, x) +g = parameters.g +w = Omega*R*u_max + (u_max**2)/2 +sigma = 0 + +Dexpr = H - (1/g)*(w)*((sin(phi))**2) +D_for_v = H - (1/g)*(w + sigma)*((sin(phi))**2) + +# though this set-up has no buoyancy, we use the expression for theta to set up +# the initial vapour +numerator = theta_0 + sigma*((cos(phi))**2) * ((w + sigma)*(cos(phi))**2 + 2*(phi_0 - w - sigma)) +denominator = phi_0**2 + (w + sigma)**2*(sin(phi))**4 - 2*phi_0*(w + sigma)*(sin(phi))**2 +theta = numerator/denominator + +initial_msat = q0/(g*Dexpr) * exp(20*theta) +vexpr = (1 - xi) * initial_msat + +u0.project(uexpr) +D0.interpolate(Dexpr) +v0.interpolate(vexpr) + +# Set reference profiles +Dbar = Function(D0.function_space()).assign(H) +stepper.set_reference_profiles([('D', Dbar)]) + +# ----------------------------------------------------------------- # +# Run +# ----------------------------------------------------------------- # + +stepper.run(t=0, tmax=tmax) diff --git a/gusto/equations.py b/gusto/equations.py index a1421d0ea..c01a0ef38 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -985,7 +985,8 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, # -------------------------------------------------------------------- # # Gravitational Term # -------------------------------------------------------------------- # - gravity_form = subject(prognostic(Term(g*inner(domain.k, w)*dx), 'u'), self.X) + gravity_form = name_label(subject(prognostic(Term(g*inner(domain.k, w)*dx), + 'u'), self.X), "gravity") residual = (mass_form + adv_form + pressure_gradient_form + gravity_form) diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index d1017aa28..4804010c6 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -24,7 +24,7 @@ from abc import ABCMeta, abstractmethod, abstractproperty -__all__ = ["IncompressibleSolver", "LinearTimesteppingSolver", "CompressibleSolver", "ThermalSWSolver"] +__all__ = ["IncompressibleSolver", "LinearTimesteppingSolver", "CompressibleSolver", "ThermalSWSolver", "MoistConvectiveSWSolver"] class TimesteppingSolver(object, metaclass=ABCMeta): @@ -762,3 +762,107 @@ def solve(self, xrhs, dy): self.xrhs.assign(xrhs) self.solver.solve() dy.assign(self.dy) + + +class MoistConvectiveSWSolver(TimesteppingSolver): + """ + Linear solver for the moist convective shallow water equations. + + This solves a linear problem for the shallow water equations with prognostic + variables u (velocity) and D (depth). The linear system is solved using a + hybridised-mixed method. + """ + + solver_parameters = { + 'ksp_type': 'preonly', + 'mat_type': 'matfree', + 'pc_type': 'python', + 'pc_python_type': 'firedrake.HybridizationPC', + 'hybridization': {'ksp_type': 'cg', + 'pc_type': 'gamg', + 'ksp_rtol': 1e-8, + 'mg_levels': {'ksp_type': 'chebyshev', + 'ksp_max_it': 2, + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu'}} + } + + @timed_function("Gusto:SolverSetup") + def _setup_solver(self): + equation = self.equations # just cutting down line length a bit + dt = self.dt + beta_ = dt*self.alpha + Vu = equation.domain.spaces("HDiv") + VD = equation.domain.spaces("DG") + + # Store time-stepping coefficients as UFL Constants + beta = Constant(beta_) + + # Split up the rhs vector + self.xrhs = Function(self.equations.function_space) + u_in = split(self.xrhs)[0] + D_in = split(self.xrhs)[1] + + # Build the reduced function space for u, D + M = MixedFunctionSpace((Vu, VD)) + w, phi = TestFunctions(M) + u, D = TrialFunctions(M) + + # Get background depth + Dbar = split(equation.X_ref)[1] + + g = equation.parameters.g + + eqn = ( + inner(w, (u - u_in)) * dx + - beta * (D - Dbar) * div(w*g) * dx + + inner(phi, (D - D_in)) * dx + + beta * phi * Dbar * div(u) * dx + ) + + aeqn = lhs(eqn) + Leqn = rhs(eqn) + + # Place to put results of (u,D) solver + self.uD = Function(M) + + # Boundary conditions + bcs = [DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u']] + + # Solver for u, D + uD_problem = LinearVariationalProblem(aeqn, Leqn, self.uD, bcs=bcs) + + # Provide callback for the nullspace of the trace system + def trace_nullsp(T): + return VectorSpaceBasis(constant=True) + + appctx = {"trace_nullspace": trace_nullsp} + self.uD_solver = LinearVariationalSolver(uD_problem, + solver_parameters=self.solver_parameters, + appctx=appctx) + + # Log residuals on hybridized solver + self.log_ksp_residuals(self.uD_solver.snes.ksp) + + @timed_function("Gusto:LinearSolve") + def solve(self, xrhs, dy): + """ + Solve the linear problem. + + Args: + xrhs (:class:`Function`): the right-hand side field in the + appropriate :class:`MixedFunctionSpace`. + dy (:class:`Function`): the resulting field in the appropriate + :class:`MixedFunctionSpace`. + """ + self.xrhs.assign(xrhs) + + with timed_region("Gusto:VelocityDepthSolve"): + logger.info('Moist convective linear solver: mixed solve') + self.uD_solver.solve() + + u1, D1 = self.uD.subfunctions + u = dy.subfunctions[0] + D = dy.subfunctions[1] + u.assign(u1) + D.assign(D1)