From 8d2860da5713266c75a012f12e63cc8a87754bb7 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Thu, 7 Aug 2025 14:47:38 +0100 Subject: [PATCH 01/14] Create branch From bc556370de6351663b3e21fd81d54a0554becdc3 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 7 Aug 2025 16:07:05 +0100 Subject: [PATCH 02/14] New moist solver for 5 variables works --- .../moist_convective_williamson_2.py | 14 +- gusto/solvers/linear_solvers.py | 150 +++++++++++++++++- 2 files changed, 160 insertions(+), 4 deletions(-) diff --git a/examples/shallow_water/moist_convective_williamson_2.py b/examples/shallow_water/moist_convective_williamson_2.py index 63038b9f4..151c21f8d 100644 --- a/examples/shallow_water/moist_convective_williamson_2.py +++ b/examples/shallow_water/moist_convective_williamson_2.py @@ -13,14 +13,14 @@ """ from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter -from firedrake import SpatialCoordinate, sin, cos, exp, Function +from firedrake import SpatialCoordinate, sin, cos, exp, Function, errornorm, norm from gusto import ( Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations, ZonalComponent, MeridionalComponent, SteadyStateError, lonlatr_from_xyz, DG1Limiter, InstantRain, MoistConvectiveSWSolver, ForwardEuler, RelativeVorticity, SWSaturationAdjustment, WaterVapour, CloudWater, Rain, - GeneralIcosahedralSphereMesh, xyz_vector_from_lonlatr + GeneralIcosahedralSphereMesh, xyz_vector_from_lonlatr, MoistConvectiveSWSolverNew ) moist_convect_williamson_2_defaults = { @@ -127,7 +127,7 @@ def sat_func(x_in): SSPRK3(domain, "rain", limiter=limiter) ] - linear_solver = MoistConvectiveSWSolver(eqns) + linear_solver = MoistConvectiveSWSolverNew(eqns) # Physics schemes sat_adj = SWSaturationAdjustment( @@ -158,6 +158,8 @@ def sat_func(x_in): D0 = stepper.fields("D") v0 = stepper.fields("water_vapour") + D_init = Function(D0.function_space()) + uexpr = xyz_vector_from_lonlatr(u_max*cos(phi), 0, 0, (x, y, z)) g = parameters.g w = Omega*radius*u_max + (u_max**2)/2 @@ -184,6 +186,8 @@ def sat_func(x_in): D0.interpolate(Dexpr) v0.interpolate(vexpr) + D_init.interpolate(Dexpr) + # Set reference profiles Dbar = Function(D0.function_space()).assign(mean_depth) stepper.set_reference_profiles([('D', Dbar)]) @@ -194,6 +198,10 @@ def sat_func(x_in): stepper.run(t=0, tmax=tmax) + error_D = errornorm(D0, D_init, mesh=mesh)/norm(D_init, mesh=mesh) + + print("Error D:", error_D) + # ---------------------------------------------------------------------------- # # MAIN # ---------------------------------------------------------------------------- # diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py index 213d09df2..67d218730 100644 --- a/gusto/solvers/linear_solvers.py +++ b/gusto/solvers/linear_solvers.py @@ -30,7 +30,7 @@ __all__ = ["BoussinesqSolver", "LinearTimesteppingSolver", "CompressibleSolver", - "ThermalSWSolver", "MoistConvectiveSWSolver"] + "ThermalSWSolver", "MoistConvectiveSWSolver", "MoistConvectiveSWSolverNew"] class TimesteppingSolver(object, metaclass=ABCMeta): @@ -970,3 +970,151 @@ def solve(self, xrhs, dy): D = dy.subfunctions[1] u.assign(u1) D.assign(D1) + +class MoistConvectiveSWSolverNew(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', + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "additive", + "pc_fieldsplit_0_fields": "0,1", + "pc_fieldsplit_1_fields": "2,3,4", + 'mat_type': 'matfree', + "fieldsplit_0": { + 'ksp_type': 'preonly', + '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' + } + } + }, + "fieldsplit_1": { + "ksp_type": "preonly", + "pc_type": "none" + }, + } + + @timed_function("Gusto:SolverSetup") + def _setup_solver(self): + equation = self.equations # just cutting down line length a bit + dt = self.dt + + # Set up Preconditioner aP + beta_u = dt*self.tau_values.get("u", self.alpha) + beta_d = dt*self.tau_values.get("D", self.alpha) + Vu = equation.domain.spaces("HDiv") + VD = equation.domain.spaces("DG") + + # Split up the rhs vector + self.xrhs = Function(self.equations.function_space) + u_in = split(self.xrhs)[0] + D_in = split(self.xrhs)[1] + p0_in = split(self.xrhs)[2] + p1_in = split(self.xrhs)[3] + p2_in = split(self.xrhs)[4] + + + # Build the reduced function space for u, D + M = self.equations.function_space + w, phi, q0, q1, q2 = TestFunctions(M) + u, D, p0, p1, p2 = TrialFunctions(M) + + # Get background depth + Dbar = split(equation.X_ref)[1] + + g = equation.parameters.g + + eqn = ( + inner(w, (u - u_in)) * dx + - beta_u * (D - Dbar) * div(w*g) * dx + + inner(phi, (D - D_in)) * dx + + beta_d * phi * div(Dbar*u) * dx + + inner(q0, (p0)) * dx + + inner(q1, (p1)) * dx + + inner(q2, (p2)) * dx + ) + + if 'coriolis' in equation.prescribed_fields._field_names: + f = equation.prescribed_fields('coriolis') + eqn += beta_u * f * inner(w, equation.domain.perp(u)) * dx + + aeqn= lhs(eqn) + Leqn = rhs(eqn) + + # # Calculated linearised problem + # residual = equation.residual.label_map( + # lambda t: t.has_label(linearisation), + # lambda t: Term(t.get(linearisation).form, t.labels), + # drop) + + # beta = dt*self.alpha + + # # Split up the rhs vector (symbolically) + # #self.xrhs = Function(M) + + # aeqn = residual.label_map( + # lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), + # map_if_false=lambda t: beta*t) + # Leqn = residual.label_map( + # lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), + # map_if_false=drop) + + # 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, + constant_jacobian=True) + + # 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:UpdateReferenceProfiles") + def update_reference_profiles(self): + self.uD_solver.invalidate_jacobian() + + @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() + + dy.assign(self.uD) From f1d81263c033528e50b696ba6b7c7ebd9be45ac1 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 7 Aug 2025 17:01:23 +0100 Subject: [PATCH 03/14] Temprary save - working toward aP --- gusto/solvers/linear_solvers.py | 59 ++++++++++++++++----------------- 1 file changed, 29 insertions(+), 30 deletions(-) diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py index 67d218730..21ce028af 100644 --- a/gusto/solvers/linear_solvers.py +++ b/gusto/solvers/linear_solvers.py @@ -13,10 +13,11 @@ BrokenElement, FunctionSpace, MixedFunctionSpace, DirichletBC, as_vector, assemble, conditional ) -from firedrake.fml import Term, drop +from firedrake.fml import Term, drop, subject from firedrake.petsc import flatten_parameters from firedrake.__future__ import interpolate from pyop2.profiling import timed_function, timed_region +from ufl import derivative as ufl_derivative from gusto.equations.active_tracers import TracerVariableType from gusto.core.logging import ( @@ -1017,20 +1018,15 @@ def _setup_solver(self): # Set up Preconditioner aP beta_u = dt*self.tau_values.get("u", self.alpha) beta_d = dt*self.tau_values.get("D", self.alpha) - Vu = equation.domain.spaces("HDiv") - VD = equation.domain.spaces("DG") - # Split up the rhs vector - self.xrhs = Function(self.equations.function_space) + # Split up the rhs vectorx + M = self.equations.function_space + self.xrhs = Function(M) u_in = split(self.xrhs)[0] D_in = split(self.xrhs)[1] - p0_in = split(self.xrhs)[2] - p1_in = split(self.xrhs)[3] - p2_in = split(self.xrhs)[4] # Build the reduced function space for u, D - M = self.equations.function_space w, phi, q0, q1, q2 = TestFunctions(M) u, D, p0, p1, p2 = TrialFunctions(M) @@ -1040,9 +1036,9 @@ def _setup_solver(self): g = equation.parameters.g eqn = ( - inner(w, (u - u_in)) * dx + inner(w, (u-u_in)) * dx - beta_u * (D - Dbar) * div(w*g) * dx - + inner(phi, (D - D_in)) * dx + + inner(phi, (D-D_in)) * dx + beta_d * phi * div(Dbar*u) * dx + inner(q0, (p0)) * dx + inner(q1, (p1)) * dx @@ -1053,26 +1049,29 @@ def _setup_solver(self): f = equation.prescribed_fields('coriolis') eqn += beta_u * f * inner(w, equation.domain.perp(u)) * dx - aeqn= lhs(eqn) - Leqn = rhs(eqn) - - # # Calculated linearised problem - # residual = equation.residual.label_map( - # lambda t: t.has_label(linearisation), - # lambda t: Term(t.get(linearisation).form, t.labels), - # drop) - - # beta = dt*self.alpha + aP= lhs(eqn) - # # Split up the rhs vector (symbolically) - # #self.xrhs = Function(M) + # Calculated linearised problem + residual = equation.residual.label_map( + lambda t: t.has_label(linearisation), + lambda t: Term(t.get(linearisation).form, t.labels), + drop) + + other_res = equation.residual.label_map( + lambda t: t.has_label(linearisation), + map_if_true=drop) + other_res = other_res.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=lambda t: Term(ufl_derivative(t.form, t.get(subject)), t.labels), + map_if_false=drop) + beta = dt*self.alpha - # aeqn = residual.label_map( - # lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), - # map_if_false=lambda t: beta*t) - # Leqn = residual.label_map( - # lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), - # map_if_false=drop) + aeqn = residual.label_map( + lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), + map_if_false=lambda t: beta*t) + other_res + Leqn = residual.label_map( + lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), + map_if_false=drop) # Place to put results of (u,D) solver self.uD = Function(M) @@ -1081,7 +1080,7 @@ def _setup_solver(self): 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, + uD_problem = LinearVariationalProblem(aeqn.form, action(Leqn.form, self.xrhs), self.uD, bcs=bcs, constant_jacobian=True) # Provide callback for the nullspace of the trace system From 99ca85b63f275a83b0dbf7c2b13231935b58d79a Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 8 Aug 2025 10:18:09 +0100 Subject: [PATCH 04/14] remove aP from ConvectiveSWESolver --- .../moist_convective_williamson_2.py | 43 +++++++++++++++++-- gusto/solvers/linear_solvers.py | 38 ++++++++-------- 2 files changed, 60 insertions(+), 21 deletions(-) diff --git a/examples/shallow_water/moist_convective_williamson_2.py b/examples/shallow_water/moist_convective_williamson_2.py index 151c21f8d..f7ce33344 100644 --- a/examples/shallow_water/moist_convective_williamson_2.py +++ b/examples/shallow_water/moist_convective_williamson_2.py @@ -20,8 +20,10 @@ ZonalComponent, MeridionalComponent, SteadyStateError, lonlatr_from_xyz, DG1Limiter, InstantRain, MoistConvectiveSWSolver, ForwardEuler, RelativeVorticity, SWSaturationAdjustment, WaterVapour, CloudWater, Rain, - GeneralIcosahedralSphereMesh, xyz_vector_from_lonlatr, MoistConvectiveSWSolverNew + GeneralIcosahedralSphereMesh, xyz_vector_from_lonlatr, + MoistConvectiveSWSolver, MoistConvectiveSWSolverNew ) +from gusto.core.labels import time_derivative, prognostic, coriolis, pressure_gradient, transport moist_convect_williamson_2_defaults = { 'ncells_per_edge': 16, # number of cells per icosahedron edge @@ -84,9 +86,19 @@ def moist_convect_williamson_2( WaterVapour(space='DG'), CloudWater(space='DG'), Rain(space='DG') ] + def linearisation_map(term): + if (term.get(prognostic) in ['u', 'D'] + and (any(term.has_label(time_derivative, coriolis, pressure_gradient)) + or (term.get(prognostic) in ['D'] and term.has_label(transport)))): + return True + elif term.has_label(time_derivative): + return True + else: + return False + eqns = ShallowWaterEquations( domain, parameters, fexpr=fexpr, u_transport_option=u_eqn_type, - active_tracers=tracers + active_tracers=tracers, linearisation_map=linearisation_map ) # IO @@ -127,7 +139,32 @@ def sat_func(x_in): SSPRK3(domain, "rain", limiter=limiter) ] - linear_solver = MoistConvectiveSWSolverNew(eqns) + import numpy as np + np.random.seed(6) + + xold = Function(eqns.function_space) + yold = Function(eqns.function_space) + + for xsub in xold.subfunctions[:2]: + with xsub.dat.vec_wo as v: + v.array[:] = np.random.random_sample(v.array.shape) + + xnew = xold.copy(deepcopy=True) + ynew = yold.copy(deepcopy=True) + + old_solver = MoistConvectiveSWSolver(eqns, options_prefix="swe_old") + new_solver = MoistConvectiveSWSolverNew(eqns, options_prefix="swe_new") + + old_solver.solve(xold, yold) + new_solver.solve(xnew, ynew) + + # for i, (xo, xn) in enumerate(zip(xold.subfunctions, xnew.subfunctions)): + # print(f"{i = }: {errornorm(xo, xn)/norm(xo) = :4e}") + for i, (yo, yn) in enumerate(zip(yold.subfunctions, ynew.subfunctions)): + print(f"{i = }| {norm(yo) = :2e} | {norm(yn) = :2e} | {errornorm(yo, yn)/norm(yo) = :4e}") + + from sys import exit; exit() + # Physics schemes sat_adj = SWSaturationAdjustment( diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py index 21ce028af..49ea330d0 100644 --- a/gusto/solvers/linear_solvers.py +++ b/gusto/solvers/linear_solvers.py @@ -38,7 +38,8 @@ class TimesteppingSolver(object, metaclass=ABCMeta): """Base class for timestepping linear solvers for Gusto.""" def __init__(self, equations, alpha=0.5, tau_values=None, - solver_parameters=None, overwrite_solver_parameters=False): + solver_parameters=None, overwrite_solver_parameters=False, + options_prefix=None): """ Args: equations (:class:`PrognosticEquation`): the model's equation. @@ -58,6 +59,7 @@ def __init__(self, equations, alpha=0.5, tau_values=None, self.dt = equations.domain.dt self.alpha = Constant(alpha) self.tau_values = tau_values if tau_values is not None else {} + self.options_prefix = options_prefix if solver_parameters is not None: if not overwrite_solver_parameters: @@ -940,7 +942,7 @@ def trace_nullsp(T): appctx = {"trace_nullspace": trace_nullsp} self.uD_solver = LinearVariationalSolver(uD_problem, solver_parameters=self.solver_parameters, - appctx=appctx) + appctx=appctx, options_prefix=self.options_prefix) # Log residuals on hybridized solver self.log_ksp_residuals(self.uD_solver.snes.ksp) @@ -972,6 +974,7 @@ def solve(self, xrhs, dy): u.assign(u1) D.assign(D1) + class MoistConvectiveSWSolverNew(TimesteppingSolver): """ Linear solver for the moist convective shallow water equations. @@ -982,13 +985,14 @@ class MoistConvectiveSWSolverNew(TimesteppingSolver): """ solver_parameters = { + 'mat_type': 'matfree', 'ksp_type': 'preonly', "pc_type": "fieldsplit", "pc_fieldsplit_type": "additive", "pc_fieldsplit_0_fields": "0,1", "pc_fieldsplit_1_fields": "2,3,4", - 'mat_type': 'matfree', "fieldsplit_0": { + 'ksp_monitor_true_residual': None, 'ksp_type': 'preonly', 'pc_type': 'python', 'pc_python_type': 'firedrake.HybridizationPC', @@ -1036,9 +1040,9 @@ def _setup_solver(self): g = equation.parameters.g eqn = ( - inner(w, (u-u_in)) * dx + inner(w, u) * dx - beta_u * (D - Dbar) * div(w*g) * dx - + inner(phi, (D-D_in)) * dx + + inner(phi, D) * dx + beta_d * phi * div(Dbar*u) * dx + inner(q0, (p0)) * dx + inner(q1, (p1)) * dx @@ -1049,26 +1053,19 @@ def _setup_solver(self): f = equation.prescribed_fields('coriolis') eqn += beta_u * f * inner(w, equation.domain.perp(u)) * dx - aP= lhs(eqn) + aP = lhs(eqn) # Calculated linearised problem residual = equation.residual.label_map( lambda t: t.has_label(linearisation), lambda t: Term(t.get(linearisation).form, t.labels), drop) - - other_res = equation.residual.label_map( - lambda t: t.has_label(linearisation), - map_if_true=drop) - other_res = other_res.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=lambda t: Term(ufl_derivative(t.form, t.get(subject)), t.labels), - map_if_false=drop) + beta = dt*self.alpha aeqn = residual.label_map( lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), - map_if_false=lambda t: beta*t) + other_res + map_if_false=lambda t: beta*t) Leqn = residual.label_map( lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), map_if_false=drop) @@ -1080,8 +1077,9 @@ def _setup_solver(self): 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.form, action(Leqn.form, self.xrhs), self.uD, bcs=bcs, - constant_jacobian=True) + uD_problem = LinearVariationalProblem( + aeqn.form, action(Leqn.form, self.xrhs), + self.uD, bcs=bcs, constant_jacobian=True) # Provide callback for the nullspace of the trace system def trace_nullsp(T): @@ -1090,7 +1088,7 @@ def trace_nullsp(T): appctx = {"trace_nullspace": trace_nullsp} self.uD_solver = LinearVariationalSolver(uD_problem, solver_parameters=self.solver_parameters, - appctx=appctx) + appctx=appctx, options_prefix=self.options_prefix) # Log residuals on hybridized solver self.log_ksp_residuals(self.uD_solver.snes.ksp) @@ -1111,6 +1109,10 @@ def solve(self, xrhs, dy): :class:`MixedFunctionSpace`. """ self.xrhs.assign(xrhs) + for xsub in self.xrhs.subfunctions[2:]: + xsub.zero() + for ysub in self.uD.subfunctions[2:]: + ysub.zero() with timed_region("Gusto:VelocityDepthSolve"): logger.info('Moist convective linear solver: mixed solve') From 9cfe8b1a0b1e2a8a3029139ee50ce06c1e3179d1 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Fri, 8 Aug 2025 14:20:18 +0100 Subject: [PATCH 05/14] update linearisation map and default solver --- gusto/equations/boussinesq_equations.py | 30 +- .../equations/compressible_euler_equations.py | 21 +- gusto/equations/prognostic_equations.py | 11 +- gusto/equations/shallow_water_equations.py | 73 ++-- .../semi_implicit_quasi_newton.py | 355 ++++++++++++++---- 5 files changed, 356 insertions(+), 134 deletions(-) diff --git a/gusto/equations/boussinesq_equations.py b/gusto/equations/boussinesq_equations.py index d94e2a038..dccd35896 100644 --- a/gusto/equations/boussinesq_equations.py +++ b/gusto/equations/boussinesq_equations.py @@ -3,7 +3,7 @@ from firedrake import inner, dx, div, cross, split, as_vector from firedrake.fml import subject from gusto.core.labels import ( - time_derivative, transport, prognostic, linearisation, + time_derivative, prognostic, linearisation, pressure_gradient, coriolis, divergence, gravity, incompressible ) from gusto.equations.common_forms import ( @@ -60,8 +60,8 @@ def __init__(self, domain, parameters, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation includes time derivatives and - scalar transport terms. + in which case the linearisation drops terms for any active + tracers. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form' and @@ -90,14 +90,11 @@ def __init__(self, domain, parameters, active_tracers = [] if linearisation_map == 'default': - # Default linearisation is time derivatives, scalar transport, - # pressure gradient, gravity and divergence terms - # Don't include active tracers - linearisation_map = lambda t: \ - t.get(prognostic) in ['u', 'p', 'b'] \ - and (any(t.has_label(time_derivative, pressure_gradient, - divergence, gravity)) - or (t.get(prognostic) not in ['u', 'p'] and t.has_label(transport))) + # Default linearisation is to include all terms + # Don't include terms for active tracers + linearisation_map = lambda t: ( + t.has_label(time_derivative) or t.get(prognostic) in field_names + ) super().__init__(field_names, domain, space_names, linearisation_map=linearisation_map, @@ -249,8 +246,8 @@ def __init__(self, domain, parameters, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation includes time derivatives and - scalar transport terms. + in which case the linearisation drops terms for any active + tracers. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form' and @@ -267,13 +264,6 @@ def __init__(self, domain, parameters, NotImplementedError: active tracers are not implemented. """ - if linearisation_map == 'default': - # Default linearisation is time derivatives, pressure gradient, - # Coriolis and transport term from depth equation - linearisation_map = lambda t: \ - (any(t.has_label(time_derivative, pressure_gradient, coriolis, - gravity, divergence, incompressible)) - or (t.get(prognostic) in ['p', 'b'] and t.has_label(transport))) super().__init__(domain=domain, parameters=parameters, compressible=compressible, diff --git a/gusto/equations/compressible_euler_equations.py b/gusto/equations/compressible_euler_equations.py index d512432cf..02a60d7ea 100644 --- a/gusto/equations/compressible_euler_equations.py +++ b/gusto/equations/compressible_euler_equations.py @@ -6,7 +6,7 @@ ) from firedrake.fml import subject, replace_subject from gusto.core.labels import ( - time_derivative, transport, prognostic, hydrostatic, linearisation, + time_derivative, prognostic, hydrostatic, linearisation, pressure_gradient, coriolis, gravity, sponge ) from gusto.equations.thermodynamics import exner_pressure @@ -58,8 +58,8 @@ def __init__(self, domain, parameters, sponge_options=None, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation includes time derivatives and - scalar transport terms. + in which case the linearisation drops terms for any active + tracers. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form' and @@ -88,12 +88,11 @@ def __init__(self, domain, parameters, sponge_options=None, active_tracers = [] if linearisation_map == 'default': - # Default linearisation is time derivatives and scalar transport terms - # Don't include active tracers - linearisation_map = lambda t: \ - t.get(prognostic) in ['u', 'rho', 'theta'] \ - and (t.has_label(time_derivative) - or (t.get(prognostic) != 'u' and t.has_label(transport))) + # Default linearisation is to include all terms + # Don't include terms for active tracers + linearisation_map = lambda t: ( + t.has_label(time_derivative) or t.get(prognostic) in field_names + ) super().__init__(field_names, domain, space_names, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, @@ -306,8 +305,8 @@ def __init__(self, domain, parameters, sponge_options=None, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation includes time derivatives and - scalar transport terms. + in which case the linearisation drops terms for any active + tracers. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form' and diff --git a/gusto/equations/prognostic_equations.py b/gusto/equations/prognostic_equations.py index a562795d6..1f6035b71 100644 --- a/gusto/equations/prognostic_equations.py +++ b/gusto/equations/prognostic_equations.py @@ -211,8 +211,15 @@ def should_linearise(term): # Linearise a term, and add the linearisation as a label def linearise(term, X, X_ref, du): - linear_term = Term(action(ufl.derivative(term.form, X), du), term.labels) - return linearisation(term, replace_subject(X_ref)(linear_term)) + try: + derivative = ufl.derivative(term.form, X) + linear_term = Term(action(derivative, du), term.labels) + return linearisation(term, replace_subject(X_ref)(linear_term)) + except IndexError: + # If the derivative is zero (e.g. for the gravity term) then + # action(derivative, du) will raise an IndexError + # In that case, just return the term with no linearisation label + return term # Add linearisations to all terms that need linearising residual = residual.label_map( diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py index 0e9aebe2a..829693f29 100644 --- a/gusto/equations/shallow_water_equations.py +++ b/gusto/equations/shallow_water_equations.py @@ -3,7 +3,7 @@ from firedrake import (inner, dx, div, FunctionSpace, FacetNormal, jump, avg, dS, split, conditional, exp) from firedrake.fml import subject, drop -from gusto.core.labels import (time_derivative, transport, prognostic, +from gusto.core.labels import (time_derivative, prognostic, linearisation, pressure_gradient, coriolis) from gusto.equations.common_forms import ( advection_form, advection_form_1d, continuity_form, @@ -51,8 +51,8 @@ def __init__(self, domain, parameters, fexpr=None, topog_expr=None, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation includes both time derivatives, - the 'D' transport term and the pressure gradient term. + in which case the linearisation drops terms for any active + tracers. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form', and @@ -70,12 +70,11 @@ def __init__(self, domain, parameters, fexpr=None, topog_expr=None, active_tracers = [] if linearisation_map == 'default': - # Default linearisation is time derivatives, pressure gradient and - # transport term from depth equation. Don't include active tracers - linearisation_map = lambda t: \ - t.get(prognostic) in ['u', 'D'] \ - and (any(t.has_label(time_derivative, coriolis, pressure_gradient)) - or (t.get(prognostic) in ['D'] and t.has_label(transport))) + # Default linearisation is to include all terms + # Don't include terms for active tracers + linearisation_map = lambda t: ( + t.has_label(time_derivative) or t.get(prognostic) in field_names + ) field_names = ['u', 'D'] space_names = {'u': 'HDiv', 'D': 'L2'} @@ -232,8 +231,8 @@ def __init__(self, domain, parameters, fexpr=None, topog_expr=None, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation includes both time derivatives, - the 'D' transport term, pressure gradient and Coriolis terms. + in which case the linearisation drops terms for any active + tracers. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form' and @@ -303,8 +302,8 @@ def __init__(self, domain, parameters, equivalent_buoyancy=False, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation includes both time derivatives, - the 'D' transport term and the pressure gradient term. + in which case the linearisation drops terms for any active + tracers. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form', and @@ -335,19 +334,21 @@ def __init__(self, domain, parameters, equivalent_buoyancy=False, active_tracers = [] if linearisation_map == 'default': - # Default linearisation is time derivatives, pressure - # gradient and transport terms from depth and buoyancy - # equations. Include q_t if equivalent buoyancy. Don't include - # active tracers. - linear_transported = ['D', self.b_name] if equivalent_buoyancy: - linear_transported.append('q_t') - linearisation_map = lambda t: \ - t.get(prognostic) in field_names \ - and (any(t.has_label(time_derivative, pressure_gradient, - coriolis)) - or (t.get(prognostic) in linear_transported - and t.has_label(transport))) + # Default linearisation is to include all terms + # Don't include qt terms + linearisation_map = lambda t: ( + t.has_label(time_derivative) + or t.get(prognostic) in ['u', 'D', 'b_e'] + ) + + else: + # Default linearisation is to include all terms + # Don't include terms for active tracers + linearisation_map = lambda t: ( + t.has_label(time_derivative) + or t.get(prognostic) in field_names + ) # Bypass ShallowWaterEquations.__init__ to avoid having to # define the field_names separately @@ -562,8 +563,8 @@ def __init__(self, domain, parameters, equivalent_buoyancy=False, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation includes both time derivatives, - the 'D' transport term, pressure gradient and Coriolis terms. + in which case the linearisation drops terms for any active + tracers. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form' and @@ -614,8 +615,8 @@ class ShallowWaterEquations_1d(PrognosticEquationSet): linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation includes both time derivatives, - the 'D' transport term, pressure gradient and Coriolis terms. + in which case the linearisation drops terms for any active + tracers. no_normal_flow_bc_ids (list, optional): a list of IDs of domain boundaries at which no normal flow will be enforced. Defaults to None. @@ -637,11 +638,11 @@ def __init__(self, domain, parameters, raise NotImplementedError('Tracers not implemented for 1D shallow water equations') if linearisation_map == 'default': - # Default linearisation is time derivatives, pressure gradient, - # Coriolis and transport term from depth equation - linearisation_map = lambda t: \ - (any(t.has_label(time_derivative, pressure_gradient, coriolis)) - or (t.get(prognostic) == 'D' and t.has_label(transport))) + # Default linearisation is to include all terms + # Don't include terms for active tracers + linearisation_map = lambda t: ( + t.has_label(time_derivative) or t.get(prognostic) in field_names + ) super().__init__(field_names, domain, space_names, linearisation_map=linearisation_map, @@ -743,8 +744,8 @@ def __init__(self, domain, parameters, fexpr=None, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation includes both time derivatives, - the 'D' transport term, pressure gradient and Coriolis terms. + in which case the linearisation drops terms for any active + tracers. no_normal_flow_bc_ids (list, optional): a list of IDs of domain boundaries at which no normal flow will be enforced. Defaults to None. diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py index d8ca50194..9d0b057f4 100644 --- a/gusto/timestepping/semi_implicit_quasi_newton.py +++ b/gusto/timestepping/semi_implicit_quasi_newton.py @@ -5,16 +5,18 @@ from firedrake import ( Function, Constant, TrialFunctions, DirichletBC, div, assemble, - LinearVariationalProblem, LinearVariationalSolver, FunctionSpace + LinearVariationalProblem, LinearVariationalSolver, FunctionSpace, action ) -from firedrake.fml import drop, replace_subject +from firedrake.fml import drop, replace_subject, Term from firedrake.__future__ import interpolate -from pyop2.profiling import timed_stage +from firedrake.petsc import flatten_parameters +from pyop2.profiling import timed_stage, timed_function from gusto.core import TimeLevelFields, StateFields -from gusto.core.labels import (transport, diffusion, time_derivative, - hydrostatic, physics_label, sponge, - incompressible) -from gusto.solvers import LinearTimesteppingSolver, mass_parameters +from gusto.core.labels import ( + transport, diffusion, time_derivative, hydrostatic, physics_label, sponge, + incompressible, linearisation, prognostic +) +from gusto.solvers import mass_parameters from gusto.core.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation from gusto.timestepping.timestepper import BaseTimestepper @@ -37,10 +39,12 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, auxiliary_equations_and_schemes=None, linear_solver=None, diffusion_schemes=None, physics_schemes=None, slow_physics_schemes=None, fast_physics_schemes=None, - alpha=0.5, off_centred_u=False, + alpha=0.5, tau_values=None, off_centred_u=False, num_outer=2, num_inner=2, accelerator=True, predictor=None, reference_update_freq=None, - spinup_steps=0): + spinup_steps=0, solver_prognostics=None, + linear_solver_parameters=None, + overwrite_linear_solver_parameters=False): """ Args: equation_set (:class:`PrognosticEquationSet`): the prognostic @@ -76,6 +80,8 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, alpha (`float, optional): the semi-implicit off-centering parameter. A value of 1 corresponds to fully implicit, while 0 corresponds to fully explicit. Defaults to 0.5. + tau_values (dict, optional): contains the semi-implicit relaxation + parameters. Defaults to None, in which case the value of alpha is used. off_centred_u (bool, optional): option to offcentre the transporting velocity. Defaults to False, in which case transporting velocity is centred. If True offcentring uses value of alpha. @@ -110,8 +116,18 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, spinup_steps (int, optional): the number of steps to run the model in "spin-up" mode, where the alpha parameter is set to 1.0. Defaults to 0, which corresponds to no spin-up. + solver_prognostics (list, optional): a list of the names of + prognostic variables to include in the solver. + linear_solver_parameters (dict, optional): contains the options to + be passed to the underlying :class:`LinearVariationalSolver`. + for the Defaults to None. + overwrite_solver_parameters (bool, optional): if True use only the + `solver_parameters` that have been passed in. If False then + update the default parameters with the `solver_parameters` + passed in. Defaults to False. """ + # Basic parts of the SIQN structure ------------------------------------ self.num_outer = num_outer self.num_inner = num_inner mesh = equation_set.domain.mesh @@ -119,8 +135,58 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, self.alpha = Function(R, val=float(alpha)) self.predictor = predictor self.accelerator = accelerator + self.implicit_terms = [incompressible, sponge] + + # default is to not offcentre transporting velocity but if it + # is offcentred then use the same value as alpha + if off_centred_u: + self.alpha_u = Function(R, val=float(alpha)) + else: + self.alpha_u = Function(R, val=0.5) + + # Set up some important fields ----------------------------------------- + self.field_name = equation_set.field_name + W = equation_set.function_space + self.xrhs = Function(W) + self.xrhs_phys = Function(W) + self.dy = Function(W) + + # Determine prognostics for solver ------------------------------------- + self.non_solver_prognostics = [] + if solver_prognostics is not None: + assert type(solver_prognostics) is list, ( + "solver_prognostics should be a list of prognostic variable " + + f"names, and not {type(solver_prognostics)}" + ) + for prognostic_name in solver_prognostics: + assert prognostic_name in equation_set.field_names, ( + f"Prognostic variable {prognostic_name} not found in " + + "equation set field names" + ) + for field_name in equation_set.field_names: + if field_name not in solver_prognostics: + self.non_solver_prognostics.append(field_name) + else: + solver_prognostics = equation_set.field_names + logger.info( + 'Semi-implicit Quasi-Newton: Using solver prognostics ' + + f'{solver_prognostics}' + ) + + # BCs, Forcing and Linear Solver --------------------------------------- + self.bcs = equation_set.bcs + self.forcing = Forcing(equation_set, self.implicit_terms, self.alpha) + if linear_solver is None: + self.linear_solver = SIQNLinearSolver( + equation_set, solver_prognostics, self.implicit_terms, + self.alpha, tau_values=tau_values, + solver_parameters=linear_solver_parameters, + overwrite_solver_parameters=overwrite_linear_solver_parameters + ) + else: + self.linear_solver = linear_solver - # Options relating to reference profiles and spin-up + # Options relating to reference profiles and spin-up ------------------- self._alpha_original = float(alpha) # float so as to not upset adjoint self.reference_update_freq = reference_update_freq self.to_update_ref_profile = False @@ -128,43 +194,18 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, self.spinup_begun = False self.spinup_done = False - # Flag for if we have simultaneous transport - self.simult = False - - # default is to not offcentre transporting velocity but if it - # is offcentred then use the same value as alpha - self.alpha_u = Function(R, val=float(alpha)) if off_centred_u else Function(R, val=0.5) - + # Transport ------------------------------------------------------------ self.spatial_methods = spatial_methods - if physics_schemes is not None: - self.final_physics_schemes = physics_schemes - else: - self.final_physics_schemes = [] - if slow_physics_schemes is not None: - self.slow_physics_schemes = slow_physics_schemes - else: - self.slow_physics_schemes = [] - if fast_physics_schemes is not None: - self.fast_physics_schemes = fast_physics_schemes - else: - self.fast_physics_schemes = [] - self.all_physics_schemes = (self.slow_physics_schemes - + self.fast_physics_schemes - + self.final_physics_schemes) - - for parametrisation, scheme in self.all_physics_schemes: - assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" - if hasattr(parametrisation, "explicit_only") and parametrisation.explicit_only: - assert isinstance(scheme, ExplicitTimeDiscretisation), \ - ("Only explicit time discretisations can be used with " - + f"physics scheme {parametrisation.label.label}") + # Flag for if we have simultaneous transport of tracers + self.simult = False self.active_transport = [] for scheme in transport_schemes: - assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" + assert scheme.nlevels == 1, \ + "multilevel schemes not supported as part of this timestepping loop" if isinstance(scheme.field_name, list): - # This means that multiple fields are being transported simultaneously + # This means that multiple fields are transported simultaneously self.simult = True for subfield in scheme.field_name: assert subfield in equation_set.field_names @@ -175,7 +216,8 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, for method in spatial_methods: if subfield == method.variable and method.term_label == transport: method_found = True - assert method_found, f'No transport method found for variable {scheme.field_name}' + assert method_found, \ + f'No transport method found for variable {scheme.field_name}' self.active_transport.append((scheme.field_name, scheme)) else: assert scheme.field_name in equation_set.field_names @@ -186,12 +228,40 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, if scheme.field_name == method.variable and method.term_label == transport: method_found = True self.active_transport.append((scheme.field_name, scheme)) - assert method_found, f'No transport method found for variable {scheme.field_name}' + assert method_found, \ + f'No transport method found for variable {scheme.field_name}' + + # Physics schemes ------------------------------------------------------ + if physics_schemes is not None: + self.final_physics_schemes = physics_schemes + else: + self.final_physics_schemes = [] + if slow_physics_schemes is not None: + self.slow_physics_schemes = slow_physics_schemes + else: + self.slow_physics_schemes = [] + if fast_physics_schemes is not None: + self.fast_physics_schemes = fast_physics_schemes + else: + self.fast_physics_schemes = [] + self.all_physics_schemes = (self.slow_physics_schemes + + self.fast_physics_schemes + + self.final_physics_schemes) + for parametrisation, scheme in self.all_physics_schemes: + assert scheme.nlevels == 1, \ + "multilevel schemes not supported as part of this timestepping loop" + if hasattr(parametrisation, "explicit_only") and parametrisation.explicit_only: + assert isinstance(scheme, ExplicitTimeDiscretisation), \ + ("Only explicit time discretisations can be used with " + + f"physics scheme {parametrisation.label.label}") + + # Diffusion and other terms -------------------------------------------- self.diffusion_schemes = [] if diffusion_schemes is not None: for scheme in diffusion_schemes: - assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" + assert scheme.nlevels == 1, \ + "multilevel schemes not supported as part of this timestepping loop" assert scheme.field_name in equation_set.field_names self.diffusion_schemes.append((scheme.field_name, scheme)) # Check that there is a corresponding transport method @@ -203,7 +273,8 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, if auxiliary_equations_and_schemes is not None: for eqn, scheme in auxiliary_equations_and_schemes: - assert not hasattr(eqn, "field_names"), 'Cannot use auxiliary schemes with multiple fields' + assert not hasattr(eqn, "field_names"), \ + 'Cannot use auxiliary schemes with multiple fields' self.auxiliary_schemes = [ (eqn.field_name, scheme) for eqn, scheme in auxiliary_equations_and_schemes] @@ -213,30 +284,20 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, self.auxiliary_schemes = [] self.auxiliary_equations_and_schemes = auxiliary_equations_and_schemes + # General set up ------------------------------------------------------- + # This sets up all of the time discretisations for transport, diffusion, + # physics schemes, etc, so must happen after those have been worked out super().__init__(equation_set, io) + # Set up auxiliary equations separately, as these are not done as part + # of the super init for aux_eqn, aux_scheme in self.auxiliary_equations_and_schemes: self.setup_equation(aux_eqn) aux_scheme.setup(aux_eqn) self.setup_transporting_velocity(aux_scheme) - self.tracers_to_copy = [] - if equation_set.active_tracers is not None: - for active_tracer in equation_set.active_tracers: - self.tracers_to_copy.append(active_tracer.name) - - self.field_name = equation_set.field_name - W = equation_set.function_space - self.xrhs = Function(W) - self.xrhs_phys = Function(W) - self.dy = Function(W) - if linear_solver is None: - self.linear_solver = LinearTimesteppingSolver(equation_set, self.alpha) - else: - self.linear_solver = linear_solver - self.forcing = Forcing(equation_set, self.alpha) - self.bcs = equation_set.bcs - + # Set up the predictor last, as it requires the state fields to already + # be created if self.predictor is not None: V_DG = equation_set.domain.spaces('DG') self.predictor_field_in = Function(V_DG) @@ -307,7 +368,7 @@ def copy_active_tracers(self, x_in, x_out): x_out: The output set of fields """ - for name in self.tracers_to_copy: + for name in self.non_solver_prognostics: x_out(name).assign(x_in(name)) def transport_fields(self, outer, xstar, xp): @@ -535,18 +596,19 @@ class Forcing(object): semi-implicit time discretisation. """ - def __init__(self, equation, alpha): + def __init__(self, equation, implicit_terms, alpha): """ Args: equation (:class:`PrognosticEquationSet`): the prognostic equations containing the forcing terms. + implicit_terms (list): list of labels for terms that are always + fully-implicit. alpha (:class:`Function`): semi-implicit off-centering factor. An alpha of 0 corresponds to fully explicit, while a factor of 1 corresponds to fully implicit. """ self.field_name = equation.field_name - implicit_terms = [incompressible, sponge] dt = equation.domain.dt W = equation.function_space @@ -678,3 +740,166 @@ def zero_forcing_terms(self, equation, x_in, x_out, field_names): logger.debug(f'Semi-Implicit Quasi Newton: Zeroing implicit forcing for {field_name}') field_index = equation.field_names.index(field_name) x_out.subfunctions[field_index].assign(x_in(field_name)) + + +class SIQNLinearSolver(object): + """ + Sets up the linear solver for the Semi-Implicit Quasi-Newton timestepper. + + This computes the approximate linearisation of the algorithm for computing + the (n+1)-th level state, and sets up matrix-vector problem for the + Quasi-Newton problem. + """ + + 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'}} + } + + solver_parameters = { + 'ksp_error_if_not_converged': None, + 'pc_factor_shift_type': 'nonzero', + 'snes_type': 'ksponly', + } + + def __init__(self, equation, solver_prognostics, implicit_terms, + alpha, tau_values=None, solver_parameters=None, + overwrite_solver_parameters=False, reference_dependent=True): + """ + Args: + equations (:class:`PrognosticEquation`): the model's equation. + solver_prognostics (list): a list of prognostic variable names to + include in the linear solver. + implicit_terms (list): a list of labels for terms that are always + fully-implicit. + alpha (float, optional): the semi-implicit off-centring factor. + Defaults to 0.5. A value of 1 is fully-implicit. + tau_values (dict, optional): contains the semi-implicit relaxation + parameters. Defaults to None, in which case the value of alpha + is used. + solver_parameters (dict, optional): contains the options to be + passed to the underlying :class:`LinearVariationalSolver`. + Defaults to None. + overwrite_solver_parameters (bool, optional): if True use only the + `solver_parameters` that have been passed in. If False then + update the default parameters with the `solver_parameters` + passed in. Defaults to False. + reference_dependent: this indicates that the solver Jacobian should + be rebuilt if the reference profiles have been updated. + """ + + # Update or set solver parameters -------------------------------------- + if solver_parameters is not None: + if not overwrite_solver_parameters: + p = flatten_parameters(self.solver_parameters) + p.update(flatten_parameters(solver_parameters)) + solver_parameters = p + self.solver_parameters = solver_parameters + + if logger.isEnabledFor(DEBUG): + self.solver_parameters["ksp_monitor_true_residual"] = None + + dt = equation.domain.dt + self.reference_dependent = reference_dependent + self.alpha = Constant(alpha) + self.tau_values = tau_values if tau_values is not None else {} + + for prognostic_name in solver_prognostics: + if prognostic_name not in self.tau_values.keys(): + # If no tau value is specified for this prognostic, use alpha + self.tau_values[prognostic_name] = self.alpha + + # Set up the linearisation of the equation ----------------------------- + # Time derivative terms: all prognostic variables + residual = equation.residual.label_map( + lambda t: t.has_label(time_derivative), + lambda t: Term(t.get(linearisation).form, t.labels), + drop + ) + + # All other linear terms (only specified prognostics) + for prognostic_name in solver_prognostics: + # Add semi-implicit terms for this prognostic + tau = self.tau_values[prognostic_name] + residual += equation.residual.label_map( + lambda t: ( + t.has_label(linearisation) + and not t.has_label(time_derivative) + and t.get(prognostic) == prognostic_name + and not any(t.has_label(*implicit_terms, return_tuple=True)) + ), + lambda t: tau*dt*Term(t.get(linearisation).form, t.labels), + drop + ) + # Add implicit terms for this prognostic + residual += equation.residual.label_map( + lambda t: ( + t.has_label(linearisation) + and not t.has_label(time_derivative) + and t.get(prognostic) == prognostic_name + and any(t.has_label(*implicit_terms, return_tuple=True)) + ), + lambda t: dt*Term(t.get(linearisation).form, t.labels), + drop + ) + + W = equation.function_space + + # Split up the rhs vector (symbolically) + self.xrhs = Function(W) + + aeqn = residual + Leqn = residual.label_map( + lambda t: t.has_label(time_derivative), map_if_false=drop + ) + + # Place to put result of solver + self.dy = Function(W) + + # Set up the solver ---------------------------------------------------- + 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, + constant_jacobian=True) + + self.solver = LinearVariationalSolver( + problem, solver_parameters=self.solver_parameters, + options_prefix='linear_solver' + ) + + @staticmethod + def log_ksp_residuals(ksp): + if logger.isEnabledFor(DEBUG): + ksp.setMonitor(logging_ksp_monitor_true_residual) + + @timed_function("Gusto:UpdateReferenceProfiles") + def update_reference_profiles(self): + if self.reference_dependent: + self.solver.invalidate_jacobian() + + @timed_function("Gusto:LinearSolve") + def solve(self, xrhs, dy): + """ + Solve the linear problem. + + Args: + xrhs (:class:`Function`): the right-hand side residual field in the + appropriate :class:`MixedFunctionSpace`. + dy (:class:`Function`): the resulting increment field in the + appropriate :class:`MixedFunctionSpace`. + """ + self.xrhs.assign(xrhs) + self.solver.solve() + dy.assign(self.dy) From 807c7c01eeec1da5d60284c97d214de7e2e79a6b Mon Sep 17 00:00:00 2001 From: atb1995 Date: Fri, 8 Aug 2025 14:34:35 +0100 Subject: [PATCH 06/14] Saving progress before merge in of linearisation map --- .../shallow_water/moist_thermal_equivb_gw.py | 48 ++- gusto/equations/shallow_water_equations.py | 2 + gusto/solvers/linear_solvers.py | 304 +++++++++++++++++- 3 files changed, 350 insertions(+), 4 deletions(-) diff --git a/examples/shallow_water/moist_thermal_equivb_gw.py b/examples/shallow_water/moist_thermal_equivb_gw.py index f2928461c..d876ff41d 100644 --- a/examples/shallow_water/moist_thermal_equivb_gw.py +++ b/examples/shallow_water/moist_thermal_equivb_gw.py @@ -5,14 +5,17 @@ from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter from firedrake import ( - SpatialCoordinate, pi, sqrt, min_value, cos, Constant, Function, exp, sin + SpatialCoordinate, pi, sqrt, min_value, cos, Constant, Function, exp, sin, + norm, errornorm ) from gusto import ( Domain, IO, OutputParameters, DGUpwind, ShallowWaterParameters, ThermalShallowWaterEquations, lonlatr_from_xyz, MeridionalComponent, GeneralIcosahedralSphereMesh, SubcyclingOptions, ZonalComponent, PartitionedCloud, RungeKuttaFormulation, SSPRK3, ThermalSWSolver, - SemiImplicitQuasiNewton, xyz_vector_from_lonlatr + SemiImplicitQuasiNewton, xyz_vector_from_lonlatr, ThermalSWSolverNew, + ThermalSWSolverMono, prognostic, time_derivative, coriolis, + pressure_gradient, transport, drop ) moist_thermal_gw_defaults = { @@ -63,11 +66,17 @@ def moist_thermal_gw( nu=nu, beta2=beta2) Omega = parameters.Omega fexpr = 2*Omega*xyz[2]/radius + # def linearisation_map(term): + # if term.has_label(time_derivative): + # return True + # else: + # return False + linearisation_map = lambda t: t.has_label(time_derivative) # Equation eqns = ThermalShallowWaterEquations( domain, parameters, fexpr=fexpr, - equivalent_buoyancy=True + equivalent_buoyancy=False, linearisation_map = linearisation_map ) # IO @@ -86,6 +95,39 @@ def moist_thermal_gw( linear_solver = ThermalSWSolver(eqns) + import numpy as np + np.random.seed(6) + + xold = Function(eqns.function_space) + yold = Function(eqns.function_space) + + + for xsub in xold.subfunctions[:3]: + with xsub.dat.vec_wo as v: + v.array[:] = np.random.random_sample(v.array.shape) + + xnew = xold.copy(deepcopy=True) + ynew = yold.copy(deepcopy=True) + xmono = xold.copy(deepcopy=True) + ymono = yold.copy(deepcopy=True) + + # old_solver = ThermalSWSolver(eqns, options_prefix="swe_old") + # new_solver = ThermalSWSolverNew(eqns, options_prefix="swe_new") + mono_solver = ThermalSWSolverMono(eqns, alpha=0.5) + + # old_solver.solve(xold, yold) + # new_solver.solve(xnew, ynew) + mono_solver.solve(xmono, ymono) + + # for i, (xo, xn) in enumerate(zip(xold.subfunctions, xnew.subfunctions)): + # print(f"{i = }: {errornorm(xo, xn)/norm(xo) = :4e}") + for i, (yo, yn) in enumerate(zip(yold.subfunctions, ynew.subfunctions)): + print(f"Hybrid New {i = }| {norm(yo) = :2e} | {norm(yn) = :2e} | {errornorm(yo, yn)/norm(yo) = :4e}") + for i, (yo, ym) in enumerate(zip(yold.subfunctions, ymono.subfunctions)): + print(f"Monolithic {i = }| {norm(yo) = :2e} | {norm(ym) = :2e} | {errornorm(yo, ym)/norm(yo) = :4e}") + + from sys import exit; exit() + # ------------------------------------------------------------------------ # # Timestepper # ------------------------------------------------------------------------ # diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py index 0e9aebe2a..51fe39d34 100644 --- a/gusto/equations/shallow_water_equations.py +++ b/gusto/equations/shallow_water_equations.py @@ -334,6 +334,8 @@ def __init__(self, domain, parameters, equivalent_buoyancy=False, if active_tracers is None: active_tracers = [] + breakpoint() + if linearisation_map == 'default': # Default linearisation is time derivatives, pressure # gradient and transport terms from depth and buoyancy diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py index 49ea330d0..18676ebbc 100644 --- a/gusto/solvers/linear_solvers.py +++ b/gusto/solvers/linear_solvers.py @@ -31,7 +31,8 @@ __all__ = ["BoussinesqSolver", "LinearTimesteppingSolver", "CompressibleSolver", - "ThermalSWSolver", "MoistConvectiveSWSolver", "MoistConvectiveSWSolverNew"] + "ThermalSWSolver", "ThermalSWSolverNew", "ThermalSWSolverMono", + "MoistConvectiveSWSolver", "MoistConvectiveSWSolverNew"] class TimesteppingSolver(object, metaclass=ABCMeta): @@ -773,6 +774,307 @@ def solve(self, xrhs, dy): b.assign(self.b) +class ThermalSWSolverNew(TimesteppingSolver): + """Linear solver object for the thermal shallow water equations. + + This solves a linear problem for the thermal shallow water + equations with prognostic variables u (velocity), D (depth) and + either b (buoyancy) or b_e (equivalent buoyancy). It follows the + following strategy: + + (1) Eliminate b + (2) Solve the resulting system for (u, D) using a hybrid-mixed method + (3) Reconstruct b + + """ + + # 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'}} + # } + + + solver_parameters = { + 'mat_type': 'matfree', + 'ksp_type': 'preonly', + "ksp_view":":ksp_view.log", + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "multiplicative", + "pc_fieldsplit_0_fields": "0,1", + "pc_fieldsplit_1_fields": "2", + "fieldsplit_0": { + 'ksp_monitor_true_residual': None, + 'ksp_type': 'preonly', + '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' + } + } + }, + "fieldsplit_L2":{ + 'ksp_monitor': None, + "ksp_type": "preonly", + "pc_type": "python", + 'pc_python_type': 'firedrake.AssembledPC', + 'assembled':{'pc_type':'lu'} + }, + } + + @timed_function("Gusto:SolverSetup") + def _setup_solver(self): + equation = self.equations # just cutting down line length a bit + equivalent_buoyancy = equation.equivalent_buoyancy + + dt = self.dt + beta_u = dt*self.tau_values.get("u", self.alpha) + beta_d = dt*self.tau_values.get("D", self.alpha) + beta_b = dt*self.tau_values.get("b", self.alpha) + Vu = equation.domain.spaces("HDiv") + VD = equation.domain.spaces("DG") + Vb = equation.domain.spaces("DG") + + # Check that the third field is buoyancy + if not equation.field_names[2] == 'b' and not (equation.field_names[2] == 'b_e' and equivalent_buoyancy): + raise NotImplementedError("Field 'b' or 'b_e' must exist to use the thermal linear solver in the SIQN scheme") + + # Split up the rhs vector + self.xrhs = Function(equation.function_space) + u_in, D_in, b_in = split(self.xrhs)[0:3] + + # Build the reduced function space for u, D + M = equation.function_space + w, phi, gamma = TestFunctions(M) + u, D, b = TrialFunctions(M) + + # Get background buoyancy and depth + Dbar, bbar = split(equation.X_ref)[1:3] + + # Approximate elimination of b + bval = -dot(u, grad(bbar))*beta_b + b_in + + if equivalent_buoyancy: + # compute q_v using q_sat to partition q_t into q_v and q_c + self.q_sat_func = Function(VD) + self.qvbar = Function(VD) + qtbar = split(equation.X_ref)[3] + + # set up interpolators that use the X_ref values for D and b_e + self.q_sat_expr_interpolate = interpolate( + equation.compute_saturation(equation.X_ref), VD) + self.q_v_interpolate = interpolate( + conditional(qtbar < self.q_sat_func, qtbar, self.q_sat_func), + VD) + + # bbar was be_bar and here we correct to become bbar + bbar += equation.parameters.beta2 * self.qvbar + + n = FacetNormal(equation.domain.mesh) + + eqn = ( + inner(w, (u - u_in)) * dx + - beta_u * (D - Dbar) * div(w*bbar) * dx + + beta_u * jump(w*bbar, n) * avg(D-Dbar) * dS + - beta_u * 0.5 * Dbar * bbar * div(w) * dx + - beta_u * 0.5 * Dbar * bval * div(w) * dx + - beta_u * 0.5 * bbar * div(w*(D-Dbar)) * dx + + beta_u * 0.5 * jump((D-Dbar)*w, n) * avg(bbar) * dS + + inner(phi, (D - D_in)) * dx + + beta_d * phi * div(Dbar*u) * dx + + gamma*(b - b_in + inner(u, grad(bbar))*beta_b) * dx + ) + + if 'coriolis' in equation.prescribed_fields._field_names: + f = equation.prescribed_fields('coriolis') + eqn += beta_u * f * inner(w, equation.domain.perp(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, + constant_jacobian=True) + + # 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) + # # Reconstruction of b + # b = TrialFunction(Vb) + # gamma = TestFunction(Vb) + + # u, D = self.uD.subfunctions + # self.b = Function(Vb) + + # b_eqn = gamma*(b - b_in + inner(u, grad(bbar))*beta_b) * dx + + # b_problem = LinearVariationalProblem(lhs(b_eqn), + # rhs(b_eqn), + # self.b, + # constant_jacobian=True) + # self.b_solver = LinearVariationalSolver(b_problem) + + # Log residuals on hybridized solver + self.log_ksp_residuals(self.uD_solver.snes.ksp) + + @timed_function("Gusto:UpdateReferenceProfiles") + def update_reference_profiles(self): + if self.equations.equivalent_buoyancy: + self.q_sat_func.assign(assemble(self.q_sat_expr_interpolate)) + self.qvbar.assign(assemble(self.q_v_interpolate)) + + @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) + + # Check that the b reference profile has been set + bbar = split(self.equations.X_ref)[2] + b = dy.subfunctions[2] + bbar_func = Function(b.function_space()).interpolate(bbar) + if bbar_func.dat.data.max() == 0 and bbar_func.dat.data.min() == 0: + logger.warning("The reference profile for b in the linear solver is zero. To set a non-zero profile add b to the set_reference_profiles argument.") + + with timed_region("Gusto:VelocityDepthSolve"): + logger.info('Thermal linear solver: mixed solve') + self.uD_solver.solve() + + # u1, D1 = self.uD.subfunctions + # u = dy.subfunctions[0] + # D = dy.subfunctions[1] + # b = dy.subfunctions[2] + # u.assign(u1) + # D.assign(D1) + dy.assign(self.uD) + + # with timed_region("Gusto:BuoyancyRecon"): + # logger.info('Thermal linear solver: buoyancy reconstruction') + # self.b_solver.solve() + + # b.assign(self.b) + +class ThermalSWSolverMono(object): + """ + A general object for solving mixed finite element linear problems. + + This linear solver object is general and is designed for use with different + equation sets, including with the non-linear shallow-water equations. It + forms the linear problem from the equations using FML. The linear system is + solved using a hybridised-mixed method. + """ + + solver_parameters = { + 'ksp_type': 'preonly', + 'pc_type': 'lu' + } + + def __init__(self, equation, alpha, reference_dependent=True): + """ + Args: + equation (:class:`PrognosticEquation`): the model's equation object. + alpha (float): the semi-implicit off-centring factor. A value of 1 + is fully-implicit. + reference_dependent: this indicates that the solver Jacobian should + be rebuilt if the reference profiles have been updated. + """ + self.reference_dependent = reference_dependent + + residual = equation.residual.label_map( + lambda t: t.has_label(linearisation), + lambda t: Term(t.get(linearisation).form, t.labels), + drop) + print(residual.form.__str__()) + breakpoint() + + self.alpha = Constant(alpha) + dt = equation.domain.dt + W = equation.function_space + beta = dt*self.alpha + + # Split up the rhs vector (symbolically) + self.xrhs = Function(W) + + aeqn = residual.label_map( + lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), + map_if_false=lambda t: beta*t) + Leqn = residual.label_map( + lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), + map_if_false=drop) + print(aeqn.form.__str__()) + breakpoint() + print(Leqn.form.__str__()) + breakpoint() + + # Place to put result of solver + self.dy = Function(W) + + # Solver + 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, + constant_jacobian=True) + + self.solver = LinearVariationalSolver(problem, + options_prefix='linear_solver') + + @timed_function("Gusto:UpdateReferenceProfiles") + def update_reference_profiles(self): + if self.reference_dependent: + self.solver.invalidate_jacobian() + + @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) + print(self.xrhs.__str__()) + print(self.dy.__str__()) + self.solver.solve() + dy.assign(self.dy) + + class LinearTimesteppingSolver(object): """ From e5084496e3acad30eb329730c53333ae59279d50 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 8 Aug 2025 15:03:07 +0100 Subject: [PATCH 07/14] PC for calculating a schur complement using slate --- .../skamarock_klemp_nonhydrostatic.py | 69 +++++++++++++++++-- gusto/solvers/linear_solvers.py | 9 ++- gusto/solvers/preconditioners.py | 64 ++++++++++++++++- 3 files changed, 133 insertions(+), 9 deletions(-) diff --git a/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py b/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py index 5ce8e5e49..9c14d5e3c 100644 --- a/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py +++ b/examples/compressible_euler/skamarock_klemp_nonhydrostatic.py @@ -13,12 +13,11 @@ """ from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter -from petsc4py import PETSc -PETSc.Sys.popErrorHandler() import itertools from firedrake import ( as_vector, SpatialCoordinate, PeriodicIntervalMesh, ExtrudedMesh, exp, sin, - Function, pi, COMM_WORLD + Function, pi, COMM_WORLD, + PETSc ) import numpy as np from gusto import ( @@ -28,6 +27,7 @@ compressible_hydrostatic_balance, RungeKuttaFormulation, CompressibleSolver, SubcyclingOptions, hydrostatic_parameters ) +PETSc.Sys.popErrorHandler() skamarock_klemp_nonhydrostatic_defaults = { 'ncolumns': 150, @@ -133,6 +133,65 @@ def skamarock_klemp_nonhydrostatic( DGUpwind(eqns, "theta", ibp=theta_opts.ibp) ] + lu_params = { + 'ksp_type': 'preonly', + 'pc_type': 'lu', + 'pc_factor_mat_solver_type': 'mumps', + } + + gamg_params = { + 'ksp_type': 'fgmres', + 'ksp_rtol': 1.0e-8, + 'ksp_atol': 1.0e-8, + 'ksp_max_it': 100, + 'pc_type': 'gamg', + 'pc_gamg_sym_graph': None, + 'mg_levels': { + 'ksp_type': 'gmres', + 'ksp_max_it': 5, + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu', + }, + 'mg_coarse': lu_params + } + + scpc_parameters = { + 'mat_type': 'matfree', + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'firedrake.SCPC', + 'pc_sc_eliminate_fields': '0, 1', + 'condensed_field': gamg_params + } + + slate_parameters = { + 'mat_type': 'matfree', + 'ksp_type': 'preonly', + 'pc_type': 'fieldsplit', + 'pc_fieldsplit_type': 'additive', + 'pc_fieldsplit_type': 'schur', + 'pc_fieldsplit_schur_fact_type': 'full', + 'pc_fieldsplit_0_fields': '0,1', + 'pc_fieldsplit_1_fields': '2', + 'fieldsplit_0': { + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'firedrake.AssembledPC', + 'assembled_pc_type': 'ilu', + }, + 'fieldsplit_1': { + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'gusto.SlateSchurPC', + 'pc_slateschur_fields': 2, + 'slateschur': gamg_params + }, + } + + solver_parameters = scpc_parameters + # solver_parameters = slate_parameters + + # Linear solver if hydrostatic: linear_solver = CompressibleSolver( @@ -140,7 +199,9 @@ def skamarock_klemp_nonhydrostatic( overwrite_solver_parameters=True ) else: - linear_solver = CompressibleSolver(eqns) + linear_solver = CompressibleSolver( + eqns, solver_parameters=solver_parameters, + overwrite_solver_parameters=True) # Time stepper stepper = SemiImplicitQuasiNewton( diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py index 49ea330d0..e3397e248 100644 --- a/gusto/solvers/linear_solvers.py +++ b/gusto/solvers/linear_solvers.py @@ -335,11 +335,14 @@ def L_tr(f): # Function for the hybridized solutions self.urhol0 = Function(M) + appctx = {'slateschur_form': aeqn} + hybridized_prb = LinearVariationalProblem(aeqn, Leqn, self.urhol0, constant_jacobian=True) hybridized_solver = LinearVariationalSolver(hybridized_prb, solver_parameters=self.solver_parameters, - options_prefix='ImplicitSolver') + options_prefix='ImplicitSolver', + appctx=appctx) self.hybridized_solver = hybridized_solver # Project broken u into the HDiv space using facet averaging. @@ -375,8 +378,8 @@ def L_tr(f): # Log residuals on hybridized solver self.log_ksp_residuals(self.hybridized_solver.snes.ksp) # Log residuals on the trace system too - python_context = self.hybridized_solver.snes.ksp.pc.getPythonContext() - attach_custom_monitor(python_context, logging_ksp_monitor_true_residual) + # python_context = self.hybridized_solver.snes.ksp.pc.getPythonContext() + # attach_custom_monitor(python_context, logging_ksp_monitor_true_residual) @timed_function("Gusto:UpdateReferenceProfiles") def update_reference_profiles(self): diff --git a/gusto/solvers/preconditioners.py b/gusto/solvers/preconditioners.py index 722104830..248b4bb5a 100644 --- a/gusto/solvers/preconditioners.py +++ b/gusto/solvers/preconditioners.py @@ -1,7 +1,8 @@ """A module containing specialised preconditioners for Gusto applications.""" from firedrake import (dot, jump, dS_h, ds_b, ds_t, ds, - FacetNormal, Tensor, AssembledVector) + FacetNormal, Tensor, AssembledVector, + AuxiliaryOperatorPC, AssembledPC, PETSc) from firedrake.preconditioners import PCBase from firedrake.matrix_free.operators import ImplicitMatrixContext @@ -9,9 +10,68 @@ from gusto.recovery.recovery_kernels import AverageKernel, AverageWeightings from pyop2.profiling import timed_region, timed_function from functools import partial +from numpy import arange -__all__ = ["VerticalHybridizationPC"] +__all__ = ["VerticalHybridizationPC", "SlateSchurPC"] + + +class SlateSchurPC(AuxiliaryOperatorPC): + _prefix = "slateschur_" + + def form(self, pc, test, trial): + appctx = self.get_appctx(pc) + + aform = appctx['slateschur_form'] + Va = aform.arguments()[0].function_space() + Vlen = len(Va) + + # which fields are in the schur complement? + pc_prefix = pc.getOptionsPrefix() + "pc_" + self._prefix + fields = PETSc.Options().getIntArray( + pc_prefix + 'fields', [Vlen-1]) + nfields = len(fields) + + if nfields > Vlen - 1: + raise ValueError("Must have at least one uneliminated field") + + first_fields = arange(nfields) + last_fields = arange(Vlen - nfields, Vlen) + + # eliminate fields not in the schur complement + eliminate_first = all(fields == last_fields) + eliminate_last = all(fields == first_fields) + + if not any((eliminate_first, eliminate_last)): + raise ValueError( + "Can only eliminate contiguous fields at the" + f" beginning {first_fields} or end {last_fields}" + f" of function space, not {fields}") + + a = Tensor(aform) + if eliminate_first: + n = Vlen - nfields + a00 = a.blocks[:n, :n] + a10 = a.blocks[n:, :n] + a01 = a.blocks[:n, n:] + a11 = a.blocks[n:, n:] + elif eliminate_last: + n = nfields + a00 = a.blocks[n:, n:] + a10 = a.blocks[:n, n:] + a01 = a.blocks[n:, :n] + a11 = a.blocks[:n, :n] + + schur_complement = a11 - a10*a00.inv*a01 + + return (schur_complement, None) + + def view(self, pc, viewer=None): + super().view(pc, viewer) + if hasattr(self, "pc"): + msg = "PC to approximate Schur complement using Slate.\n" + viewer.printfASCII(msg) + self.pc.view(viewer) class VerticalHybridizationPC(PCBase): From 593f7bae60c32c771c8590039e7488129c1e4201 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 8 Aug 2025 15:42:14 +0100 Subject: [PATCH 08/14] remove MoistConvectiveSWSolver --- .../moist_convective_williamson_2.py | 72 +++-- gusto/solvers/linear_solvers.py | 280 +----------------- gusto/solvers/preconditioners.py | 3 +- 3 files changed, 58 insertions(+), 297 deletions(-) diff --git a/examples/shallow_water/moist_convective_williamson_2.py b/examples/shallow_water/moist_convective_williamson_2.py index f7ce33344..e639a9c77 100644 --- a/examples/shallow_water/moist_convective_williamson_2.py +++ b/examples/shallow_water/moist_convective_williamson_2.py @@ -13,15 +13,14 @@ """ from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter -from firedrake import SpatialCoordinate, sin, cos, exp, Function, errornorm, norm +from firedrake import SpatialCoordinate, sin, cos, exp, Function, errornorm, norm, VectorSpaceBasis from gusto import ( Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations, ZonalComponent, MeridionalComponent, SteadyStateError, lonlatr_from_xyz, - DG1Limiter, InstantRain, MoistConvectiveSWSolver, ForwardEuler, + DG1Limiter, InstantRain, ForwardEuler, LinearTimesteppingSolver, RelativeVorticity, SWSaturationAdjustment, WaterVapour, CloudWater, Rain, - GeneralIcosahedralSphereMesh, xyz_vector_from_lonlatr, - MoistConvectiveSWSolver, MoistConvectiveSWSolverNew + GeneralIcosahedralSphereMesh, xyz_vector_from_lonlatr ) from gusto.core.labels import time_derivative, prognostic, coriolis, pressure_gradient, transport @@ -139,32 +138,45 @@ def sat_func(x_in): SSPRK3(domain, "rain", limiter=limiter) ] - import numpy as np - np.random.seed(6) - - xold = Function(eqns.function_space) - yold = Function(eqns.function_space) - - for xsub in xold.subfunctions[:2]: - with xsub.dat.vec_wo as v: - v.array[:] = np.random.random_sample(v.array.shape) - - xnew = xold.copy(deepcopy=True) - ynew = yold.copy(deepcopy=True) - - old_solver = MoistConvectiveSWSolver(eqns, options_prefix="swe_old") - new_solver = MoistConvectiveSWSolverNew(eqns, options_prefix="swe_new") - - old_solver.solve(xold, yold) - new_solver.solve(xnew, ynew) - - # for i, (xo, xn) in enumerate(zip(xold.subfunctions, xnew.subfunctions)): - # print(f"{i = }: {errornorm(xo, xn)/norm(xo) = :4e}") - for i, (yo, yn) in enumerate(zip(yold.subfunctions, ynew.subfunctions)): - print(f"{i = }| {norm(yo) = :2e} | {norm(yn) = :2e} | {errornorm(yo, yn)/norm(yo) = :4e}") - - from sys import exit; exit() - + solver_parameters = { + 'mat_type': 'matfree', + 'ksp_type': 'preonly', + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "additive", + "pc_fieldsplit_0_fields": "0,1", + "pc_fieldsplit_1_fields": "2,3,4", + "fieldsplit_0": { + 'ksp_monitor_true_residual': None, + 'ksp_type': 'preonly', + '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' + } + } + }, + "fieldsplit_1": { + "ksp_type": "preonly", + "pc_type": "none" + }, + } + + # Provide callback for the nullspace of the trace system + def trace_nullsp(T): + return VectorSpaceBasis(constant=True) + appctx = {"trace_nullspace": trace_nullsp} + linear_solver = LinearTimesteppingSolver( + eqns, alpha=0.5, + reference_dependent=True, + solver_parameters=solver_parameters, + options_prefix="swe_gen", appctx=appctx) # Physics schemes sat_adj = SWSaturationAdjustment( diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py index e3397e248..e3f54ab70 100644 --- a/gusto/solvers/linear_solvers.py +++ b/gusto/solvers/linear_solvers.py @@ -13,11 +13,10 @@ BrokenElement, FunctionSpace, MixedFunctionSpace, DirichletBC, as_vector, assemble, conditional ) -from firedrake.fml import Term, drop, subject +from firedrake.fml import Term, drop from firedrake.petsc import flatten_parameters from firedrake.__future__ import interpolate from pyop2.profiling import timed_function, timed_region -from ufl import derivative as ufl_derivative from gusto.equations.active_tracers import TracerVariableType from gusto.core.logging import ( @@ -30,8 +29,7 @@ from abc import ABCMeta, abstractmethod, abstractproperty -__all__ = ["BoussinesqSolver", "LinearTimesteppingSolver", "CompressibleSolver", - "ThermalSWSolver", "MoistConvectiveSWSolver", "MoistConvectiveSWSolverNew"] +__all__ = ["BoussinesqSolver", "LinearTimesteppingSolver", "CompressibleSolver", "ThermalSWSolver"] class TimesteppingSolver(object, metaclass=ABCMeta): @@ -801,7 +799,8 @@ class LinearTimesteppingSolver(object): 'sub_pc_type': 'ilu'}} } - def __init__(self, equation, alpha, reference_dependent=True): + def __init__(self, equation, alpha, reference_dependent=True, + solver_parameters=None, options_prefix=None, appctx=None): """ Args: equation (:class:`PrognosticEquation`): the model's equation object. @@ -809,9 +808,17 @@ def __init__(self, equation, alpha, reference_dependent=True): is fully-implicit. reference_dependent: this indicates that the solver Jacobian should be rebuilt if the reference profiles have been updated. + solver_parameters (dict, optional): contains the options to be + passed to the underlying :class:`LinearVariationalSolver`. + Defaults to None. + options_prefix: options_prefix for the underlying :class:`LinearVariationalSolver`. + appctx: appctx for the underlying :class:`LinearVariationalSolver`. """ self.reference_dependent = reference_dependent + params = solver_parameters or self.solver_parameters + prefix = options_prefix or 'linear_solver' + residual = equation.residual.label_map( lambda t: t.has_label(linearisation), lambda t: Term(t.get(linearisation).form, t.labels), @@ -842,9 +849,9 @@ def __init__(self, equation, alpha, reference_dependent=True): self.dy, bcs=bcs, constant_jacobian=True) - self.solver = LinearVariationalSolver(problem, - solver_parameters=self.solver_parameters, - options_prefix='linear_solver') + self.solver = LinearVariationalSolver(problem, appctx=appctx, + solver_parameters=params, + options_prefix=prefix) @timed_function("Gusto:UpdateReferenceProfiles") def update_reference_profiles(self): @@ -865,260 +872,3 @@ 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_u = dt*self.tau_values.get("u", self.alpha) - beta_d = dt*self.tau_values.get("D", self.alpha) - Vu = equation.domain.spaces("HDiv") - VD = equation.domain.spaces("DG") - - # 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_u * (D - Dbar) * div(w*g) * dx - + inner(phi, (D - D_in)) * dx - + beta_d * phi * div(Dbar*u) * dx - ) - - if 'coriolis' in equation.prescribed_fields._field_names: - f = equation.prescribed_fields('coriolis') - eqn += beta_u * f * inner(w, equation.domain.perp(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, - constant_jacobian=True) - - # 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, options_prefix=self.options_prefix) - - # Log residuals on hybridized solver - self.log_ksp_residuals(self.uD_solver.snes.ksp) - - @timed_function("Gusto:UpdateReferenceProfiles") - def update_reference_profiles(self): - self.uD_solver.invalidate_jacobian() - - @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) - - -class MoistConvectiveSWSolverNew(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 = { - 'mat_type': 'matfree', - 'ksp_type': 'preonly', - "pc_type": "fieldsplit", - "pc_fieldsplit_type": "additive", - "pc_fieldsplit_0_fields": "0,1", - "pc_fieldsplit_1_fields": "2,3,4", - "fieldsplit_0": { - 'ksp_monitor_true_residual': None, - 'ksp_type': 'preonly', - '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' - } - } - }, - "fieldsplit_1": { - "ksp_type": "preonly", - "pc_type": "none" - }, - } - - @timed_function("Gusto:SolverSetup") - def _setup_solver(self): - equation = self.equations # just cutting down line length a bit - dt = self.dt - - # Set up Preconditioner aP - beta_u = dt*self.tau_values.get("u", self.alpha) - beta_d = dt*self.tau_values.get("D", self.alpha) - - # Split up the rhs vectorx - M = self.equations.function_space - self.xrhs = Function(M) - u_in = split(self.xrhs)[0] - D_in = split(self.xrhs)[1] - - - # Build the reduced function space for u, D - w, phi, q0, q1, q2 = TestFunctions(M) - u, D, p0, p1, p2 = TrialFunctions(M) - - # Get background depth - Dbar = split(equation.X_ref)[1] - - g = equation.parameters.g - - eqn = ( - inner(w, u) * dx - - beta_u * (D - Dbar) * div(w*g) * dx - + inner(phi, D) * dx - + beta_d * phi * div(Dbar*u) * dx - + inner(q0, (p0)) * dx - + inner(q1, (p1)) * dx - + inner(q2, (p2)) * dx - ) - - if 'coriolis' in equation.prescribed_fields._field_names: - f = equation.prescribed_fields('coriolis') - eqn += beta_u * f * inner(w, equation.domain.perp(u)) * dx - - aP = lhs(eqn) - - # Calculated linearised problem - residual = equation.residual.label_map( - lambda t: t.has_label(linearisation), - lambda t: Term(t.get(linearisation).form, t.labels), - drop) - - beta = dt*self.alpha - - aeqn = residual.label_map( - lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), - map_if_false=lambda t: beta*t) - Leqn = residual.label_map( - lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), - map_if_false=drop) - - # 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.form, action(Leqn.form, self.xrhs), - self.uD, bcs=bcs, constant_jacobian=True) - - # 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, options_prefix=self.options_prefix) - - # Log residuals on hybridized solver - self.log_ksp_residuals(self.uD_solver.snes.ksp) - - @timed_function("Gusto:UpdateReferenceProfiles") - def update_reference_profiles(self): - self.uD_solver.invalidate_jacobian() - - @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) - for xsub in self.xrhs.subfunctions[2:]: - xsub.zero() - for ysub in self.uD.subfunctions[2:]: - ysub.zero() - - with timed_region("Gusto:VelocityDepthSolve"): - logger.info('Moist convective linear solver: mixed solve') - self.uD_solver.solve() - - dy.assign(self.uD) diff --git a/gusto/solvers/preconditioners.py b/gusto/solvers/preconditioners.py index 248b4bb5a..62b4cf0c6 100644 --- a/gusto/solvers/preconditioners.py +++ b/gusto/solvers/preconditioners.py @@ -2,11 +2,10 @@ from firedrake import (dot, jump, dS_h, ds_b, ds_t, ds, FacetNormal, Tensor, AssembledVector, - AuxiliaryOperatorPC, AssembledPC, PETSc) + AuxiliaryOperatorPC, PETSc) from firedrake.preconditioners import PCBase from firedrake.matrix_free.operators import ImplicitMatrixContext -from firedrake.petsc import PETSc from gusto.recovery.recovery_kernels import AverageKernel, AverageWeightings from pyop2.profiling import timed_region, timed_function from functools import partial From 32ce053094c89a3496d6b05b04aacd1c2839a8ab Mon Sep 17 00:00:00 2001 From: atb1995 Date: Fri, 8 Aug 2025 16:26:04 +0100 Subject: [PATCH 09/14] Revert "Merge remote-tracking branch 'origin/TBendall/SolverEqns' into auxiliary-equation-pc" This reverts commit 5e51f2c9117134df89b8350cc69bfa78da8ae1ef, reversing changes made to 807c7c01eeec1da5d60284c97d214de7e2e79a6b. --- gusto/equations/boussinesq_equations.py | 30 +- .../equations/compressible_euler_equations.py | 21 +- gusto/equations/prognostic_equations.py | 11 +- gusto/equations/shallow_water_equations.py | 73 ++-- .../semi_implicit_quasi_newton.py | 355 ++++-------------- 5 files changed, 134 insertions(+), 356 deletions(-) diff --git a/gusto/equations/boussinesq_equations.py b/gusto/equations/boussinesq_equations.py index dccd35896..d94e2a038 100644 --- a/gusto/equations/boussinesq_equations.py +++ b/gusto/equations/boussinesq_equations.py @@ -3,7 +3,7 @@ from firedrake import inner, dx, div, cross, split, as_vector from firedrake.fml import subject from gusto.core.labels import ( - time_derivative, prognostic, linearisation, + time_derivative, transport, prognostic, linearisation, pressure_gradient, coriolis, divergence, gravity, incompressible ) from gusto.equations.common_forms import ( @@ -60,8 +60,8 @@ def __init__(self, domain, parameters, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation drops terms for any active - tracers. + in which case the linearisation includes time derivatives and + scalar transport terms. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form' and @@ -90,11 +90,14 @@ def __init__(self, domain, parameters, active_tracers = [] if linearisation_map == 'default': - # Default linearisation is to include all terms - # Don't include terms for active tracers - linearisation_map = lambda t: ( - t.has_label(time_derivative) or t.get(prognostic) in field_names - ) + # Default linearisation is time derivatives, scalar transport, + # pressure gradient, gravity and divergence terms + # Don't include active tracers + linearisation_map = lambda t: \ + t.get(prognostic) in ['u', 'p', 'b'] \ + and (any(t.has_label(time_derivative, pressure_gradient, + divergence, gravity)) + or (t.get(prognostic) not in ['u', 'p'] and t.has_label(transport))) super().__init__(field_names, domain, space_names, linearisation_map=linearisation_map, @@ -246,8 +249,8 @@ def __init__(self, domain, parameters, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation drops terms for any active - tracers. + in which case the linearisation includes time derivatives and + scalar transport terms. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form' and @@ -264,6 +267,13 @@ def __init__(self, domain, parameters, NotImplementedError: active tracers are not implemented. """ + if linearisation_map == 'default': + # Default linearisation is time derivatives, pressure gradient, + # Coriolis and transport term from depth equation + linearisation_map = lambda t: \ + (any(t.has_label(time_derivative, pressure_gradient, coriolis, + gravity, divergence, incompressible)) + or (t.get(prognostic) in ['p', 'b'] and t.has_label(transport))) super().__init__(domain=domain, parameters=parameters, compressible=compressible, diff --git a/gusto/equations/compressible_euler_equations.py b/gusto/equations/compressible_euler_equations.py index 02a60d7ea..d512432cf 100644 --- a/gusto/equations/compressible_euler_equations.py +++ b/gusto/equations/compressible_euler_equations.py @@ -6,7 +6,7 @@ ) from firedrake.fml import subject, replace_subject from gusto.core.labels import ( - time_derivative, prognostic, hydrostatic, linearisation, + time_derivative, transport, prognostic, hydrostatic, linearisation, pressure_gradient, coriolis, gravity, sponge ) from gusto.equations.thermodynamics import exner_pressure @@ -58,8 +58,8 @@ def __init__(self, domain, parameters, sponge_options=None, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation drops terms for any active - tracers. + in which case the linearisation includes time derivatives and + scalar transport terms. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form' and @@ -88,11 +88,12 @@ def __init__(self, domain, parameters, sponge_options=None, active_tracers = [] if linearisation_map == 'default': - # Default linearisation is to include all terms - # Don't include terms for active tracers - linearisation_map = lambda t: ( - t.has_label(time_derivative) or t.get(prognostic) in field_names - ) + # Default linearisation is time derivatives and scalar transport terms + # Don't include active tracers + linearisation_map = lambda t: \ + t.get(prognostic) in ['u', 'rho', 'theta'] \ + and (t.has_label(time_derivative) + or (t.get(prognostic) != 'u' and t.has_label(transport))) super().__init__(field_names, domain, space_names, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, @@ -305,8 +306,8 @@ def __init__(self, domain, parameters, sponge_options=None, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation drops terms for any active - tracers. + in which case the linearisation includes time derivatives and + scalar transport terms. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form' and diff --git a/gusto/equations/prognostic_equations.py b/gusto/equations/prognostic_equations.py index 1f6035b71..a562795d6 100644 --- a/gusto/equations/prognostic_equations.py +++ b/gusto/equations/prognostic_equations.py @@ -211,15 +211,8 @@ def should_linearise(term): # Linearise a term, and add the linearisation as a label def linearise(term, X, X_ref, du): - try: - derivative = ufl.derivative(term.form, X) - linear_term = Term(action(derivative, du), term.labels) - return linearisation(term, replace_subject(X_ref)(linear_term)) - except IndexError: - # If the derivative is zero (e.g. for the gravity term) then - # action(derivative, du) will raise an IndexError - # In that case, just return the term with no linearisation label - return term + linear_term = Term(action(ufl.derivative(term.form, X), du), term.labels) + return linearisation(term, replace_subject(X_ref)(linear_term)) # Add linearisations to all terms that need linearising residual = residual.label_map( diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py index 61236a138..51fe39d34 100644 --- a/gusto/equations/shallow_water_equations.py +++ b/gusto/equations/shallow_water_equations.py @@ -3,7 +3,7 @@ from firedrake import (inner, dx, div, FunctionSpace, FacetNormal, jump, avg, dS, split, conditional, exp) from firedrake.fml import subject, drop -from gusto.core.labels import (time_derivative, prognostic, +from gusto.core.labels import (time_derivative, transport, prognostic, linearisation, pressure_gradient, coriolis) from gusto.equations.common_forms import ( advection_form, advection_form_1d, continuity_form, @@ -51,8 +51,8 @@ def __init__(self, domain, parameters, fexpr=None, topog_expr=None, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation drops terms for any active - tracers. + in which case the linearisation includes both time derivatives, + the 'D' transport term and the pressure gradient term. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form', and @@ -70,11 +70,12 @@ def __init__(self, domain, parameters, fexpr=None, topog_expr=None, active_tracers = [] if linearisation_map == 'default': - # Default linearisation is to include all terms - # Don't include terms for active tracers - linearisation_map = lambda t: ( - t.has_label(time_derivative) or t.get(prognostic) in field_names - ) + # Default linearisation is time derivatives, pressure gradient and + # transport term from depth equation. Don't include active tracers + linearisation_map = lambda t: \ + t.get(prognostic) in ['u', 'D'] \ + and (any(t.has_label(time_derivative, coriolis, pressure_gradient)) + or (t.get(prognostic) in ['D'] and t.has_label(transport))) field_names = ['u', 'D'] space_names = {'u': 'HDiv', 'D': 'L2'} @@ -231,8 +232,8 @@ def __init__(self, domain, parameters, fexpr=None, topog_expr=None, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation drops terms for any active - tracers. + in which case the linearisation includes both time derivatives, + the 'D' transport term, pressure gradient and Coriolis terms. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form' and @@ -302,8 +303,8 @@ def __init__(self, domain, parameters, equivalent_buoyancy=False, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation drops terms for any active - tracers. + in which case the linearisation includes both time derivatives, + the 'D' transport term and the pressure gradient term. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form', and @@ -336,21 +337,19 @@ def __init__(self, domain, parameters, equivalent_buoyancy=False, breakpoint() if linearisation_map == 'default': + # Default linearisation is time derivatives, pressure + # gradient and transport terms from depth and buoyancy + # equations. Include q_t if equivalent buoyancy. Don't include + # active tracers. + linear_transported = ['D', self.b_name] if equivalent_buoyancy: - # Default linearisation is to include all terms - # Don't include qt terms - linearisation_map = lambda t: ( - t.has_label(time_derivative) - or t.get(prognostic) in ['u', 'D', 'b_e'] - ) - - else: - # Default linearisation is to include all terms - # Don't include terms for active tracers - linearisation_map = lambda t: ( - t.has_label(time_derivative) - or t.get(prognostic) in field_names - ) + linear_transported.append('q_t') + linearisation_map = lambda t: \ + t.get(prognostic) in field_names \ + and (any(t.has_label(time_derivative, pressure_gradient, + coriolis)) + or (t.get(prognostic) in linear_transported + and t.has_label(transport))) # Bypass ShallowWaterEquations.__init__ to avoid having to # define the field_names separately @@ -565,8 +564,8 @@ def __init__(self, domain, parameters, equivalent_buoyancy=False, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation drops terms for any active - tracers. + in which case the linearisation includes both time derivatives, + the 'D' transport term, pressure gradient and Coriolis terms. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: 'vector_invariant_form', 'vector_advection_form' and @@ -617,8 +616,8 @@ class ShallowWaterEquations_1d(PrognosticEquationSet): linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation drops terms for any active - tracers. + in which case the linearisation includes both time derivatives, + the 'D' transport term, pressure gradient and Coriolis terms. no_normal_flow_bc_ids (list, optional): a list of IDs of domain boundaries at which no normal flow will be enforced. Defaults to None. @@ -640,11 +639,11 @@ def __init__(self, domain, parameters, raise NotImplementedError('Tracers not implemented for 1D shallow water equations') if linearisation_map == 'default': - # Default linearisation is to include all terms - # Don't include terms for active tracers - linearisation_map = lambda t: ( - t.has_label(time_derivative) or t.get(prognostic) in field_names - ) + # Default linearisation is time derivatives, pressure gradient, + # Coriolis and transport term from depth equation + linearisation_map = lambda t: \ + (any(t.has_label(time_derivative, pressure_gradient, coriolis)) + or (t.get(prognostic) == 'D' and t.has_label(transport))) super().__init__(field_names, domain, space_names, linearisation_map=linearisation_map, @@ -746,8 +745,8 @@ def __init__(self, domain, parameters, fexpr=None, linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', - in which case the linearisation drops terms for any active - tracers. + in which case the linearisation includes both time derivatives, + the 'D' transport term, pressure gradient and Coriolis terms. no_normal_flow_bc_ids (list, optional): a list of IDs of domain boundaries at which no normal flow will be enforced. Defaults to None. diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py index 9d0b057f4..d8ca50194 100644 --- a/gusto/timestepping/semi_implicit_quasi_newton.py +++ b/gusto/timestepping/semi_implicit_quasi_newton.py @@ -5,18 +5,16 @@ from firedrake import ( Function, Constant, TrialFunctions, DirichletBC, div, assemble, - LinearVariationalProblem, LinearVariationalSolver, FunctionSpace, action + LinearVariationalProblem, LinearVariationalSolver, FunctionSpace ) -from firedrake.fml import drop, replace_subject, Term +from firedrake.fml import drop, replace_subject from firedrake.__future__ import interpolate -from firedrake.petsc import flatten_parameters -from pyop2.profiling import timed_stage, timed_function +from pyop2.profiling import timed_stage from gusto.core import TimeLevelFields, StateFields -from gusto.core.labels import ( - transport, diffusion, time_derivative, hydrostatic, physics_label, sponge, - incompressible, linearisation, prognostic -) -from gusto.solvers import mass_parameters +from gusto.core.labels import (transport, diffusion, time_derivative, + hydrostatic, physics_label, sponge, + incompressible) +from gusto.solvers import LinearTimesteppingSolver, mass_parameters from gusto.core.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.time_discretisation.time_discretisation import ExplicitTimeDiscretisation from gusto.timestepping.timestepper import BaseTimestepper @@ -39,12 +37,10 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, auxiliary_equations_and_schemes=None, linear_solver=None, diffusion_schemes=None, physics_schemes=None, slow_physics_schemes=None, fast_physics_schemes=None, - alpha=0.5, tau_values=None, off_centred_u=False, + alpha=0.5, off_centred_u=False, num_outer=2, num_inner=2, accelerator=True, predictor=None, reference_update_freq=None, - spinup_steps=0, solver_prognostics=None, - linear_solver_parameters=None, - overwrite_linear_solver_parameters=False): + spinup_steps=0): """ Args: equation_set (:class:`PrognosticEquationSet`): the prognostic @@ -80,8 +76,6 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, alpha (`float, optional): the semi-implicit off-centering parameter. A value of 1 corresponds to fully implicit, while 0 corresponds to fully explicit. Defaults to 0.5. - tau_values (dict, optional): contains the semi-implicit relaxation - parameters. Defaults to None, in which case the value of alpha is used. off_centred_u (bool, optional): option to offcentre the transporting velocity. Defaults to False, in which case transporting velocity is centred. If True offcentring uses value of alpha. @@ -116,18 +110,8 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, spinup_steps (int, optional): the number of steps to run the model in "spin-up" mode, where the alpha parameter is set to 1.0. Defaults to 0, which corresponds to no spin-up. - solver_prognostics (list, optional): a list of the names of - prognostic variables to include in the solver. - linear_solver_parameters (dict, optional): contains the options to - be passed to the underlying :class:`LinearVariationalSolver`. - for the Defaults to None. - overwrite_solver_parameters (bool, optional): if True use only the - `solver_parameters` that have been passed in. If False then - update the default parameters with the `solver_parameters` - passed in. Defaults to False. """ - # Basic parts of the SIQN structure ------------------------------------ self.num_outer = num_outer self.num_inner = num_inner mesh = equation_set.domain.mesh @@ -135,58 +119,8 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, self.alpha = Function(R, val=float(alpha)) self.predictor = predictor self.accelerator = accelerator - self.implicit_terms = [incompressible, sponge] - - # default is to not offcentre transporting velocity but if it - # is offcentred then use the same value as alpha - if off_centred_u: - self.alpha_u = Function(R, val=float(alpha)) - else: - self.alpha_u = Function(R, val=0.5) - - # Set up some important fields ----------------------------------------- - self.field_name = equation_set.field_name - W = equation_set.function_space - self.xrhs = Function(W) - self.xrhs_phys = Function(W) - self.dy = Function(W) - - # Determine prognostics for solver ------------------------------------- - self.non_solver_prognostics = [] - if solver_prognostics is not None: - assert type(solver_prognostics) is list, ( - "solver_prognostics should be a list of prognostic variable " - + f"names, and not {type(solver_prognostics)}" - ) - for prognostic_name in solver_prognostics: - assert prognostic_name in equation_set.field_names, ( - f"Prognostic variable {prognostic_name} not found in " - + "equation set field names" - ) - for field_name in equation_set.field_names: - if field_name not in solver_prognostics: - self.non_solver_prognostics.append(field_name) - else: - solver_prognostics = equation_set.field_names - logger.info( - 'Semi-implicit Quasi-Newton: Using solver prognostics ' - + f'{solver_prognostics}' - ) - - # BCs, Forcing and Linear Solver --------------------------------------- - self.bcs = equation_set.bcs - self.forcing = Forcing(equation_set, self.implicit_terms, self.alpha) - if linear_solver is None: - self.linear_solver = SIQNLinearSolver( - equation_set, solver_prognostics, self.implicit_terms, - self.alpha, tau_values=tau_values, - solver_parameters=linear_solver_parameters, - overwrite_solver_parameters=overwrite_linear_solver_parameters - ) - else: - self.linear_solver = linear_solver - # Options relating to reference profiles and spin-up ------------------- + # Options relating to reference profiles and spin-up self._alpha_original = float(alpha) # float so as to not upset adjoint self.reference_update_freq = reference_update_freq self.to_update_ref_profile = False @@ -194,18 +128,43 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, self.spinup_begun = False self.spinup_done = False - # Transport ------------------------------------------------------------ + # Flag for if we have simultaneous transport + self.simult = False + + # default is to not offcentre transporting velocity but if it + # is offcentred then use the same value as alpha + self.alpha_u = Function(R, val=float(alpha)) if off_centred_u else Function(R, val=0.5) + self.spatial_methods = spatial_methods - # Flag for if we have simultaneous transport of tracers - self.simult = False + if physics_schemes is not None: + self.final_physics_schemes = physics_schemes + else: + self.final_physics_schemes = [] + if slow_physics_schemes is not None: + self.slow_physics_schemes = slow_physics_schemes + else: + self.slow_physics_schemes = [] + if fast_physics_schemes is not None: + self.fast_physics_schemes = fast_physics_schemes + else: + self.fast_physics_schemes = [] + self.all_physics_schemes = (self.slow_physics_schemes + + self.fast_physics_schemes + + self.final_physics_schemes) + + for parametrisation, scheme in self.all_physics_schemes: + assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" + if hasattr(parametrisation, "explicit_only") and parametrisation.explicit_only: + assert isinstance(scheme, ExplicitTimeDiscretisation), \ + ("Only explicit time discretisations can be used with " + + f"physics scheme {parametrisation.label.label}") self.active_transport = [] for scheme in transport_schemes: - assert scheme.nlevels == 1, \ - "multilevel schemes not supported as part of this timestepping loop" + assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" if isinstance(scheme.field_name, list): - # This means that multiple fields are transported simultaneously + # This means that multiple fields are being transported simultaneously self.simult = True for subfield in scheme.field_name: assert subfield in equation_set.field_names @@ -216,8 +175,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, for method in spatial_methods: if subfield == method.variable and method.term_label == transport: method_found = True - assert method_found, \ - f'No transport method found for variable {scheme.field_name}' + assert method_found, f'No transport method found for variable {scheme.field_name}' self.active_transport.append((scheme.field_name, scheme)) else: assert scheme.field_name in equation_set.field_names @@ -228,40 +186,12 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, if scheme.field_name == method.variable and method.term_label == transport: method_found = True self.active_transport.append((scheme.field_name, scheme)) - assert method_found, \ - f'No transport method found for variable {scheme.field_name}' - - # Physics schemes ------------------------------------------------------ - if physics_schemes is not None: - self.final_physics_schemes = physics_schemes - else: - self.final_physics_schemes = [] - if slow_physics_schemes is not None: - self.slow_physics_schemes = slow_physics_schemes - else: - self.slow_physics_schemes = [] - if fast_physics_schemes is not None: - self.fast_physics_schemes = fast_physics_schemes - else: - self.fast_physics_schemes = [] - self.all_physics_schemes = (self.slow_physics_schemes - + self.fast_physics_schemes - + self.final_physics_schemes) + assert method_found, f'No transport method found for variable {scheme.field_name}' - for parametrisation, scheme in self.all_physics_schemes: - assert scheme.nlevels == 1, \ - "multilevel schemes not supported as part of this timestepping loop" - if hasattr(parametrisation, "explicit_only") and parametrisation.explicit_only: - assert isinstance(scheme, ExplicitTimeDiscretisation), \ - ("Only explicit time discretisations can be used with " - + f"physics scheme {parametrisation.label.label}") - - # Diffusion and other terms -------------------------------------------- self.diffusion_schemes = [] if diffusion_schemes is not None: for scheme in diffusion_schemes: - assert scheme.nlevels == 1, \ - "multilevel schemes not supported as part of this timestepping loop" + assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" assert scheme.field_name in equation_set.field_names self.diffusion_schemes.append((scheme.field_name, scheme)) # Check that there is a corresponding transport method @@ -273,8 +203,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, if auxiliary_equations_and_schemes is not None: for eqn, scheme in auxiliary_equations_and_schemes: - assert not hasattr(eqn, "field_names"), \ - 'Cannot use auxiliary schemes with multiple fields' + assert not hasattr(eqn, "field_names"), 'Cannot use auxiliary schemes with multiple fields' self.auxiliary_schemes = [ (eqn.field_name, scheme) for eqn, scheme in auxiliary_equations_and_schemes] @@ -284,20 +213,30 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, self.auxiliary_schemes = [] self.auxiliary_equations_and_schemes = auxiliary_equations_and_schemes - # General set up ------------------------------------------------------- - # This sets up all of the time discretisations for transport, diffusion, - # physics schemes, etc, so must happen after those have been worked out super().__init__(equation_set, io) - # Set up auxiliary equations separately, as these are not done as part - # of the super init for aux_eqn, aux_scheme in self.auxiliary_equations_and_schemes: self.setup_equation(aux_eqn) aux_scheme.setup(aux_eqn) self.setup_transporting_velocity(aux_scheme) - # Set up the predictor last, as it requires the state fields to already - # be created + self.tracers_to_copy = [] + if equation_set.active_tracers is not None: + for active_tracer in equation_set.active_tracers: + self.tracers_to_copy.append(active_tracer.name) + + self.field_name = equation_set.field_name + W = equation_set.function_space + self.xrhs = Function(W) + self.xrhs_phys = Function(W) + self.dy = Function(W) + if linear_solver is None: + self.linear_solver = LinearTimesteppingSolver(equation_set, self.alpha) + else: + self.linear_solver = linear_solver + self.forcing = Forcing(equation_set, self.alpha) + self.bcs = equation_set.bcs + if self.predictor is not None: V_DG = equation_set.domain.spaces('DG') self.predictor_field_in = Function(V_DG) @@ -368,7 +307,7 @@ def copy_active_tracers(self, x_in, x_out): x_out: The output set of fields """ - for name in self.non_solver_prognostics: + for name in self.tracers_to_copy: x_out(name).assign(x_in(name)) def transport_fields(self, outer, xstar, xp): @@ -596,19 +535,18 @@ class Forcing(object): semi-implicit time discretisation. """ - def __init__(self, equation, implicit_terms, alpha): + def __init__(self, equation, alpha): """ Args: equation (:class:`PrognosticEquationSet`): the prognostic equations containing the forcing terms. - implicit_terms (list): list of labels for terms that are always - fully-implicit. alpha (:class:`Function`): semi-implicit off-centering factor. An alpha of 0 corresponds to fully explicit, while a factor of 1 corresponds to fully implicit. """ self.field_name = equation.field_name + implicit_terms = [incompressible, sponge] dt = equation.domain.dt W = equation.function_space @@ -740,166 +678,3 @@ def zero_forcing_terms(self, equation, x_in, x_out, field_names): logger.debug(f'Semi-Implicit Quasi Newton: Zeroing implicit forcing for {field_name}') field_index = equation.field_names.index(field_name) x_out.subfunctions[field_index].assign(x_in(field_name)) - - -class SIQNLinearSolver(object): - """ - Sets up the linear solver for the Semi-Implicit Quasi-Newton timestepper. - - This computes the approximate linearisation of the algorithm for computing - the (n+1)-th level state, and sets up matrix-vector problem for the - Quasi-Newton problem. - """ - - 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'}} - } - - solver_parameters = { - 'ksp_error_if_not_converged': None, - 'pc_factor_shift_type': 'nonzero', - 'snes_type': 'ksponly', - } - - def __init__(self, equation, solver_prognostics, implicit_terms, - alpha, tau_values=None, solver_parameters=None, - overwrite_solver_parameters=False, reference_dependent=True): - """ - Args: - equations (:class:`PrognosticEquation`): the model's equation. - solver_prognostics (list): a list of prognostic variable names to - include in the linear solver. - implicit_terms (list): a list of labels for terms that are always - fully-implicit. - alpha (float, optional): the semi-implicit off-centring factor. - Defaults to 0.5. A value of 1 is fully-implicit. - tau_values (dict, optional): contains the semi-implicit relaxation - parameters. Defaults to None, in which case the value of alpha - is used. - solver_parameters (dict, optional): contains the options to be - passed to the underlying :class:`LinearVariationalSolver`. - Defaults to None. - overwrite_solver_parameters (bool, optional): if True use only the - `solver_parameters` that have been passed in. If False then - update the default parameters with the `solver_parameters` - passed in. Defaults to False. - reference_dependent: this indicates that the solver Jacobian should - be rebuilt if the reference profiles have been updated. - """ - - # Update or set solver parameters -------------------------------------- - if solver_parameters is not None: - if not overwrite_solver_parameters: - p = flatten_parameters(self.solver_parameters) - p.update(flatten_parameters(solver_parameters)) - solver_parameters = p - self.solver_parameters = solver_parameters - - if logger.isEnabledFor(DEBUG): - self.solver_parameters["ksp_monitor_true_residual"] = None - - dt = equation.domain.dt - self.reference_dependent = reference_dependent - self.alpha = Constant(alpha) - self.tau_values = tau_values if tau_values is not None else {} - - for prognostic_name in solver_prognostics: - if prognostic_name not in self.tau_values.keys(): - # If no tau value is specified for this prognostic, use alpha - self.tau_values[prognostic_name] = self.alpha - - # Set up the linearisation of the equation ----------------------------- - # Time derivative terms: all prognostic variables - residual = equation.residual.label_map( - lambda t: t.has_label(time_derivative), - lambda t: Term(t.get(linearisation).form, t.labels), - drop - ) - - # All other linear terms (only specified prognostics) - for prognostic_name in solver_prognostics: - # Add semi-implicit terms for this prognostic - tau = self.tau_values[prognostic_name] - residual += equation.residual.label_map( - lambda t: ( - t.has_label(linearisation) - and not t.has_label(time_derivative) - and t.get(prognostic) == prognostic_name - and not any(t.has_label(*implicit_terms, return_tuple=True)) - ), - lambda t: tau*dt*Term(t.get(linearisation).form, t.labels), - drop - ) - # Add implicit terms for this prognostic - residual += equation.residual.label_map( - lambda t: ( - t.has_label(linearisation) - and not t.has_label(time_derivative) - and t.get(prognostic) == prognostic_name - and any(t.has_label(*implicit_terms, return_tuple=True)) - ), - lambda t: dt*Term(t.get(linearisation).form, t.labels), - drop - ) - - W = equation.function_space - - # Split up the rhs vector (symbolically) - self.xrhs = Function(W) - - aeqn = residual - Leqn = residual.label_map( - lambda t: t.has_label(time_derivative), map_if_false=drop - ) - - # Place to put result of solver - self.dy = Function(W) - - # Set up the solver ---------------------------------------------------- - 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, - constant_jacobian=True) - - self.solver = LinearVariationalSolver( - problem, solver_parameters=self.solver_parameters, - options_prefix='linear_solver' - ) - - @staticmethod - def log_ksp_residuals(ksp): - if logger.isEnabledFor(DEBUG): - ksp.setMonitor(logging_ksp_monitor_true_residual) - - @timed_function("Gusto:UpdateReferenceProfiles") - def update_reference_profiles(self): - if self.reference_dependent: - self.solver.invalidate_jacobian() - - @timed_function("Gusto:LinearSolve") - def solve(self, xrhs, dy): - """ - Solve the linear problem. - - Args: - xrhs (:class:`Function`): the right-hand side residual field in the - appropriate :class:`MixedFunctionSpace`. - dy (:class:`Function`): the resulting increment field in the - appropriate :class:`MixedFunctionSpace`. - """ - self.xrhs.assign(xrhs) - self.solver.solve() - dy.assign(self.dy) From 6ae2e4b115e9c955fea275c4b9681a69fa53124f Mon Sep 17 00:00:00 2001 From: atb1995 Date: Fri, 8 Aug 2025 16:38:00 +0100 Subject: [PATCH 10/14] Monolithic solver for hard coded linear thermal eqns added --- .../shallow_water/moist_thermal_equivb_gw.py | 17 +- gusto/equations/shallow_water_equations.py | 2 - gusto/solvers/linear_solvers.py | 159 +++++++++++++++++- 3 files changed, 157 insertions(+), 21 deletions(-) diff --git a/examples/shallow_water/moist_thermal_equivb_gw.py b/examples/shallow_water/moist_thermal_equivb_gw.py index d876ff41d..ebe900a5d 100644 --- a/examples/shallow_water/moist_thermal_equivb_gw.py +++ b/examples/shallow_water/moist_thermal_equivb_gw.py @@ -66,17 +66,10 @@ def moist_thermal_gw( nu=nu, beta2=beta2) Omega = parameters.Omega fexpr = 2*Omega*xyz[2]/radius - # def linearisation_map(term): - # if term.has_label(time_derivative): - # return True - # else: - # return False - linearisation_map = lambda t: t.has_label(time_derivative) - # Equation eqns = ThermalShallowWaterEquations( domain, parameters, fexpr=fexpr, - equivalent_buoyancy=False, linearisation_map = linearisation_map + equivalent_buoyancy=False ) # IO @@ -111,12 +104,12 @@ def moist_thermal_gw( xmono = xold.copy(deepcopy=True) ymono = yold.copy(deepcopy=True) - # old_solver = ThermalSWSolver(eqns, options_prefix="swe_old") - # new_solver = ThermalSWSolverNew(eqns, options_prefix="swe_new") + old_solver = ThermalSWSolver(eqns, options_prefix="swe_old") + new_solver = ThermalSWSolverNew(eqns, options_prefix="swe_new") mono_solver = ThermalSWSolverMono(eqns, alpha=0.5) - # old_solver.solve(xold, yold) - # new_solver.solve(xnew, ynew) + old_solver.solve(xold, yold) + new_solver.solve(xnew, ynew) mono_solver.solve(xmono, ymono) # for i, (xo, xn) in enumerate(zip(xold.subfunctions, xnew.subfunctions)): diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py index 51fe39d34..0e9aebe2a 100644 --- a/gusto/equations/shallow_water_equations.py +++ b/gusto/equations/shallow_water_equations.py @@ -334,8 +334,6 @@ def __init__(self, domain, parameters, equivalent_buoyancy=False, if active_tracers is None: active_tracers = [] - breakpoint() - if linearisation_map == 'default': # Default linearisation is time derivatives, pressure # gradient and transport terms from depth and buoyancy diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py index 18676ebbc..a70f6a668 100644 --- a/gusto/solvers/linear_solvers.py +++ b/gusto/solvers/linear_solvers.py @@ -987,7 +987,158 @@ def solve(self, xrhs, dy): # b.assign(self.b) -class ThermalSWSolverMono(object): +class ThermalSWSolverMono(TimesteppingSolver): + """Linear solver object for the thermal shallow water equations. + + This solves a linear problem for the thermal shallow water + equations with prognostic variables u (velocity), D (depth) and + either b (buoyancy) or b_e (equivalent buoyancy). It follows the + following strategy: + + (1) Eliminate b + (2) Solve the resulting system for (u, D) using a hybrid-mixed method + (3) Reconstruct b + + """ + + # 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'}} + # } + + + solver_parameters = { + 'ksp_error_if_not_converged': None, + 'pc_factor_shift_type': 'nonzero', + 'snes_type': 'ksponly', + } + + @timed_function("Gusto:SolverSetup") + def _setup_solver(self): + equation = self.equations # just cutting down line length a bit + equivalent_buoyancy = equation.equivalent_buoyancy + + dt = self.dt + beta_u = dt*self.tau_values.get("u", self.alpha) + beta_d = dt*self.tau_values.get("D", self.alpha) + beta_b = dt*self.tau_values.get("b", self.alpha) + Vu = equation.domain.spaces("HDiv") + VD = equation.domain.spaces("DG") + Vb = equation.domain.spaces("DG") + + # Check that the third field is buoyancy + if not equation.field_names[2] == 'b' and not (equation.field_names[2] == 'b_e' and equivalent_buoyancy): + raise NotImplementedError("Field 'b' or 'b_e' must exist to use the thermal linear solver in the SIQN scheme") + + # Split up the rhs vector + self.xrhs = Function(equation.function_space) + u_in, D_in, b_in = split(self.xrhs)[0:3] + + # Build the reduced function space for u, D + M = equation.function_space + w, phi, gamma = TestFunctions(M) + u, D, b = TrialFunctions(M) + + # Get background buoyancy and depth + Dbar, bbar = split(equation.X_ref)[1:3] + + + if equivalent_buoyancy: + # compute q_v using q_sat to partition q_t into q_v and q_c + self.q_sat_func = Function(VD) + self.qvbar = Function(VD) + qtbar = split(equation.X_ref)[3] + + # set up interpolators that use the X_ref values for D and b_e + self.q_sat_expr_interpolate = interpolate( + equation.compute_saturation(equation.X_ref), VD) + self.q_v_interpolate = interpolate( + conditional(qtbar < self.q_sat_func, qtbar, self.q_sat_func), + VD) + + # bbar was be_bar and here we correct to become bbar + bbar += equation.parameters.beta2 * self.qvbar + + n = FacetNormal(equation.domain.mesh) + + eqn = ( + inner(w, (u - u_in)) * dx + - beta_u * (D - Dbar) * div(w*bbar) * dx + + beta_u * jump(w*bbar, n) * avg(D-Dbar) * dS + - beta_u * 0.5 * Dbar * bbar * div(w) * dx + - beta_u * 0.5 * Dbar * b * div(w) * dx + - beta_u * 0.5 * bbar * div(w*(D-Dbar)) * dx + + beta_u * 0.5 * jump((D-Dbar)*w, n) * avg(bbar) * dS + + inner(phi, (D - D_in)) * dx + + beta_d * phi * div(Dbar*u) * dx + + gamma*(b - b_in + inner(u, grad(bbar))*beta_b) * dx + ) + + if 'coriolis' in equation.prescribed_fields._field_names: + f = equation.prescribed_fields('coriolis') + eqn += beta_u * f * inner(w, equation.domain.perp(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, + constant_jacobian=True) + self.uD_solver = LinearVariationalSolver(uD_problem, + solver_parameters=self.solver_parameters) + + # Log residuals on hybridized solver + self.log_ksp_residuals(self.uD_solver.snes.ksp) + + @timed_function("Gusto:UpdateReferenceProfiles") + def update_reference_profiles(self): + if self.equations.equivalent_buoyancy: + self.q_sat_func.assign(assemble(self.q_sat_expr_interpolate)) + self.qvbar.assign(assemble(self.q_v_interpolate)) + + @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) + + # Check that the b reference profile has been set + bbar = split(self.equations.X_ref)[2] + b = dy.subfunctions[2] + bbar_func = Function(b.function_space()).interpolate(bbar) + if bbar_func.dat.data.max() == 0 and bbar_func.dat.data.min() == 0: + logger.warning("The reference profile for b in the linear solver is zero. To set a non-zero profile add b to the set_reference_profiles argument.") + + with timed_region("Gusto:VelocityDepthSolve"): + logger.info('Thermal linear solver: mixed solve') + self.uD_solver.solve() + + dy.assign(self.uD) + + +class ThermalSWSolverMonoNew(object): """ A general object for solving mixed finite element linear problems. @@ -1017,8 +1168,6 @@ def __init__(self, equation, alpha, reference_dependent=True): lambda t: t.has_label(linearisation), lambda t: Term(t.get(linearisation).form, t.labels), drop) - print(residual.form.__str__()) - breakpoint() self.alpha = Constant(alpha) dt = equation.domain.dt @@ -1034,10 +1183,6 @@ def __init__(self, equation, alpha, reference_dependent=True): Leqn = residual.label_map( lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), map_if_false=drop) - print(aeqn.form.__str__()) - breakpoint() - print(Leqn.form.__str__()) - breakpoint() # Place to put result of solver self.dy = Function(W) From 1f636e48b32d84a65150d52b29949147c6db2847 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Fri, 8 Aug 2025 16:43:10 +0100 Subject: [PATCH 11/14] Remove declaration.. --- gusto/solvers/linear_solvers.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py index 35d900a84..cc9c38e05 100644 --- a/gusto/solvers/linear_solvers.py +++ b/gusto/solvers/linear_solvers.py @@ -30,8 +30,7 @@ __all__ = ["BoussinesqSolver", "LinearTimesteppingSolver", "CompressibleSolver", - "ThermalSWSolver", "ThermalSWSolverNew", "ThermalSWSolverMono", - "MoistConvectiveSWSolver", "MoistConvectiveSWSolverNew"] + "ThermalSWSolver", "ThermalSWSolverNew", "ThermalSWSolverMono"] class TimesteppingSolver(object, metaclass=ABCMeta): @@ -1004,7 +1003,7 @@ class ThermalSWSolverMono(TimesteppingSolver): """ # solver_parameters = { - # 'ksp_type': 'preonly', + #C 'ksp_type': 'preonly', # 'mat_type': 'matfree', # 'pc_type': 'python', # 'pc_python_type': 'firedrake.HybridizationPC', @@ -1020,6 +1019,8 @@ class ThermalSWSolverMono(TimesteppingSolver): solver_parameters = { 'ksp_error_if_not_converged': None, + 'ksp_rtol': 1e-11, + 'ksp_atol': 1e-11, 'pc_factor_shift_type': 'nonzero', 'snes_type': 'ksponly', } From a326eda519213b2a6d5b1eb28e1888f6c78ad954 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Thu, 14 Aug 2025 17:14:26 +0100 Subject: [PATCH 12/14] refactor monolithic thermal swe linear solver --- .../moist_convective_williamson_2.py | 14 +- .../shallow_water/moist_thermal_equivb_gw.py | 101 +++-- gusto/solvers/linear_solvers.py | 420 +++++++----------- gusto/solvers/preconditioners.py | 8 +- 4 files changed, 242 insertions(+), 301 deletions(-) diff --git a/examples/shallow_water/moist_convective_williamson_2.py b/examples/shallow_water/moist_convective_williamson_2.py index e639a9c77..c26fa4efa 100644 --- a/examples/shallow_water/moist_convective_williamson_2.py +++ b/examples/shallow_water/moist_convective_williamson_2.py @@ -13,7 +13,10 @@ """ from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter -from firedrake import SpatialCoordinate, sin, cos, exp, Function, errornorm, norm, VectorSpaceBasis +from firedrake import ( + SpatialCoordinate, sin, cos, exp, Function, errornorm, norm, + VectorSpaceBasis, COMM_WORLD +) from gusto import ( Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations, @@ -145,7 +148,7 @@ def sat_func(x_in): "pc_fieldsplit_type": "additive", "pc_fieldsplit_0_fields": "0,1", "pc_fieldsplit_1_fields": "2,3,4", - "fieldsplit_0": { + "fieldsplit_0": { # hybridisation on the (u,D) system 'ksp_monitor_true_residual': None, 'ksp_type': 'preonly', 'pc_type': 'python', @@ -162,7 +165,7 @@ def sat_func(x_in): } } }, - "fieldsplit_1": { + "fieldsplit_1": { # Don't touch the transported fields "ksp_type": "preonly", "pc_type": "none" }, @@ -170,13 +173,14 @@ def sat_func(x_in): # Provide callback for the nullspace of the trace system def trace_nullsp(T): - return VectorSpaceBasis(constant=True) + return VectorSpaceBasis(constant=True, comm=COMM_WORLD) appctx = {"trace_nullspace": trace_nullsp} + linear_solver = LinearTimesteppingSolver( eqns, alpha=0.5, reference_dependent=True, solver_parameters=solver_parameters, - options_prefix="swe_gen", appctx=appctx) + options_prefix="swe", appctx=appctx) # Physics schemes sat_adj = SWSaturationAdjustment( diff --git a/examples/shallow_water/moist_thermal_equivb_gw.py b/examples/shallow_water/moist_thermal_equivb_gw.py index ebe900a5d..45203f4c7 100644 --- a/examples/shallow_water/moist_thermal_equivb_gw.py +++ b/examples/shallow_water/moist_thermal_equivb_gw.py @@ -13,7 +13,7 @@ ThermalShallowWaterEquations, lonlatr_from_xyz, MeridionalComponent, GeneralIcosahedralSphereMesh, SubcyclingOptions, ZonalComponent, PartitionedCloud, RungeKuttaFormulation, SSPRK3, ThermalSWSolver, - SemiImplicitQuasiNewton, xyz_vector_from_lonlatr, ThermalSWSolverNew, + SemiImplicitQuasiNewton, xyz_vector_from_lonlatr, ThermalSWSolverMono, prognostic, time_derivative, coriolis, pressure_gradient, transport, drop ) @@ -69,7 +69,7 @@ def moist_thermal_gw( # Equation eqns = ThermalShallowWaterEquations( domain, parameters, fexpr=fexpr, - equivalent_buoyancy=False + equivalent_buoyancy=True ) # IO @@ -86,40 +86,69 @@ def moist_thermal_gw( DGUpwind(eqns, field_name) for field_name in eqns.field_names ] - linear_solver = ThermalSWSolver(eqns) - - import numpy as np - np.random.seed(6) - - xold = Function(eqns.function_space) - yold = Function(eqns.function_space) - - - for xsub in xold.subfunctions[:3]: - with xsub.dat.vec_wo as v: - v.array[:] = np.random.random_sample(v.array.shape) - - xnew = xold.copy(deepcopy=True) - ynew = yold.copy(deepcopy=True) - xmono = xold.copy(deepcopy=True) - ymono = yold.copy(deepcopy=True) - - old_solver = ThermalSWSolver(eqns, options_prefix="swe_old") - new_solver = ThermalSWSolverNew(eqns, options_prefix="swe_new") - mono_solver = ThermalSWSolverMono(eqns, alpha=0.5) - - old_solver.solve(xold, yold) - new_solver.solve(xnew, ynew) - mono_solver.solve(xmono, ymono) - - # for i, (xo, xn) in enumerate(zip(xold.subfunctions, xnew.subfunctions)): - # print(f"{i = }: {errornorm(xo, xn)/norm(xo) = :4e}") - for i, (yo, yn) in enumerate(zip(yold.subfunctions, ynew.subfunctions)): - print(f"Hybrid New {i = }| {norm(yo) = :2e} | {norm(yn) = :2e} | {errornorm(yo, yn)/norm(yo) = :4e}") - for i, (yo, ym) in enumerate(zip(yold.subfunctions, ymono.subfunctions)): - print(f"Monolithic {i = }| {norm(yo) = :2e} | {norm(ym) = :2e} | {errornorm(yo, ym)/norm(yo) = :4e}") - - from sys import exit; exit() + params = { + 'ksp_view': ':ksp_view.log', + 'ksp_error_if_not_converged': None, + 'mat_type': 'matfree', + 'ksp_type': 'preonly', + 'pc_type': 'fieldsplit', + 'pc_fieldsplit_type': 'additive', + 'pc_fieldsplit_0_fields': '0,1,2', + 'pc_fieldsplit_1_fields': '3', # do nothing for scalar field + 'fieldsplit_0': { + 'pc_type': 'fieldsplit', + 'pc_fieldsplit_type': 'schur', + 'pc_fieldsplit_schur_fact_type': 'full', + 'pc_fieldsplit_1_fields': '0,1', + 'pc_fieldsplit_0_fields': '2', # eliminate temperature + 'fieldsplit_L2': { + 'ksp_monitor': None, + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'firedrake.AssembledPC', + 'assembled_pc_type': 'bjacobi', + 'assembled_sub_pc_type': 'ilu', + }, + 'fieldsplit_1': { + 'ksp_monitor': None, + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'gusto.AuxiliaryPC', + 'aux': { + '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' + }, + 'mg_coarse': { + 'ksp_type': 'preonly', + 'pc_type': 'lu', + 'pc_factor_mat_solver_type': 'mumps', + }, + }, + }, + }, + }, + 'fieldsplit_L2': { + 'ksp_type': 'preonly', + 'pc_type': 'none', + }, + } + + linear_solver = ThermalSWSolverMono( + eqns, alpha=0.5, + options_prefix="swe", + # solver_parameters=params, + # overwrite_solver_parameters=True + ) # ------------------------------------------------------------------------ # # Timestepper diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py index cc9c38e05..0790af775 100644 --- a/gusto/solvers/linear_solvers.py +++ b/gusto/solvers/linear_solvers.py @@ -29,8 +29,12 @@ from abc import ABCMeta, abstractmethod, abstractproperty -__all__ = ["BoussinesqSolver", "LinearTimesteppingSolver", "CompressibleSolver", - "ThermalSWSolver", "ThermalSWSolverNew", "ThermalSWSolverMono"] +__all__ = [ + "LinearTimesteppingSolver", + "CompressibleSolver", + "BoussinesqSolver", + "ThermalSWSolver", + "ThermalSWSolverMono"] class TimesteppingSolver(object, metaclass=ABCMeta): @@ -60,6 +64,8 @@ def __init__(self, equations, alpha=0.5, tau_values=None, self.tau_values = tau_values if tau_values is not None else {} self.options_prefix = options_prefix + self._specialise_parameters() + if solver_parameters is not None: if not overwrite_solver_parameters: p = flatten_parameters(self.solver_parameters) @@ -102,6 +108,9 @@ def update_reference_profiles(self): def solve(self): pass + def _specialise_parameters(self): + pass + class CompressibleSolver(TimesteppingSolver): """ @@ -613,18 +622,31 @@ class ThermalSWSolver(TimesteppingSolver): """ + # solver_parameters = { + # 'ksp_type': 'preonly', + # 'mat_type': 'matfree', + # 'pc_type': 'python', + # 'pc_python_type': 'firedrake.HybridizationPC', + # 'hybridization': { + # 'ksp_type': 'preonly', + # 'pc_type': 'lu', + # 'pc_factor_mat_solver_type': 'mumps' + # } + # # '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'}} + # } + solver_parameters = { + 'ksp_error_if_not_converged': None, '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'}} + 'pc_type': 'lu', + 'pc_factor_mat_solver_type': 'mumps', + 'pc_factor_shift_type': 'nonzero', } @timed_function("Gusto:SolverSetup") @@ -775,7 +797,8 @@ def solve(self, xrhs, dy): b.assign(self.b) -class ThermalSWSolverNew(TimesteppingSolver): + +class ThermalSWSolverMono(TimesteppingSolver): """Linear solver object for the thermal shallow water equations. This solves a linear problem for the thermal shallow water @@ -789,254 +812,105 @@ class ThermalSWSolverNew(TimesteppingSolver): """ - # 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'}} - # } - - solver_parameters = { + 'ksp_error_if_not_converged': None, 'mat_type': 'matfree', 'ksp_type': 'preonly', - "ksp_view":":ksp_view.log", - "pc_type": "fieldsplit", - "pc_fieldsplit_type": "multiplicative", - "pc_fieldsplit_0_fields": "0,1", - "pc_fieldsplit_1_fields": "2", - "fieldsplit_0": { - 'ksp_monitor_true_residual': None, + 'pc_type': 'fieldsplit', + 'pc_fieldsplit_type': 'schur', + 'pc_fieldsplit_schur_fact_type': 'full', + 'pc_fieldsplit_1_fields': '0,1', + 'pc_fieldsplit_0_fields': '2', + 'fieldsplit_L2': { + 'ksp_monitor': None, 'ksp_type': 'preonly', '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' - } - } - }, - "fieldsplit_L2":{ - 'ksp_monitor': None, - "ksp_type": "preonly", - "pc_type": "python", 'pc_python_type': 'firedrake.AssembledPC', - 'assembled':{'pc_type':'lu'} + 'assembled_pc_type': 'bjacobi', + 'assembled_sub_pc_type': 'ilu', }, + 'fieldsplit_1': { + 'ksp_monitor': None, + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'gusto.AuxiliaryPC', + 'aux': { + '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', + }, + 'mg_coarse': { + 'ksp_type': 'preonly', + 'pc_type': 'lu', + 'pc_factor_mat_solver_type': 'mumps', + }, + }, + }, + } } + def _specialise_parameters(self): + if self.equations.equivalent_buoyancy: + uDb_parameters = self.solver_parameters + uDb_parameters.pop("mat_type") + self.solver_parameters = { + 'ksp_error_if_not_converged': None, + 'mat_type': 'matfree', + 'ksp_type': 'preonly', + 'pc_type': 'fieldsplit', + 'pc_fieldsplit_type': 'additive', + 'pc_fieldsplit_0_fields': '0,1,2', # split out (u, D, b) system + 'pc_fieldsplit_1_fields': '3', # do nothing for scalar field + 'fieldsplit_0': uDb_parameters, + 'fieldsplit_L2': { + 'ksp_type': 'preonly', + 'pc_type': 'none', + }, + } + @timed_function("Gusto:SolverSetup") def _setup_solver(self): - equation = self.equations # just cutting down line length a bit - equivalent_buoyancy = equation.equivalent_buoyancy - - dt = self.dt - beta_u = dt*self.tau_values.get("u", self.alpha) - beta_d = dt*self.tau_values.get("D", self.alpha) - beta_b = dt*self.tau_values.get("b", self.alpha) - Vu = equation.domain.spaces("HDiv") - VD = equation.domain.spaces("DG") - Vb = equation.domain.spaces("DG") - - # Check that the third field is buoyancy - if not equation.field_names[2] == 'b' and not (equation.field_names[2] == 'b_e' and equivalent_buoyancy): - raise NotImplementedError("Field 'b' or 'b_e' must exist to use the thermal linear solver in the SIQN scheme") - - # Split up the rhs vector - self.xrhs = Function(equation.function_space) - u_in, D_in, b_in = split(self.xrhs)[0:3] - - # Build the reduced function space for u, D - M = equation.function_space - w, phi, gamma = TestFunctions(M) - u, D, b = TrialFunctions(M) - - # Get background buoyancy and depth - Dbar, bbar = split(equation.X_ref)[1:3] - - # Approximate elimination of b - bval = -dot(u, grad(bbar))*beta_b + b_in - - if equivalent_buoyancy: - # compute q_v using q_sat to partition q_t into q_v and q_c - self.q_sat_func = Function(VD) - self.qvbar = Function(VD) - qtbar = split(equation.X_ref)[3] - - # set up interpolators that use the X_ref values for D and b_e - self.q_sat_expr_interpolate = interpolate( - equation.compute_saturation(equation.X_ref), VD) - self.q_v_interpolate = interpolate( - conditional(qtbar < self.q_sat_func, qtbar, self.q_sat_func), - VD) - - # bbar was be_bar and here we correct to become bbar - bbar += equation.parameters.beta2 * self.qvbar - - n = FacetNormal(equation.domain.mesh) - - eqn = ( - inner(w, (u - u_in)) * dx - - beta_u * (D - Dbar) * div(w*bbar) * dx - + beta_u * jump(w*bbar, n) * avg(D-Dbar) * dS - - beta_u * 0.5 * Dbar * bbar * div(w) * dx - - beta_u * 0.5 * Dbar * bval * div(w) * dx - - beta_u * 0.5 * bbar * div(w*(D-Dbar)) * dx - + beta_u * 0.5 * jump((D-Dbar)*w, n) * avg(bbar) * dS - + inner(phi, (D - D_in)) * dx - + beta_d * phi * div(Dbar*u) * dx - + gamma*(b - b_in + inner(u, grad(bbar))*beta_b) * dx - ) - - if 'coriolis' in equation.prescribed_fields._field_names: - f = equation.prescribed_fields('coriolis') - eqn += beta_u * f * inner(w, equation.domain.perp(u)) * dx + # Form for the full (u, D, b) space + eqn, bcs = self._build_monolithic_form() aeqn = lhs(eqn) Leqn = rhs(eqn) - # Place to put results of (u,D) solver - self.uD = Function(M) + # Form after analytically eliminating bouyancy + seqn, sbcs = self._build_schur_complement_form() + appctx = {"auxform": (seqn, sbcs)} - # Boundary conditions - bcs = [DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u']] + # Place to put results of (u,D) solver + self.uD = Function(self.equations.function_space) # Solver for u, D uD_problem = LinearVariationalProblem(aeqn, Leqn, self.uD, bcs=bcs, constant_jacobian=True) - - # 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) - # # Reconstruction of b - # b = TrialFunction(Vb) - # gamma = TestFunction(Vb) - - # u, D = self.uD.subfunctions - # self.b = Function(Vb) - - # b_eqn = gamma*(b - b_in + inner(u, grad(bbar))*beta_b) * dx - - # b_problem = LinearVariationalProblem(lhs(b_eqn), - # rhs(b_eqn), - # self.b, - # constant_jacobian=True) - # self.b_solver = LinearVariationalSolver(b_problem) + self.uD_solver = LinearVariationalSolver(uD_problem, appctx=appctx, + options_prefix=self.options_prefix, + solver_parameters=self.solver_parameters) # Log residuals on hybridized solver self.log_ksp_residuals(self.uD_solver.snes.ksp) - @timed_function("Gusto:UpdateReferenceProfiles") - def update_reference_profiles(self): - if self.equations.equivalent_buoyancy: - self.q_sat_func.assign(assemble(self.q_sat_expr_interpolate)) - self.qvbar.assign(assemble(self.q_v_interpolate)) - - @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) - - # Check that the b reference profile has been set - bbar = split(self.equations.X_ref)[2] - b = dy.subfunctions[2] - bbar_func = Function(b.function_space()).interpolate(bbar) - if bbar_func.dat.data.max() == 0 and bbar_func.dat.data.min() == 0: - logger.warning("The reference profile for b in the linear solver is zero. To set a non-zero profile add b to the set_reference_profiles argument.") - - with timed_region("Gusto:VelocityDepthSolve"): - logger.info('Thermal linear solver: mixed solve') - self.uD_solver.solve() - - # u1, D1 = self.uD.subfunctions - # u = dy.subfunctions[0] - # D = dy.subfunctions[1] - # b = dy.subfunctions[2] - # u.assign(u1) - # D.assign(D1) - dy.assign(self.uD) - - # with timed_region("Gusto:BuoyancyRecon"): - # logger.info('Thermal linear solver: buoyancy reconstruction') - # self.b_solver.solve() - - # b.assign(self.b) - -class ThermalSWSolverMono(TimesteppingSolver): - """Linear solver object for the thermal shallow water equations. - - This solves a linear problem for the thermal shallow water - equations with prognostic variables u (velocity), D (depth) and - either b (buoyancy) or b_e (equivalent buoyancy). It follows the - following strategy: - - (1) Eliminate b - (2) Solve the resulting system for (u, D) using a hybrid-mixed method - (3) Reconstruct b - - """ - - # solver_parameters = { - #C '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'}} - # } - - - solver_parameters = { - 'ksp_error_if_not_converged': None, - 'ksp_rtol': 1e-11, - 'ksp_atol': 1e-11, - 'pc_factor_shift_type': 'nonzero', - 'snes_type': 'ksponly', - } - - @timed_function("Gusto:SolverSetup") - def _setup_solver(self): - equation = self.equations # just cutting down line length a bit + def _build_monolithic_form(self): + equation = self.equations equivalent_buoyancy = equation.equivalent_buoyancy dt = self.dt beta_u = dt*self.tau_values.get("u", self.alpha) beta_d = dt*self.tau_values.get("D", self.alpha) beta_b = dt*self.tau_values.get("b", self.alpha) - Vu = equation.domain.spaces("HDiv") VD = equation.domain.spaces("DG") - Vb = equation.domain.spaces("DG") # Check that the third field is buoyancy if not equation.field_names[2] == 'b' and not (equation.field_names[2] == 'b_e' and equivalent_buoyancy): @@ -1048,13 +922,22 @@ def _setup_solver(self): # Build the reduced function space for u, D M = equation.function_space - w, phi, gamma = TestFunctions(M) - u, D, b = TrialFunctions(M) + tests = TestFunctions(M) + trials = TrialFunctions(M) + + w, phi, gamma = tests[:3] + u, D, b = trials[:3] + + if len(tests) > 3: + has_scalar_fields = True + q_tests = tests[3:] + q_trials = trials[3:] + else: + has_scalar_fields = False # Get background buoyancy and depth Dbar, bbar = split(equation.X_ref)[1:3] - if equivalent_buoyancy: # compute q_v using q_sat to partition q_t into q_v and q_c self.q_sat_func = Function(VD) @@ -1086,27 +969,58 @@ def _setup_solver(self): + gamma*(b - b_in + inner(u, grad(bbar))*beta_b) * dx ) + if has_scalar_fields: + for qu, qv in zip(q_trials, q_tests): + eqn += qu * qv * dx + if 'coriolis' in equation.prescribed_fields._field_names: f = equation.prescribed_fields('coriolis') eqn += beta_u * f * inner(w, equation.domain.perp(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, - constant_jacobian=True) - self.uD_solver = LinearVariationalSolver(uD_problem, - solver_parameters=self.solver_parameters) + return eqn, bcs - # Log residuals on hybridized solver - self.log_ksp_residuals(self.uD_solver.snes.ksp) + def _build_schur_complement_form(self): + equation = self.equations + + # Approximate elimination of b + dt = self.dt + beta_u = dt*self.tau_values.get("u", self.alpha) + beta_d = dt*self.tau_values.get("D", self.alpha) + beta_b = dt*self.tau_values.get("b", self.alpha) + + n = FacetNormal(equation.domain.mesh) + Vu = equation.domain.spaces("HDiv") + VD = equation.domain.spaces("DG") + M_ = Vu*VD + u_, D_ = TrialFunctions(M_) + w_, phi_ = TestFunctions(M_) + + Dref, bref = equation.X_ref.subfunctions[1:3] + b_ = -dot(u_, grad(bref))*beta_b + + seqn = ( + inner(w_, u_) * dx + - beta_u * D_ * div(w_*bref) * dx + + beta_u * jump(w_*bref, n) * avg(D_) * dS + # - beta_u * 0.5 * Dref * bref * div(w_) * dx + - beta_u * 0.5 * Dref * b_ * div(w_) * dx + - beta_u * 0.5 * bref * div(w_*D_) * dx + + beta_u * 0.5 * jump(D_*w_, n) * avg(bref) * dS + + inner(phi_, D_) * dx + + beta_d * phi_ * div(Dref*u_) * dx + ) + + if 'coriolis' in equation.prescribed_fields._field_names: + f = equation.prescribed_fields('coriolis') + seqn += beta_u * f * inner(w_, equation.domain.perp(u_)) * dx + + # Boundary conditions + sbcs = [DirichletBC(M_.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u']] + + return seqn, sbcs @timed_function("Gusto:UpdateReferenceProfiles") def update_reference_profiles(self): @@ -1126,18 +1040,7 @@ def solve(self, xrhs, dy): :class:`MixedFunctionSpace`. """ self.xrhs.assign(xrhs) - - # Check that the b reference profile has been set - bbar = split(self.equations.X_ref)[2] - b = dy.subfunctions[2] - bbar_func = Function(b.function_space()).interpolate(bbar) - if bbar_func.dat.data.max() == 0 and bbar_func.dat.data.min() == 0: - logger.warning("The reference profile for b in the linear solver is zero. To set a non-zero profile add b to the set_reference_profiles argument.") - - with timed_region("Gusto:VelocityDepthSolve"): - logger.info('Thermal linear solver: mixed solve') - self.uD_solver.solve() - + self.uD_solver.solve() dy.assign(self.uD) @@ -1223,7 +1126,6 @@ def solve(self, xrhs, dy): dy.assign(self.dy) - class LinearTimesteppingSolver(object): """ A general object for solving mixed finite element linear problems. diff --git a/gusto/solvers/preconditioners.py b/gusto/solvers/preconditioners.py index 62b4cf0c6..054b1c5c3 100644 --- a/gusto/solvers/preconditioners.py +++ b/gusto/solvers/preconditioners.py @@ -12,7 +12,13 @@ from numpy import arange -__all__ = ["VerticalHybridizationPC", "SlateSchurPC"] +__all__ = ["VerticalHybridizationPC", "SlateSchurPC", "AuxiliaryPC"] + + +class AuxiliaryPC(AuxiliaryOperatorPC): + def form(self, pc, test, trial): + a, bcs = self.get_appctx(pc)['auxform'] + return (a(test, trial), bcs) class SlateSchurPC(AuxiliaryOperatorPC): From fd6c7553fd0bcb9ed9ab1aa398e489905485f9c8 Mon Sep 17 00:00:00 2001 From: JHopeCollins Date: Fri, 12 Sep 2025 12:06:10 +0100 Subject: [PATCH 13/14] expunge shallow water linear solvers --- .../linear_thermal_galewsky_jet.py | 58 +- .../moist_convective_williamson_2.py | 17 +- .../shallow_water/moist_thermal_equivb_gw.py | 16 +- .../shallow_water/thermal_williamson_2.py | 59 +- gusto/equations/shallow_water_equations.py | 44 +- gusto/solvers/linear_solvers.py | 528 +----------------- 6 files changed, 167 insertions(+), 555 deletions(-) diff --git a/examples/shallow_water/linear_thermal_galewsky_jet.py b/examples/shallow_water/linear_thermal_galewsky_jet.py index a5622c055..6f92a5242 100644 --- a/examples/shallow_water/linear_thermal_galewsky_jet.py +++ b/examples/shallow_water/linear_thermal_galewsky_jet.py @@ -15,7 +15,7 @@ Domain, IO, OutputParameters, SemiImplicitQuasiNewton, DefaultTransport, DGUpwind, ForwardEuler, ShallowWaterParameters, NumericalIntegral, LinearThermalShallowWaterEquations, GeneralIcosahedralSphereMesh, - ZonalComponent, ThermalSWSolver, lonlatr_from_xyz, xyz_vector_from_lonlatr, + ZonalComponent, LinearTimesteppingSolver, lonlatr_from_xyz, xyz_vector_from_lonlatr, RelativeVorticity, MeridionalComponent ) @@ -84,7 +84,61 @@ def linear_thermal_galewsky_jet( transport_methods = [DefaultTransport(eqns, "D"), DGUpwind(eqns, "b")] # Linear solver - linear_solver = ThermalSWSolver(eqns) + solver_parameters = { + 'ksp_monitor_true_residual': None, + 'ksp_view': ':ksp_view.log', + 'ksp_error_if_not_converged': None, + 'mat_type': 'matfree', + 'ksp_type': 'preonly', + 'pc_type': 'fieldsplit', + 'pc_fieldsplit_type': 'schur', + 'pc_fieldsplit_schur_fact_type': 'full', + 'pc_fieldsplit_1_fields': '0,1', + 'pc_fieldsplit_0_fields': '2', # eliminate temperature + 'fieldsplit_L2': { + 'ksp_monitor': None, + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'firedrake.AssembledPC', + 'assembled_pc_type': 'bjacobi', + 'assembled_sub_pc_type': 'ilu', + }, + 'fieldsplit_1': { + 'ksp_monitor': None, + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'gusto.AuxiliaryPC', + 'aux': { + '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' + }, + 'mg_coarse': { + 'ksp_type': 'preonly', + 'pc_type': 'lu', + 'pc_factor_mat_solver_type': 'mumps', + }, + }, + }, + }, + } + + appctx = {'auxform': eqns.schur_complement_form(alpha=0.5)} + + linear_solver = LinearTimesteppingSolver( + eqns, alpha=0.5, options_prefix="swe", + solver_parameters=solver_parameters, + appctx=appctx + ) # Time stepper stepper = SemiImplicitQuasiNewton( diff --git a/examples/shallow_water/moist_convective_williamson_2.py b/examples/shallow_water/moist_convective_williamson_2.py index c26fa4efa..1ccecf297 100644 --- a/examples/shallow_water/moist_convective_williamson_2.py +++ b/examples/shallow_water/moist_convective_williamson_2.py @@ -25,7 +25,6 @@ RelativeVorticity, SWSaturationAdjustment, WaterVapour, CloudWater, Rain, GeneralIcosahedralSphereMesh, xyz_vector_from_lonlatr ) -from gusto.core.labels import time_derivative, prognostic, coriolis, pressure_gradient, transport moist_convect_williamson_2_defaults = { 'ncells_per_edge': 16, # number of cells per icosahedron edge @@ -88,19 +87,9 @@ def moist_convect_williamson_2( WaterVapour(space='DG'), CloudWater(space='DG'), Rain(space='DG') ] - def linearisation_map(term): - if (term.get(prognostic) in ['u', 'D'] - and (any(term.has_label(time_derivative, coriolis, pressure_gradient)) - or (term.get(prognostic) in ['D'] and term.has_label(transport)))): - return True - elif term.has_label(time_derivative): - return True - else: - return False - eqns = ShallowWaterEquations( domain, parameters, fexpr=fexpr, u_transport_option=u_eqn_type, - active_tracers=tracers, linearisation_map=linearisation_map + active_tracers=tracers ) # IO @@ -146,8 +135,8 @@ def sat_func(x_in): 'ksp_type': 'preonly', "pc_type": "fieldsplit", "pc_fieldsplit_type": "additive", - "pc_fieldsplit_0_fields": "0,1", - "pc_fieldsplit_1_fields": "2,3,4", + "pc_fieldsplit_0_fields": "0,1", # (u, D) + "pc_fieldsplit_1_fields": "2,3,4", # (scalars,) "fieldsplit_0": { # hybridisation on the (u,D) system 'ksp_monitor_true_residual': None, 'ksp_type': 'preonly', diff --git a/examples/shallow_water/moist_thermal_equivb_gw.py b/examples/shallow_water/moist_thermal_equivb_gw.py index 45203f4c7..410dfe397 100644 --- a/examples/shallow_water/moist_thermal_equivb_gw.py +++ b/examples/shallow_water/moist_thermal_equivb_gw.py @@ -6,16 +6,13 @@ from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter from firedrake import ( SpatialCoordinate, pi, sqrt, min_value, cos, Constant, Function, exp, sin, - norm, errornorm ) from gusto import ( Domain, IO, OutputParameters, DGUpwind, ShallowWaterParameters, ThermalShallowWaterEquations, lonlatr_from_xyz, MeridionalComponent, GeneralIcosahedralSphereMesh, SubcyclingOptions, ZonalComponent, - PartitionedCloud, RungeKuttaFormulation, SSPRK3, ThermalSWSolver, + PartitionedCloud, RungeKuttaFormulation, SSPRK3, LinearTimesteppingSolver, SemiImplicitQuasiNewton, xyz_vector_from_lonlatr, - ThermalSWSolverMono, prognostic, time_derivative, coriolis, - pressure_gradient, transport, drop ) moist_thermal_gw_defaults = { @@ -87,6 +84,7 @@ def moist_thermal_gw( ] params = { + 'ksp_monitor_true_residual': None, 'ksp_view': ':ksp_view.log', 'ksp_error_if_not_converged': None, 'mat_type': 'matfree', @@ -143,11 +141,11 @@ def moist_thermal_gw( }, } - linear_solver = ThermalSWSolverMono( - eqns, alpha=0.5, - options_prefix="swe", - # solver_parameters=params, - # overwrite_solver_parameters=True + appctx = {'auxform': eqns.schur_complement_form(alpha=0.5)} + + linear_solver = LinearTimesteppingSolver( + eqns, alpha=0.5, options_prefix="swe", + solver_parameters=params, appctx=appctx ) # ------------------------------------------------------------------------ # diff --git a/examples/shallow_water/thermal_williamson_2.py b/examples/shallow_water/thermal_williamson_2.py index 318f44962..16376ce77 100644 --- a/examples/shallow_water/thermal_williamson_2.py +++ b/examples/shallow_water/thermal_williamson_2.py @@ -17,7 +17,7 @@ Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, TrapeziumRule, ShallowWaterParameters, ThermalShallowWaterEquations, RelativeVorticity, PotentialVorticity, SteadyStateError, - ZonalComponent, MeridionalComponent, ThermalSWSolver, + ZonalComponent, MeridionalComponent, LinearTimesteppingSolver, xyz_vector_from_lonlatr, lonlatr_from_xyz, GeneralIcosahedralSphereMesh, SubcyclingOptions ) @@ -101,8 +101,61 @@ def thermal_williamson_2( DGUpwind(eqns, "b") ] - # Linear solver - linear_solver = ThermalSWSolver(eqns) + solver_parameters = { + 'ksp_monitor_true_residual': None, + 'ksp_view': ':ksp_view.log', + 'ksp_error_if_not_converged': None, + 'mat_type': 'matfree', + 'ksp_type': 'preonly', + 'pc_type': 'fieldsplit', + 'pc_fieldsplit_type': 'schur', + 'pc_fieldsplit_schur_fact_type': 'full', + 'pc_fieldsplit_1_fields': '0,1', + 'pc_fieldsplit_0_fields': '2', # eliminate temperature + 'fieldsplit_L2': { + 'ksp_monitor': None, + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'firedrake.AssembledPC', + 'assembled_pc_type': 'bjacobi', + 'assembled_sub_pc_type': 'ilu', + }, + 'fieldsplit_1': { + 'ksp_monitor': None, + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'gusto.AuxiliaryPC', + 'aux': { + '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' + }, + 'mg_coarse': { + 'ksp_type': 'preonly', + 'pc_type': 'lu', + 'pc_factor_mat_solver_type': 'mumps', + }, + }, + }, + }, + } + + appctx = {'auxform': eqns.schur_complement_form(alpha=0.5)} + + linear_solver = LinearTimesteppingSolver( + eqns, alpha=0.5, options_prefix="swe", + solver_parameters=solver_parameters, + appctx=appctx + ) # Time stepper stepper = SemiImplicitQuasiNewton( diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py index c152e5198..49cfe1b62 100644 --- a/gusto/equations/shallow_water_equations.py +++ b/gusto/equations/shallow_water_equations.py @@ -1,7 +1,7 @@ """Classes for defining variants of the shallow-water equations.""" -from firedrake import (inner, dx, div, FunctionSpace, FacetNormal, jump, avg, - dS, split, conditional, exp) +from firedrake import (inner, dx, div, FunctionSpace, FacetNormal, jump, avg, DirichletBC, + dS, split, conditional, exp, TrialFunctions, TestFunctions, dot, grad) from firedrake.fml import subject, drop, all_terms from gusto.core.labels import ( linearisation, pressure_gradient, coriolis, prognostic @@ -533,6 +533,46 @@ def compute_saturation(self, X): sat_expr = q0*H/(D+topog) * exp(nu*(1-b/g)) return sat_expr + def schur_complement_form(self, alpha=0.5, tau_values=None): + # Approximate elimination of b + domain = self.domain + dt = domain.dt + tau_values = tau_values or {} + beta_u = dt*tau_values.get("u", alpha) + beta_d = dt*tau_values.get("D", alpha) + beta_b = dt*tau_values.get("b", alpha) + + n = FacetNormal(domain.mesh) + Vu = domain.spaces("HDiv") + VD = domain.spaces("DG") + M_ = Vu*VD + u_, D_ = TrialFunctions(M_) + w_, phi_ = TestFunctions(M_) + + Dref, bref = self.X_ref.subfunctions[1:3] + b_ = -dot(u_, grad(bref))*beta_b + + seqn = ( + inner(w_, u_) * dx + - beta_u * D_ * div(w_*bref) * dx + + beta_u * jump(w_*bref, n) * avg(D_) * dS + # - beta_u * 0.5 * Dref * bref * div(w_) * dx + - beta_u * 0.5 * Dref * b_ * div(w_) * dx + - beta_u * 0.5 * bref * div(w_*D_) * dx + + beta_u * 0.5 * jump(D_*w_, n) * avg(bref) * dS + + inner(phi_, D_) * dx + + beta_d * phi_ * div(Dref*u_) * dx + ) + + if 'coriolis' in self.prescribed_fields._field_names: + f = self.prescribed_fields('coriolis') + seqn += beta_u * f * inner(w_, domain.perp(u_)) * dx + + # Boundary conditions + sbcs = [DirichletBC(M_.sub(0), bc.function_arg, bc.sub_domain) for bc in self.bcs['u']] + + return seqn, sbcs + class LinearThermalShallowWaterEquations(ThermalShallowWaterEquations): u""" diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py index 0790af775..f0b7e8288 100644 --- a/gusto/solvers/linear_solvers.py +++ b/gusto/solvers/linear_solvers.py @@ -8,18 +8,16 @@ from firedrake import ( split, LinearVariationalProblem, Constant, LinearVariationalSolver, TestFunctions, TrialFunctions, TestFunction, TrialFunction, lhs, - rhs, FacetNormal, div, dx, jump, avg, dS, dS_v, dS_h, ds_v, ds_t, ds_b, + rhs, FacetNormal, div, dx, jump, avg, dS_v, dS_h, ds_v, ds_t, ds_b, ds_tb, inner, action, dot, grad, Function, VectorSpaceBasis, cross, BrokenElement, FunctionSpace, MixedFunctionSpace, DirichletBC, as_vector, - assemble, conditional ) from firedrake.fml import Term, drop from firedrake.petsc import flatten_parameters -from firedrake.__future__ import interpolate from pyop2.profiling import timed_function, timed_region from gusto.equations.active_tracers import TracerVariableType -from gusto.core.logging import ( +from gusto.core.logging import ( # noqa: F401 logger, DEBUG, logging_ksp_monitor_true_residual, attach_custom_monitor ) @@ -32,9 +30,7 @@ __all__ = [ "LinearTimesteppingSolver", "CompressibleSolver", - "BoussinesqSolver", - "ThermalSWSolver", - "ThermalSWSolverMono"] + "BoussinesqSolver"] class TimesteppingSolver(object, metaclass=ABCMeta): @@ -608,524 +604,6 @@ def solve(self, xrhs, dy): b.assign(self.b) -class ThermalSWSolver(TimesteppingSolver): - """Linear solver object for the thermal shallow water equations. - - This solves a linear problem for the thermal shallow water - equations with prognostic variables u (velocity), D (depth) and - either b (buoyancy) or b_e (equivalent buoyancy). It follows the - following strategy: - - (1) Eliminate b - (2) Solve the resulting system for (u, D) using a hybrid-mixed method - (3) Reconstruct b - - """ - - # solver_parameters = { - # 'ksp_type': 'preonly', - # 'mat_type': 'matfree', - # 'pc_type': 'python', - # 'pc_python_type': 'firedrake.HybridizationPC', - # 'hybridization': { - # 'ksp_type': 'preonly', - # 'pc_type': 'lu', - # 'pc_factor_mat_solver_type': 'mumps' - # } - # # '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'}} - # } - - solver_parameters = { - 'ksp_error_if_not_converged': None, - 'ksp_type': 'preonly', - 'pc_type': 'lu', - 'pc_factor_mat_solver_type': 'mumps', - 'pc_factor_shift_type': 'nonzero', - } - - @timed_function("Gusto:SolverSetup") - def _setup_solver(self): - equation = self.equations # just cutting down line length a bit - equivalent_buoyancy = equation.equivalent_buoyancy - - dt = self.dt - beta_u = dt*self.tau_values.get("u", self.alpha) - beta_d = dt*self.tau_values.get("D", self.alpha) - beta_b = dt*self.tau_values.get("b", self.alpha) - Vu = equation.domain.spaces("HDiv") - VD = equation.domain.spaces("DG") - Vb = equation.domain.spaces("DG") - - # Check that the third field is buoyancy - if not equation.field_names[2] == 'b' and not (equation.field_names[2] == 'b_e' and equivalent_buoyancy): - raise NotImplementedError("Field 'b' or 'b_e' must exist to use the thermal linear solver in the SIQN scheme") - - # Split up the rhs vector - self.xrhs = Function(equation.function_space) - u_in, D_in, b_in = split(self.xrhs)[0:3] - - # Build the reduced function space for u, D - M = MixedFunctionSpace((Vu, VD)) - w, phi = TestFunctions(M) - u, D = TrialFunctions(M) - - # Get background buoyancy and depth - Dbar, bbar = split(equation.X_ref)[1:3] - - # Approximate elimination of b - b = -dot(u, grad(bbar))*beta_b + b_in - - if equivalent_buoyancy: - # compute q_v using q_sat to partition q_t into q_v and q_c - self.q_sat_func = Function(VD) - self.qvbar = Function(VD) - qtbar = split(equation.X_ref)[3] - - # set up interpolators that use the X_ref values for D and b_e - self.q_sat_expr_interpolate = interpolate( - equation.compute_saturation(equation.X_ref), VD) - self.q_v_interpolate = interpolate( - conditional(qtbar < self.q_sat_func, qtbar, self.q_sat_func), - VD) - - # bbar was be_bar and here we correct to become bbar - bbar += equation.parameters.beta2 * self.qvbar - - n = FacetNormal(equation.domain.mesh) - - eqn = ( - inner(w, (u - u_in)) * dx - - beta_u * (D - Dbar) * div(w*bbar) * dx - + beta_u * jump(w*bbar, n) * avg(D-Dbar) * dS - - beta_u * 0.5 * Dbar * bbar * div(w) * dx - - beta_u * 0.5 * Dbar * b * div(w) * dx - - beta_u * 0.5 * bbar * div(w*(D-Dbar)) * dx - + beta_u * 0.5 * jump((D-Dbar)*w, n) * avg(bbar) * dS - + inner(phi, (D - D_in)) * dx - + beta_d * phi * div(Dbar*u) * dx - ) - - if 'coriolis' in equation.prescribed_fields._field_names: - f = equation.prescribed_fields('coriolis') - eqn += beta_u * f * inner(w, equation.domain.perp(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, - constant_jacobian=True) - - # 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) - # Reconstruction of b - b = TrialFunction(Vb) - gamma = TestFunction(Vb) - - u, D = self.uD.subfunctions - self.b = Function(Vb) - - b_eqn = gamma*(b - b_in + inner(u, grad(bbar))*beta_b) * dx - - b_problem = LinearVariationalProblem(lhs(b_eqn), - rhs(b_eqn), - self.b, - constant_jacobian=True) - self.b_solver = LinearVariationalSolver(b_problem) - - # Log residuals on hybridized solver - self.log_ksp_residuals(self.uD_solver.snes.ksp) - - @timed_function("Gusto:UpdateReferenceProfiles") - def update_reference_profiles(self): - if self.equations.equivalent_buoyancy: - self.q_sat_func.assign(assemble(self.q_sat_expr_interpolate)) - self.qvbar.assign(assemble(self.q_v_interpolate)) - - @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) - - # Check that the b reference profile has been set - bbar = split(self.equations.X_ref)[2] - b = dy.subfunctions[2] - bbar_func = Function(b.function_space()).interpolate(bbar) - if bbar_func.dat.data.max() == 0 and bbar_func.dat.data.min() == 0: - logger.warning("The reference profile for b in the linear solver is zero. To set a non-zero profile add b to the set_reference_profiles argument.") - - with timed_region("Gusto:VelocityDepthSolve"): - logger.info('Thermal linear solver: mixed solve') - self.uD_solver.solve() - - u1, D1 = self.uD.subfunctions - u = dy.subfunctions[0] - D = dy.subfunctions[1] - b = dy.subfunctions[2] - u.assign(u1) - D.assign(D1) - - with timed_region("Gusto:BuoyancyRecon"): - logger.info('Thermal linear solver: buoyancy reconstruction') - self.b_solver.solve() - - b.assign(self.b) - - -class ThermalSWSolverMono(TimesteppingSolver): - """Linear solver object for the thermal shallow water equations. - - This solves a linear problem for the thermal shallow water - equations with prognostic variables u (velocity), D (depth) and - either b (buoyancy) or b_e (equivalent buoyancy). It follows the - following strategy: - - (1) Eliminate b - (2) Solve the resulting system for (u, D) using a hybrid-mixed method - (3) Reconstruct b - - """ - - solver_parameters = { - 'ksp_error_if_not_converged': None, - 'mat_type': 'matfree', - 'ksp_type': 'preonly', - 'pc_type': 'fieldsplit', - 'pc_fieldsplit_type': 'schur', - 'pc_fieldsplit_schur_fact_type': 'full', - 'pc_fieldsplit_1_fields': '0,1', - 'pc_fieldsplit_0_fields': '2', - 'fieldsplit_L2': { - 'ksp_monitor': None, - 'ksp_type': 'preonly', - 'pc_type': 'python', - 'pc_python_type': 'firedrake.AssembledPC', - 'assembled_pc_type': 'bjacobi', - 'assembled_sub_pc_type': 'ilu', - }, - 'fieldsplit_1': { - 'ksp_monitor': None, - 'ksp_type': 'preonly', - 'pc_type': 'python', - 'pc_python_type': 'gusto.AuxiliaryPC', - 'aux': { - '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', - }, - 'mg_coarse': { - 'ksp_type': 'preonly', - 'pc_type': 'lu', - 'pc_factor_mat_solver_type': 'mumps', - }, - }, - }, - } - } - - def _specialise_parameters(self): - if self.equations.equivalent_buoyancy: - uDb_parameters = self.solver_parameters - uDb_parameters.pop("mat_type") - self.solver_parameters = { - 'ksp_error_if_not_converged': None, - 'mat_type': 'matfree', - 'ksp_type': 'preonly', - 'pc_type': 'fieldsplit', - 'pc_fieldsplit_type': 'additive', - 'pc_fieldsplit_0_fields': '0,1,2', # split out (u, D, b) system - 'pc_fieldsplit_1_fields': '3', # do nothing for scalar field - 'fieldsplit_0': uDb_parameters, - 'fieldsplit_L2': { - 'ksp_type': 'preonly', - 'pc_type': 'none', - }, - } - - @timed_function("Gusto:SolverSetup") - def _setup_solver(self): - - # Form for the full (u, D, b) space - eqn, bcs = self._build_monolithic_form() - aeqn = lhs(eqn) - Leqn = rhs(eqn) - - # Form after analytically eliminating bouyancy - seqn, sbcs = self._build_schur_complement_form() - appctx = {"auxform": (seqn, sbcs)} - - # Place to put results of (u,D) solver - self.uD = Function(self.equations.function_space) - - # Solver for u, D - uD_problem = LinearVariationalProblem(aeqn, Leqn, self.uD, bcs=bcs, - constant_jacobian=True) - self.uD_solver = LinearVariationalSolver(uD_problem, appctx=appctx, - options_prefix=self.options_prefix, - solver_parameters=self.solver_parameters) - - # Log residuals on hybridized solver - self.log_ksp_residuals(self.uD_solver.snes.ksp) - - def _build_monolithic_form(self): - equation = self.equations - equivalent_buoyancy = equation.equivalent_buoyancy - - dt = self.dt - beta_u = dt*self.tau_values.get("u", self.alpha) - beta_d = dt*self.tau_values.get("D", self.alpha) - beta_b = dt*self.tau_values.get("b", self.alpha) - VD = equation.domain.spaces("DG") - - # Check that the third field is buoyancy - if not equation.field_names[2] == 'b' and not (equation.field_names[2] == 'b_e' and equivalent_buoyancy): - raise NotImplementedError("Field 'b' or 'b_e' must exist to use the thermal linear solver in the SIQN scheme") - - # Split up the rhs vector - self.xrhs = Function(equation.function_space) - u_in, D_in, b_in = split(self.xrhs)[0:3] - - # Build the reduced function space for u, D - M = equation.function_space - tests = TestFunctions(M) - trials = TrialFunctions(M) - - w, phi, gamma = tests[:3] - u, D, b = trials[:3] - - if len(tests) > 3: - has_scalar_fields = True - q_tests = tests[3:] - q_trials = trials[3:] - else: - has_scalar_fields = False - - # Get background buoyancy and depth - Dbar, bbar = split(equation.X_ref)[1:3] - - if equivalent_buoyancy: - # compute q_v using q_sat to partition q_t into q_v and q_c - self.q_sat_func = Function(VD) - self.qvbar = Function(VD) - qtbar = split(equation.X_ref)[3] - - # set up interpolators that use the X_ref values for D and b_e - self.q_sat_expr_interpolate = interpolate( - equation.compute_saturation(equation.X_ref), VD) - self.q_v_interpolate = interpolate( - conditional(qtbar < self.q_sat_func, qtbar, self.q_sat_func), - VD) - - # bbar was be_bar and here we correct to become bbar - bbar += equation.parameters.beta2 * self.qvbar - - n = FacetNormal(equation.domain.mesh) - - eqn = ( - inner(w, (u - u_in)) * dx - - beta_u * (D - Dbar) * div(w*bbar) * dx - + beta_u * jump(w*bbar, n) * avg(D-Dbar) * dS - - beta_u * 0.5 * Dbar * bbar * div(w) * dx - - beta_u * 0.5 * Dbar * b * div(w) * dx - - beta_u * 0.5 * bbar * div(w*(D-Dbar)) * dx - + beta_u * 0.5 * jump((D-Dbar)*w, n) * avg(bbar) * dS - + inner(phi, (D - D_in)) * dx - + beta_d * phi * div(Dbar*u) * dx - + gamma*(b - b_in + inner(u, grad(bbar))*beta_b) * dx - ) - - if has_scalar_fields: - for qu, qv in zip(q_trials, q_tests): - eqn += qu * qv * dx - - if 'coriolis' in equation.prescribed_fields._field_names: - f = equation.prescribed_fields('coriolis') - eqn += beta_u * f * inner(w, equation.domain.perp(u)) * dx - - # Boundary conditions - bcs = [DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u']] - - return eqn, bcs - - def _build_schur_complement_form(self): - equation = self.equations - - # Approximate elimination of b - dt = self.dt - beta_u = dt*self.tau_values.get("u", self.alpha) - beta_d = dt*self.tau_values.get("D", self.alpha) - beta_b = dt*self.tau_values.get("b", self.alpha) - - n = FacetNormal(equation.domain.mesh) - Vu = equation.domain.spaces("HDiv") - VD = equation.domain.spaces("DG") - M_ = Vu*VD - u_, D_ = TrialFunctions(M_) - w_, phi_ = TestFunctions(M_) - - Dref, bref = equation.X_ref.subfunctions[1:3] - b_ = -dot(u_, grad(bref))*beta_b - - seqn = ( - inner(w_, u_) * dx - - beta_u * D_ * div(w_*bref) * dx - + beta_u * jump(w_*bref, n) * avg(D_) * dS - # - beta_u * 0.5 * Dref * bref * div(w_) * dx - - beta_u * 0.5 * Dref * b_ * div(w_) * dx - - beta_u * 0.5 * bref * div(w_*D_) * dx - + beta_u * 0.5 * jump(D_*w_, n) * avg(bref) * dS - + inner(phi_, D_) * dx - + beta_d * phi_ * div(Dref*u_) * dx - ) - - if 'coriolis' in equation.prescribed_fields._field_names: - f = equation.prescribed_fields('coriolis') - seqn += beta_u * f * inner(w_, equation.domain.perp(u_)) * dx - - # Boundary conditions - sbcs = [DirichletBC(M_.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u']] - - return seqn, sbcs - - @timed_function("Gusto:UpdateReferenceProfiles") - def update_reference_profiles(self): - if self.equations.equivalent_buoyancy: - self.q_sat_func.assign(assemble(self.q_sat_expr_interpolate)) - self.qvbar.assign(assemble(self.q_v_interpolate)) - - @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) - self.uD_solver.solve() - dy.assign(self.uD) - - -class ThermalSWSolverMonoNew(object): - """ - A general object for solving mixed finite element linear problems. - - This linear solver object is general and is designed for use with different - equation sets, including with the non-linear shallow-water equations. It - forms the linear problem from the equations using FML. The linear system is - solved using a hybridised-mixed method. - """ - - solver_parameters = { - 'ksp_type': 'preonly', - 'pc_type': 'lu' - } - - def __init__(self, equation, alpha, reference_dependent=True): - """ - Args: - equation (:class:`PrognosticEquation`): the model's equation object. - alpha (float): the semi-implicit off-centring factor. A value of 1 - is fully-implicit. - reference_dependent: this indicates that the solver Jacobian should - be rebuilt if the reference profiles have been updated. - """ - self.reference_dependent = reference_dependent - - residual = equation.residual.label_map( - lambda t: t.has_label(linearisation), - lambda t: Term(t.get(linearisation).form, t.labels), - drop) - - self.alpha = Constant(alpha) - dt = equation.domain.dt - W = equation.function_space - beta = dt*self.alpha - - # Split up the rhs vector (symbolically) - self.xrhs = Function(W) - - aeqn = residual.label_map( - lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), - map_if_false=lambda t: beta*t) - Leqn = residual.label_map( - lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), - map_if_false=drop) - - # Place to put result of solver - self.dy = Function(W) - - # Solver - 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, - constant_jacobian=True) - - self.solver = LinearVariationalSolver(problem, - options_prefix='linear_solver') - - @timed_function("Gusto:UpdateReferenceProfiles") - def update_reference_profiles(self): - if self.reference_dependent: - self.solver.invalidate_jacobian() - - @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) - print(self.xrhs.__str__()) - print(self.dy.__str__()) - self.solver.solve() - dy.assign(self.dy) - - class LinearTimesteppingSolver(object): """ A general object for solving mixed finite element linear problems. From 86da3dcc5ab232993e722c3cf647d451c6311f52 Mon Sep 17 00:00:00 2001 From: JHopeCollins Date: Fri, 12 Sep 2025 13:30:41 +0100 Subject: [PATCH 14/14] boussinesq schur complement? --- .../skamarock_klemp_compressible.py | 63 ++++++++++++++++++- gusto/equations/boussinesq_equations.py | 58 ++++++++++++++++- 2 files changed, 118 insertions(+), 3 deletions(-) diff --git a/examples/boussinesq/skamarock_klemp_compressible.py b/examples/boussinesq/skamarock_klemp_compressible.py index 5fb9ce075..697498bc6 100644 --- a/examples/boussinesq/skamarock_klemp_compressible.py +++ b/examples/boussinesq/skamarock_klemp_compressible.py @@ -15,7 +15,7 @@ from gusto import ( Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, TrapeziumRule, SUPGOptions, Divergence, Perturbation, CourantNumber, - BoussinesqParameters, BoussinesqEquations, BoussinesqSolver, + BoussinesqParameters, BoussinesqEquations, BoussinesqSolver, LinearTimesteppingSolver, boussinesq_hydrostatic_balance ) @@ -91,7 +91,66 @@ def skamarock_klemp_compressible_bouss( ] # Linear solver - linear_solver = BoussinesqSolver(eqns) + solver_parameters = { + 'ksp_monitor_true_residual': None, + 'ksp_view': ':ksp_view.log', + 'ksp_error_if_not_converged': None, + 'mat_type': 'matfree', + 'ksp_type': 'preonly', + 'pc_type': 'fieldsplit', + 'pc_fieldsplit_type': 'schur', + 'pc_fieldsplit_schur_fact_type': 'full', + 'pc_fieldsplit_1_fields': '0,1', + 'pc_fieldsplit_0_fields': '2', # eliminate temperature + 'fieldsplit_theta': { + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'firedrake.AssembledPC', + 'assembled_pc_type': 'bjacobi', + 'assembled_sub_pc_type': 'ilu', + }, + 'fieldsplit_1': { + 'ksp_monitor': None, + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'gusto.AuxiliaryPC', + 'aux': { + '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' + }, + 'mg_coarse': { + 'ksp_type': 'preonly', + 'pc_type': 'lu', + 'pc_factor_mat_solver_type': 'mumps', + }, + }, + }, + }, + } + + def trace_nullsp(T): + return VectorSpaceBasis(constant=True) + + appctx = { + 'auxform': eqns.schur_complement_form(alpha=0.5), + "trace_nullspace": trace_nullsp, + } + + linear_solver = LinearTimesteppingSolver( + eqns, alpha=0.5, options_prefix="boussinesq", + solver_parameters=solver_parameters, + appctx=appctx + ) # Time stepper stepper = SemiImplicitQuasiNewton( diff --git a/gusto/equations/boussinesq_equations.py b/gusto/equations/boussinesq_equations.py index 24964bd3a..feb2cf967 100644 --- a/gusto/equations/boussinesq_equations.py +++ b/gusto/equations/boussinesq_equations.py @@ -1,6 +1,7 @@ """Defines the Boussinesq equations.""" -from firedrake import inner, dx, div, cross, split, as_vector +from firedrake import ( + inner, dx, div, cross, split, as_vector, TestFunctions, TrialFunctions, dot, grad, DirichletBC) from firedrake.fml import subject, all_terms from gusto.core.labels import ( prognostic, linearisation, @@ -217,6 +218,61 @@ def __init__(self, domain, parameters, # Add linearisations to equations self.residual = self.generate_linear_terms(residual, self.linearisation_map) + def schur_complement_form(self, alpha=0.5, tau_values=None): + domain = self.domain + dt = domain.dt + tau_values = tau_values or {} + # Set relaxation parameters. If an alternative has not been given, set + # to semi-implicit off-centering factor + beta_u = dt*tau_values.get("u", alpha) + beta_p = dt*tau_values.get("p", alpha) + beta_b = dt*tau_values.get("b", alpha) + Vu = domain.spaces("HDiv") + Vb = domain.spaces("theta") + Vp = domain.spaces("DG") + + # Build the reduced function space for u,p + M = Vu*Vp + w, phi = TestFunctions(M) + u, p = TrialFunctions(M) + + # Get background fields + bbar = split(self.X_ref)[2] + + # Analytical (approximate) elimination of theta + k = self.domain.k # Upward pointing unit vector + b = -dot(k, u)*dot(k, grad(bbar))*beta_b + + # vertical projection + def V(u): + return k*inner(u, k) + + seqn = ( + inner(w, u)*dx + - beta_u*div(w)*p*dx + - beta_u*inner(w, k)*b*dx + ) + + if self.compressible: + cs = self.parameters.cs + seqn += phi * p * dx + beta_p * phi * cs**2 * div(u) * dx + else: + seqn += phi * div(u) * dx + + if hasattr(self, "mu"): + seqn += dt*self.mu*inner(w, k)*inner(u, k)*dx + + if self.parameters.Omega is not None: + Omega = as_vector((0, 0, self.parameter.Omega)) + seqn += inner(w, cross(2*Omega, u))*dx + + # Boundary conditions (assumes extruded mesh) + # BCs are declared for the plain velocity space. As we need them in + # a mixed problem, we replicate the BCs but for subspace of M + sbcs = [DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.bcs['u']] + + return seqn, sbcs + class LinearBoussinesqEquations(BoussinesqEquations): """