Coverage for pySDC/implementations/problem_classes/GenericGusto.py: 93%
110 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1from pySDC.core.problem import Problem, WorkCounter
2from pySDC.implementations.datatype_classes.firedrake_mesh import firedrake_mesh, IMEX_firedrake_mesh
3from gusto.core.labels import (
4 time_derivative,
5 implicit,
6 explicit,
7)
8from firedrake.fml import replace_subject, all_terms, drop
9import firedrake as fd
10import numpy as np
13class GenericGusto(Problem):
14 """
15 Problem class wrapping a Gusto (Firedrake) equation, with all terms of its residual treated implicitly.
17 Keep in mind that you probably want to use the pySDC-Gusto coupling via
18 the `pySDC_integrator` class in the helpers in order to get spatial methods rather than interfacing with this
19 class directly.
21 Gusto equations work by a residual, which is minimized in nonlinear solvers to obtain the right hand side
22 evaluation or the solution to (IMEX) Euler steps. You control what you solve for by manipulating labeled parts
23 of the residual.
24 """
26 dtype_u = firedrake_mesh
27 dtype_f = firedrake_mesh
28 rhs_n_labels = 1
30 def __init__(
31 self,
32 equation,
33 apply_bcs=True,
34 solver_parameters=None,
35 stop_at_divergence=False,
36 LHS_cache_size=12,
37 residual=None,
38 *active_labels,
39 ):
40 """
41 Initialisation
43 Args:
44 equation (:class:`PrognosticEquation`): the model's equation.
45 apply_bcs (bool, optional): whether to apply the equation's boundary
46 conditions. Defaults to True.
47 solver_parameters (dict, optional): Solver parameters for the nonlinear variational problems.
48 Defaults to GMRES with a block Jacobi preconditioner with ILU on the blocks.
49 stop_at_divergence (bool, optional): Whether to raise an error when the variational problems do not converge. Defaults to False
50 LHS_cache_size (int, optional): Size of the cache for solvers. Defaults to 12.
51 residual (Firedrake.form, optional): Overwrite the residual of the equation, e.g. after adding spatial methods. Defaults to None.
52 *active_labels (:class:`Label`): labels indicating which terms of
53 the equation to include.
54 """
56 self.equation = equation
57 self.residual = equation.residual if residual is None else residual
58 self.field_name = equation.field_name
59 self.fs = equation.function_space
60 self.idx = None
61 if solver_parameters is None:
62 # default solver parameters
63 solver_parameters = {'ksp_type': 'gmres', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}
64 self.solver_parameters = solver_parameters
65 self.stop_at_divergence = stop_at_divergence
67 # -------------------------------------------------------------------- #
68 # Setup caches
69 # -------------------------------------------------------------------- #
71 self.x_out = fd.Function(self.fs)
72 self.solvers = {}
73 self._u = fd.Function(self.fs)
75 super().__init__(self.fs)
76 self._makeAttributeAndRegister('LHS_cache_size', 'apply_bcs', localVars=locals(), readOnly=True)
77 self.work_counters['rhs'] = WorkCounter()
78 self.work_counters['ksp'] = WorkCounter()
79 self.work_counters['solver_setup'] = WorkCounter()
80 self.work_counters['solver'] = WorkCounter()
82 @property
83 def bcs(self):
84 """
85 Boundary conditions of the prognostic field of the equation.
87 Returns:
88 list of :class:`DirichletBC`: The boundary conditions, or None if `apply_bcs` is False
89 """
90 if not self.apply_bcs:
91 return None
92 else:
93 return self.equation.bcs[self.equation.field_name]
95 def eval_f(self, u, *args):
96 """
97 Evaluate the right hand side f(u) by solving M f = -R(u), where M is the mass form (the time derivative terms of
98 the residual) and R(u) are all other terms of the residual evaluated at `u`. The solver is set up once and
99 cached.
101 Args:
102 u (dtype_u): Solution at which to evaluate
103 *args: Not used, the time is not passed on to Gusto
105 Returns:
106 dtype_f: The right hand side
107 """
108 self._u.assign(u.functionspace)
110 if 'eval_rhs' not in self.solvers.keys():
111 residual = self.residual.label_map(
112 lambda t: t.has_label(time_derivative),
113 map_if_false=replace_subject(self._u, old_idx=self.idx),
114 map_if_true=drop,
115 )
116 mass_form = self.residual.label_map(
117 lambda t: t.has_label(time_derivative),
118 map_if_true=replace_subject(self.x_out, old_idx=self.idx),
119 map_if_false=drop,
120 )
122 problem = fd.NonlinearVariationalProblem((mass_form + residual).form, self.x_out, bcs=self.bcs)
123 solver_name = self.field_name + self.__class__.__name__
124 self.solvers['eval_rhs'] = fd.NonlinearVariationalSolver(
125 problem, solver_parameters=self.solver_parameters, options_prefix=solver_name
126 )
127 self.work_counters['solver_setup']()
129 self.solvers['eval_rhs'].solve()
130 self.work_counters['rhs']()
132 return self.dtype_f(self.x_out)
134 def solve_system(self, rhs, factor, u0, *args):
135 """
136 Solve the nonlinear system (M - factor f)(u) = M rhs, i.e. M u + factor R(u) = M rhs with R the terms of the
137 residual other than the time derivative, with a Firedrake nonlinear variational solver starting from `u0`.
139 The solvers are cached per `factor`. When the cache is full, the oldest one for a factor is evicted. If the
140 solver does not converge, the error is raised if `stop_at_divergence` is set and only logged otherwise.
142 Args:
143 rhs (dtype_u): Right hand side
144 factor (float): Prefactor of the right hand side f, e.g. the node-to-node step size
145 u0 (dtype_u): Initial guess
146 *args: Not used, typically the time
148 Returns:
149 dtype_u: The solution
150 """
151 self.x_out.assign(u0.functionspace) # set initial guess
152 self._u.assign(rhs.functionspace)
154 if factor not in self.solvers.keys():
155 if len(self.solvers) >= self.LHS_cache_size + self.rhs_n_labels:
156 self.solvers.pop(
157 [me for me in self.solvers.keys() if type(me) in [float, int, np.float64, np.float32]][0]
158 )
160 # setup left hand side (M - factor*f)(u)
161 # put in output variable
162 residual = self.residual.label_map(all_terms, map_if_true=replace_subject(self.x_out, old_idx=self.idx))
163 # multiply f by factor
164 residual = residual.label_map(
165 lambda t: t.has_label(time_derivative), map_if_false=lambda t: fd.Constant(factor) * t
166 )
168 # subtract right hand side
169 mass_form = self.residual.label_map(lambda t: t.has_label(time_derivative), map_if_false=drop)
170 residual -= mass_form.label_map(all_terms, map_if_true=replace_subject(self._u, old_idx=self.idx))
172 # construct solver
173 problem = fd.NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs)
174 solver_name = f'{self.field_name}-{self.__class__.__name__}-{factor}'
175 self.solvers[factor] = fd.NonlinearVariationalSolver(
176 problem, solver_parameters=self.solver_parameters, options_prefix=solver_name
177 )
178 self.work_counters['solver_setup']()
180 try:
181 self.solvers[factor].solve()
182 except fd.exceptions.ConvergenceError as error:
183 if self.stop_at_divergence:
184 raise error
185 else:
186 self.logger.debug(error)
188 self.work_counters['ksp'].niter += self.solvers[factor].snes.getLinearSolveIterations()
189 self.work_counters['solver']()
190 return self.dtype_u(self.x_out)
193class GenericGustoImex(GenericGusto):
194 """
195 Problem class wrapping a Gusto (Firedrake) equation, IMEX with the terms labeled implicit and explicit split.
196 """
198 dtype_f = IMEX_firedrake_mesh
199 rhs_n_labels = 2
201 def evaluate_labeled_term(self, u, label):
202 """
203 Evaluate the part of the right hand side made of the terms carrying `label`, by solving M x = -R_label(u), where
204 M is the mass form and R_label are the terms of the residual with `label` and without the time derivative label.
205 One solver per label is set up and cached.
207 Args:
208 u (dtype_u): Solution at which to evaluate
209 label (:class:`Label`): Gusto label, e.g. `implicit` or `explicit`
211 Returns:
212 firedrake.Function: The result, in the output buffer `self.x_out` that is overwritten by the next solve
213 """
214 self._u.assign(u.functionspace)
216 if label not in self.solvers.keys():
217 residual = self.residual.label_map(
218 lambda t: t.has_label(label) and not t.has_label(time_derivative),
219 map_if_true=replace_subject(self._u, old_idx=self.idx),
220 map_if_false=drop,
221 )
222 mass_form = self.residual.label_map(
223 lambda t: t.has_label(time_derivative),
224 map_if_true=replace_subject(self.x_out, old_idx=self.idx),
225 map_if_false=drop,
226 )
228 problem = fd.NonlinearVariationalProblem((mass_form + residual).form, self.x_out, bcs=self.bcs)
229 solver_name = self.field_name + self.__class__.__name__
230 self.solvers[label] = fd.NonlinearVariationalSolver(
231 problem, solver_parameters=self.solver_parameters, options_prefix=solver_name
232 )
233 self.work_counters['solver_setup']()
235 self.solvers[label].solve()
236 return self.x_out
238 def eval_f(self, u, *args):
239 """
240 Evaluate the right hand side with the terms labeled `implicit` in the implicit part and the terms labeled
241 `explicit` in the explicit part. Terms carrying neither label are left out.
243 Args:
244 u (dtype_u): Solution at which to evaluate
245 *args: Not used, the time is not passed on to Gusto
247 Returns:
248 dtype_f: The right hand side, with `impl` and `expl` parts
249 """
250 me = self.dtype_f(self.init)
251 me.impl.assign(self.evaluate_labeled_term(u, implicit))
252 me.expl.assign(self.evaluate_labeled_term(u, explicit))
253 self.work_counters['rhs']()
254 return me
256 def solve_system(self, rhs, factor, u0, *args):
257 """
258 Solve the nonlinear system (M - factor f_I)(u) = M rhs, i.e. M u + factor R_I(u) = M rhs with R_I the terms of
259 the residual labeled `implicit`, with a Firedrake nonlinear variational solver starting from `u0`. The explicit
260 terms are left out. Caching and convergence failures are handled as in `GenericGusto.solve_system`.
262 Args:
263 rhs (dtype_u): Right hand side
264 factor (float): Prefactor of the implicit part f_I, e.g. the node-to-node step size
265 u0 (dtype_u): Initial guess
266 *args: Not used, typically the time
268 Returns:
269 dtype_u: The solution
270 """
271 self.x_out.assign(u0.functionspace) # set initial guess
272 self._u.assign(rhs.functionspace)
274 if factor not in self.solvers.keys():
275 if len(self.solvers) >= self.LHS_cache_size + self.rhs_n_labels:
276 self.solvers.pop(
277 [me for me in self.solvers.keys() if type(me) in [float, int, np.float64, np.float32]][0]
278 )
280 # setup left hand side (M - factor*f_I)(u)
281 # put in output variable
282 residual = self.residual.label_map(
283 lambda t: t.has_label(time_derivative) or t.has_label(implicit),
284 map_if_true=replace_subject(self.x_out, old_idx=self.idx),
285 map_if_false=drop,
286 )
287 # multiply f_I by factor
288 residual = residual.label_map(
289 lambda t: t.has_label(implicit) and not t.has_label(time_derivative),
290 map_if_true=lambda t: fd.Constant(factor) * t,
291 )
293 # subtract right hand side
294 mass_form = self.residual.label_map(lambda t: t.has_label(time_derivative), map_if_false=drop)
295 residual -= mass_form.label_map(all_terms, map_if_true=replace_subject(self._u, old_idx=self.idx))
297 # construct solver
298 problem = fd.NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs)
299 solver_name = f'{self.field_name}-{self.__class__.__name__}-{factor}'
300 self.solvers[factor] = fd.NonlinearVariationalSolver(
301 problem, solver_parameters=self.solver_parameters, options_prefix=solver_name
302 )
303 self.work_counters['solver_setup']()
305 try:
306 self.solvers[factor].solve()
307 except fd.exceptions.ConvergenceError as error:
308 if self.stop_at_divergence:
309 raise error
310 else:
311 self.logger.debug(error)
313 self.work_counters['ksp'].niter += self.solvers[factor].snes.getLinearSolveIterations()
314 self.work_counters['solver']()
315 return self.dtype_u(self.x_out)