From cf6b462c9741898c25d40b4c3d90b41d1fdb0ff4 Mon Sep 17 00:00:00 2001 From: nhartney Date: Thu, 9 Nov 2023 20:22:11 +0000 Subject: [PATCH 1/8] start of linear solver for the thermal shallow water equations --- gusto/linear_solvers.py | 51 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index 55bee0616..b321b5692 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -550,6 +550,57 @@ 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 b (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'}} + } + + @timed_function("Gusto:SolverSetup") + def _setup_solver(self): + equation = self.equations # just cutting down line length a bit + dt = self.dt + beta_ = dt*self.alpha + Vu = equation.domain.spaces("HDiv") + VD = equation.domain.spaces("DG") + Vb = equation.domain.spaces("DG") + + # Store time-stepping coefficients as UFL Constants + beta = Constant(beta_) + + # Split up the rhs vector (symbolically) + self.xrhs = Function(self.equations.function_space) + u_in, D_in, b_in = split(self.xrhs) + + # Build the reduced function space for u, D + M = MixedFunctionSpace((Vu, VD)) + w, phi = TestFunctions(M) + u, D = TrialFunctions(M) + + + + class LinearTimesteppingSolver(object): """ A general object for solving mixed finite element linear problems. From c22719e6fafd7ceb85ac1425e9d175088bacfe8b Mon Sep 17 00:00:00 2001 From: nhartney Date: Fri, 10 Nov 2023 17:09:40 +0000 Subject: [PATCH 2/8] more work on the thermal SWEs linear solver --- gusto/linear_solvers.py | 90 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 88 insertions(+), 2 deletions(-) diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index b321b5692..ce0325511 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -9,7 +9,7 @@ split, LinearVariationalProblem, Constant, LinearVariationalSolver, TestFunctions, TrialFunctions, TestFunction, TrialFunction, lhs, rhs, FacetNormal, div, dx, jump, avg, dS_v, dS_h, ds_v, ds_t, ds_b, - ds_tb, inner, action, dot, grad, Function, VectorSpaceBasis, + ds_tb, dS, inner, action, dot, grad, Function, VectorSpaceBasis, BrokenElement, FunctionSpace, MixedFunctionSpace, DirichletBC ) from firedrake.fml import Term, drop @@ -24,7 +24,7 @@ from abc import ABCMeta, abstractmethod, abstractproperty -__all__ = ["IncompressibleSolver", "LinearTimesteppingSolver", "CompressibleSolver"] +__all__ = ["IncompressibleSolver", "LinearTimesteppingSolver", "CompressibleSolver", "ThermalSWSolver"] class TimesteppingSolver(object, metaclass=ABCMeta): @@ -598,7 +598,93 @@ def _setup_solver(self): w, phi = TestFunctions(M) u, D = TrialFunctions(M) + # Get background buoynacy + bbar = split(equation.X_ref)[2] + + # Approximate elimination of b + b = -u*grad(bbar)*beta + b_in + + # TODO: check about surface terms, and about D - D_in + n = FacetNormal(equation.domain.mesh) + H = equation.parameters('H') + eqn = ( + inner(w, (u - u_in)) * dx + - beta * D * div(w*bbar) * dx + + beta * avg(D) * jump(bbar*w, n) * dS + + beta * 0.5 * H * w * grad(bbar) * dx + - beta * 0.5 * H * b * div(w) * dx + + beta * 0.5 * w * D * grad(bbar) * dx + + inner(phi, (D - D_in)) * dx + + beta * phi * H * div(u) * dx + ) + + aeqn = lhs(eqn) + Leqn = rhs(eqn) + + # Place to put results of (u,D) solver + self.uD = Function(M) + + # Boundary conditions + # TODO: this is taken from the IncompressibleSolver - does it assume an + # extruded mesh? + bcs = [DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u']] + + # Solver for u, D + uD_problem = LinearVariationalProblem(aeqn, Leqn, self.uD, bcs=bcs) + + # Provide callback for the nullspace of the trace system + # TODO: what is this? + 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 + u*grad(bbar)*beta) * dx + + b_problem = LinearVariationalProblem(lhs(b_eqn), + rhs(b_eqn), + self.b) + self.b_solver = LinearVariationalSolver(b_problem) + + # Log residuals on hybridized solver + self.log_ksp_residuals(self.up_solver.snes.ksp) + @timed_function("Gusto:LinearSolve") + def solve(self, xrhs, dy): + """ + Solve the linear problem. + + Args: + xrhs (:class:`Function`): the right-hand side field in the + appropriate :class:`MixedFunctionSpace`. + dy (:class:`Function`): the resulting field in the appropriate + :class:`MixedFunctionSpace`. + """ + self.xrhs.assign(xrhs) + + with timed_region("Gusto:VelocityDepthSolve"): + logger.info('Thermal linear solver: mixed solve') + self.uD_solver.solve() + + u1, D1 = self.uD.subfunctions + u, D, b = dy.subfunctions + 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 LinearTimesteppingSolver(object): From 56b681ab438439b197365796c08fe39668903501 Mon Sep 17 00:00:00 2001 From: nhartney Date: Mon, 13 Nov 2023 17:11:59 +0000 Subject: [PATCH 3/8] thermal W2 with semi-implicit quasi-Newton stepper to test thermal SW linear solver --- examples/shallow_water/thermal_williamson2.py | 25 ++++++++++++++++--- gusto/linear_solvers.py | 15 +++++------ 2 files changed, 29 insertions(+), 11 deletions(-) diff --git a/examples/shallow_water/thermal_williamson2.py b/examples/shallow_water/thermal_williamson2.py index e9fcd94ac..9e1333b5b 100644 --- a/examples/shallow_water/thermal_williamson2.py +++ b/examples/shallow_water/thermal_williamson2.py @@ -6,7 +6,7 @@ # Test case parameters # ----------------------------------------------------------------- # -dt = 100 +dt = 500 if '--running-tests' in sys.argv: tmax = dt @@ -42,7 +42,7 @@ eqns = ShallowWaterEquations(domain, params, fexpr=fexpr, u_transport_option='vector_advection_form', thermal=True) # IO -dirname = "thermal_williamson2" +dirname = "thermal_williamson2_SIQN" output = OutputParameters( dirname=dirname, dumpfreq=dumpfreq, @@ -54,14 +54,26 @@ ShallowWaterPotentialEnergy(params), ShallowWaterPotentialEnstrophy(), SteadyStateError('u'), SteadyStateError('D'), - MeridionalComponent('u'), ZonalComponent('u')] + SteadyStateError('b'), MeridionalComponent('u'), + ZonalComponent('u')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) + +# Transport schemes +transported_fields = [TrapeziumRule(domain, "u"), + SSPRK3(domain, "D", fixed_subcycles=2), + SSPRK3(domain, "b", fixed_subcycles=2)] transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D"), DGUpwind(eqns, "b")] +# Linear solver +linear_solver = ThermalSWSolver(eqns) + # Time stepper -stepper = Timestepper(eqns, RK4(domain), io, spatial_methods=transport_methods) +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_methods, + linear_solver=linear_solver) +# stepper = Timestepper(eqns, RK4(domain), io, spatial_methods=transport_methods) # ----------------------------------------------------------------- # # Initial conditions @@ -92,6 +104,11 @@ D0.interpolate(Dexpr) b0.interpolate(bexpr) +# Set reference profiles +Dbar = Function(D0.function_space()).assign(H) +bbar = Function(b0.function_space()).interpolate(bexpr) +stepper.set_reference_profiles([('D', Dbar), ('b', bbar)]) + # ----------------------------------------------------------------- # # Run # ----------------------------------------------------------------- # diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index ce0325511..fd5d1409e 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -598,22 +598,23 @@ def _setup_solver(self): w, phi = TestFunctions(M) u, D = TrialFunctions(M) - # Get background buoynacy + # Get background buoyancy and depth bbar = split(equation.X_ref)[2] + Dbar = split(equation.X_ref)[1] # Approximate elimination of b - b = -u*grad(bbar)*beta + b_in + b = -dot(u, grad(bbar))*beta + b_in # TODO: check about surface terms, and about D - D_in n = FacetNormal(equation.domain.mesh) - H = equation.parameters('H') + H = Dbar eqn = ( inner(w, (u - u_in)) * dx - beta * D * div(w*bbar) * dx + beta * avg(D) * jump(bbar*w, n) * dS - + beta * 0.5 * H * w * grad(bbar) * dx + + beta * 0.5 * H * inner(w, grad(bbar)) * dx - beta * 0.5 * H * b * div(w) * dx - + beta * 0.5 * w * D * grad(bbar) * dx + + beta * 0.5 * D * inner(w, grad(bbar)) * dx + inner(phi, (D - D_in)) * dx + beta * phi * H * div(u) * dx ) @@ -648,7 +649,7 @@ def trace_nullsp(T): u, D = self.uD.subfunctions self.b = Function(Vb) - b_eqn = gamma*(b - b_in + u*grad(bbar)*beta) * dx + b_eqn = gamma*(b - b_in + inner(u, grad(bbar))*beta) * dx b_problem = LinearVariationalProblem(lhs(b_eqn), rhs(b_eqn), @@ -656,7 +657,7 @@ def trace_nullsp(T): self.b_solver = LinearVariationalSolver(b_problem) # Log residuals on hybridized solver - self.log_ksp_residuals(self.up_solver.snes.ksp) + self.log_ksp_residuals(self.uD_solver.snes.ksp) @timed_function("Gusto:LinearSolve") def solve(self, xrhs, dy): From edbff9729006c37feab6ee8fbbfd9441c013bcc0 Mon Sep 17 00:00:00 2001 From: nhartney Date: Tue, 14 Nov 2023 16:42:14 +0000 Subject: [PATCH 4/8] taking out boundary conditions and surface term we think should be zero --- examples/shallow_water/thermal_williamson2.py | 2 +- gusto/linear_solvers.py | 5 ++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/examples/shallow_water/thermal_williamson2.py b/examples/shallow_water/thermal_williamson2.py index 9e1333b5b..8b6338c78 100644 --- a/examples/shallow_water/thermal_williamson2.py +++ b/examples/shallow_water/thermal_williamson2.py @@ -42,7 +42,7 @@ eqns = ShallowWaterEquations(domain, params, fexpr=fexpr, u_transport_option='vector_advection_form', thermal=True) # IO -dirname = "thermal_williamson2_SIQN" +dirname = "thermal_williamson2_SIQN_Tues" output = OutputParameters( dirname=dirname, dumpfreq=dumpfreq, diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index fd5d1409e..a229d45c7 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -611,7 +611,7 @@ def _setup_solver(self): eqn = ( inner(w, (u - u_in)) * dx - beta * D * div(w*bbar) * dx - + beta * avg(D) * jump(bbar*w, n) * dS + # + beta * avg(D) * jump(bbar*w, n) * dS + beta * 0.5 * H * inner(w, grad(bbar)) * dx - beta * 0.5 * H * b * div(w) * dx + beta * 0.5 * D * inner(w, grad(bbar)) * dx @@ -631,10 +631,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, Leqn, self.uD, bcs=bcs) + uD_problem = LinearVariationalProblem(aeqn, Leqn, self.uD) # , bcs=bcs) # Provide callback for the nullspace of the trace system - # TODO: what is this? def trace_nullsp(T): return VectorSpaceBasis(constant=True) From efe720e1124c65e2903cc330d735ae9847ef074d Mon Sep 17 00:00:00 2001 From: nhartney Date: Tue, 21 Nov 2023 17:24:08 +0000 Subject: [PATCH 5/8] the perturbation to D is D-Dbar --- examples/shallow_water/thermal_williamson2.py | 1 + gusto/linear_solvers.py | 14 ++++++-------- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/examples/shallow_water/thermal_williamson2.py b/examples/shallow_water/thermal_williamson2.py index 8b6338c78..7d0ffd3e9 100644 --- a/examples/shallow_water/thermal_williamson2.py +++ b/examples/shallow_water/thermal_williamson2.py @@ -107,6 +107,7 @@ # Set reference profiles Dbar = Function(D0.function_space()).assign(H) bbar = Function(b0.function_space()).interpolate(bexpr) +ref_g = Function(b0.function_space()).assign(g) stepper.set_reference_profiles([('D', Dbar), ('b', bbar)]) # ----------------------------------------------------------------- # diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index a229d45c7..202aca71b 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -605,18 +605,16 @@ def _setup_solver(self): # Approximate elimination of b b = -dot(u, grad(bbar))*beta + b_in - # TODO: check about surface terms, and about D - D_in n = FacetNormal(equation.domain.mesh) H = Dbar eqn = ( inner(w, (u - u_in)) * dx - - beta * D * div(w*bbar) * dx - # + beta * avg(D) * jump(bbar*w, n) * dS - + beta * 0.5 * H * inner(w, grad(bbar)) * dx - - beta * 0.5 * H * b * div(w) * dx - + beta * 0.5 * D * inner(w, grad(bbar)) * dx + - beta * (D - Dbar) * div(w*bbar) * dx + + beta * 0.5 * Dbar * inner(w, grad(bbar)) * dx + - beta * 0.5 * Dbar * b * div(w) * dx + + beta * 0.5 * (D - Dbar) * inner(w, grad(bbar)) * dx + inner(phi, (D - D_in)) * dx - + beta * phi * H * div(u) * dx + + beta * phi * Dbar * div(u) * dx ) aeqn = lhs(eqn) @@ -631,7 +629,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, Leqn, self.uD, bcs=bcs) # Provide callback for the nullspace of the trace system def trace_nullsp(T): From 2d1f799898fdadbd7f185a57457e002aa5ff8b94 Mon Sep 17 00:00:00 2001 From: nhartney Date: Mon, 27 Nov 2023 12:20:31 +0000 Subject: [PATCH 6/8] allows the field to be more than u, D, b (so works when additional moist fields are included) --- gusto/linear_solvers.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index 202aca71b..7a93af40a 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -591,7 +591,9 @@ def _setup_solver(self): # Split up the rhs vector (symbolically) self.xrhs = Function(self.equations.function_space) - u_in, D_in, b_in = split(self.xrhs) + u_in = split(self.xrhs)[0] + D_in = split(self.xrhs)[1] + b_in = split(self.xrhs)[2] # Build the reduced function space for u, D M = MixedFunctionSpace((Vu, VD)) @@ -674,7 +676,9 @@ def solve(self, xrhs, dy): self.uD_solver.solve() u1, D1 = self.uD.subfunctions - u, D, b = dy.subfunctions + u = dy.subfunctions[0] + D = dy.subfunctions[1] + b = dy.subfunctions[2] u.assign(u1) D.assign(D1) From 971cb97cb546b7493f26cbc2c671a32ad214ae35 Mon Sep 17 00:00:00 2001 From: nhartney Date: Wed, 29 Nov 2023 14:48:19 +0000 Subject: [PATCH 7/8] use SIQN in thermal Williamson 2 --- examples/shallow_water/thermal_williamson2.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/examples/shallow_water/thermal_williamson2.py b/examples/shallow_water/thermal_williamson2.py index 7d0ffd3e9..342ee8b07 100644 --- a/examples/shallow_water/thermal_williamson2.py +++ b/examples/shallow_water/thermal_williamson2.py @@ -6,7 +6,7 @@ # Test case parameters # ----------------------------------------------------------------- # -dt = 500 +dt = 4000 if '--running-tests' in sys.argv: tmax = dt @@ -42,7 +42,7 @@ eqns = ShallowWaterEquations(domain, params, fexpr=fexpr, u_transport_option='vector_advection_form', thermal=True) # IO -dirname = "thermal_williamson2_SIQN_Tues" +dirname = "thermal_williamson2" output = OutputParameters( dirname=dirname, dumpfreq=dumpfreq, @@ -73,7 +73,6 @@ stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, transport_methods, linear_solver=linear_solver) -# stepper = Timestepper(eqns, RK4(domain), io, spatial_methods=transport_methods) # ----------------------------------------------------------------- # # Initial conditions @@ -107,7 +106,6 @@ # Set reference profiles Dbar = Function(D0.function_space()).assign(H) bbar = Function(b0.function_space()).interpolate(bexpr) -ref_g = Function(b0.function_space()).assign(g) stepper.set_reference_profiles([('D', Dbar), ('b', bbar)]) # ----------------------------------------------------------------- # From 0bf68c9e7cf9240e130af8997b586a6ecb3fcace Mon Sep 17 00:00:00 2001 From: nhartney Date: Wed, 29 Nov 2023 14:48:37 +0000 Subject: [PATCH 8/8] some lint fixes --- gusto/linear_solvers.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index 7a93af40a..d1017aa28 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -9,7 +9,7 @@ split, LinearVariationalProblem, Constant, LinearVariationalSolver, TestFunctions, TrialFunctions, TestFunction, TrialFunction, lhs, rhs, FacetNormal, div, dx, jump, avg, dS_v, dS_h, ds_v, ds_t, ds_b, - ds_tb, dS, inner, action, dot, grad, Function, VectorSpaceBasis, + ds_tb, inner, action, dot, grad, Function, VectorSpaceBasis, BrokenElement, FunctionSpace, MixedFunctionSpace, DirichletBC ) from firedrake.fml import Term, drop @@ -589,7 +589,7 @@ def _setup_solver(self): # Store time-stepping coefficients as UFL Constants beta = Constant(beta_) - # Split up the rhs vector (symbolically) + # Split up the rhs vector self.xrhs = Function(self.equations.function_space) u_in = split(self.xrhs)[0] D_in = split(self.xrhs)[1] @@ -607,8 +607,6 @@ def _setup_solver(self): # Approximate elimination of b b = -dot(u, grad(bbar))*beta + b_in - n = FacetNormal(equation.domain.mesh) - H = Dbar eqn = ( inner(w, (u - u_in)) * dx - beta * (D - Dbar) * div(w*bbar) * dx @@ -626,8 +624,6 @@ def _setup_solver(self): self.uD = Function(M) # Boundary conditions - # TODO: this is taken from the IncompressibleSolver - does it assume an - # extruded mesh? bcs = [DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u']] # Solver for u, D