Coverage for pySDC/helpers/pySDC_as_gusto_time_discretization.py: 98%
92 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1import firedrake as fd
3from gusto.time_discretisation.time_discretisation import TimeDiscretisation, wrapper_apply
4from gusto.core.labels import explicit
6from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
7from pySDC.implementations.controller_classes.controller_MPI import controller_MPI
8from pySDC.implementations.problem_classes.GenericGusto import GenericGusto, GenericGustoImex
9from pySDC.core.hooks import Hooks
10from pySDC.helpers.stats_helper import get_sorted
12import logging
13import numpy as np
16class LogTime(Hooks):
17 """
18 Utility hook for knowing how far we got when using adaptive step size selection.
19 """
21 def post_step(self, step, level_number):
22 """
23 Record the time at the end of the step, `L.time + L.dt`, with type `_time`.
25 Args:
26 step (pySDC.Step.step): the current step
27 level_number (int): the current level number
28 """
29 super().post_step(step, level_number)
30 L = step.levels[level_number]
31 self.add_to_stats(
32 process=step.status.slot,
33 process_sweeper=L.sweep.rank,
34 time=L.time,
35 level=-1,
36 iter=-1,
37 sweep=-1,
38 type='_time',
39 value=L.time + L.dt,
40 )
43class pySDC_integrator(TimeDiscretisation):
44 """
45 This class can be entered into Gusto as a time discretization scheme and will solve steps using pySDC.
46 It will construct a pySDC controller which can be used by itself and will be used within the time step when called
47 from Gusto. Access the controller via `pySDC_integrator.controller`. This class also has `pySDC_integrator.stats`,
48 which gathers all of the pySDC stats recorded in the hooks during every time step when used within Gusto.
50 This class supports subcycling with multi-step SDC. You can use pseudo-parallelism by simply giving `n_steps` > 1 or
51 do proper parallelism by giving a `controller_communicator` of kind `pySDC.FiredrakeEnsembleCommunicator` with the
52 appropriate size. You also have to toggle between pseudo and proper parallelism with `useMPIController`.
53 """
55 def __init__(
56 self,
57 description,
58 controller_params,
59 domain,
60 field_name=None,
61 solver_parameters=None,
62 options=None,
63 imex=False,
64 useMPIController=False,
65 n_steps=1,
66 controller_communicator=None,
67 ):
68 """
69 Initialization
71 Args:
72 description (dict): pySDC description
73 controller_params (dict): pySDC controller params
74 domain (:class:`Domain`): the model's domain object, containing the
75 mesh and the compatible function spaces.
76 field_name (str, optional): name of the field to be evolved.
77 Defaults to None.
78 solver_parameters (dict, optional): dictionary of parameters to
79 pass to the underlying solver. Defaults to None.
80 options (:class:`AdvectionOptions`, optional): an object containing
81 options to either be passed to the spatial discretisation, or
82 to control the "wrapper" methods, such as Embedded DG or a
83 recovery method. Defaults to None.
84 imex (bool): Whether to use IMEX splitting
85 useMPIController (bool): Whether to use the pseudo-parallel or proper parallel pySDC controller
86 n_steps (int): Number of steps done in parallel when using pseudo-parallel pySDC controller
87 controller_communicator (pySDC.FiredrakeEnsembleCommunicator, optional): Communicator for the proper parallel controller
88 """
90 self._residual = None
92 super().__init__(
93 domain=domain,
94 field_name=field_name,
95 solver_parameters=solver_parameters,
96 options=options,
97 )
99 self.description = description
100 self.controller_params = controller_params
101 self.timestepper = None
102 self.dt_next = None
103 self.imex = imex
104 self.useMPIController = useMPIController
105 self.controller_communicator = controller_communicator
107 if useMPIController:
108 assert (
109 controller_communicator is not None
110 ), 'You need to supply a communicator when using the MPI controller!'
111 assert (
112 type(self.controller_communicator).__name__ == 'FiredrakeEnsembleCommunicator'
113 ), f'Need to give a FiredrakeEnsembleCommunicator here, not {type(self.controller_communicator)}'
114 if n_steps > 1:
115 logging.getLogger(type(self).__name__).warning(
116 f'Warning: You selected {n_steps=}, which will be ignored when using the MPI controller!'
117 )
118 self.n_steps = controller_communicator.size
119 else:
120 self.n_steps = n_steps
122 def setup(self, equation, apply_bcs=True, *active_labels):
123 """
124 Set up the time discretisation for the equation and construct the pySDC controller.
126 Modifies `self.description` and `self.controller_params` in place: sets the problem class, `GenericGustoImex` if
127 any
128 term is labelled explicit or `imex` was set and `GenericGusto` otherwise, the problem parameters and the step
129 size
130 `domain.dt / n_steps`, and adds the `LogTime` hook.
132 Args:
133 equation (:class:`PrognosticEquation`): the model's equation
134 apply_bcs (bool, optional): whether to apply the equation's boundary conditions. Defaults to True.
135 *active_labels (:class:`Label`): labels indicating which terms of the equation to include
136 """
137 super().setup(equation, apply_bcs, *active_labels)
139 # Check if any terms are explicit
140 imex = any(t.has_label(explicit) for t in equation.residual) or self.imex
141 if imex:
142 self.description['problem_class'] = GenericGustoImex
143 else:
144 self.description['problem_class'] = GenericGusto
146 self.description['problem_params'] = {
147 'equation': equation,
148 'solver_parameters': self.solver_parameters,
149 'residual': self._residual,
150 **self.description['problem_params'],
151 }
152 self.description['level_params']['dt'] = float(self.domain.dt) / self.n_steps
154 # add utility hook required for step size adaptivity
155 hook_class = self.controller_params.get('hook_class', [])
156 if not type(hook_class) == list:
157 hook_class = [hook_class]
158 hook_class.append(LogTime)
159 self.controller_params['hook_class'] = hook_class
161 # prepare controller and variables
162 if self.useMPIController:
163 self.controller = controller_MPI(
164 comm=self.controller_communicator,
165 description=self.description,
166 controller_params=self.controller_params,
167 )
168 else:
169 self.controller = controller_nonMPI(
170 self.n_steps, description=self.description, controller_params=self.controller_params
171 )
173 self.prob = self.level.prob
174 self.sweeper = self.level.sweep
175 self.x0_pySDC = self.prob.dtype_u(self.prob.init)
176 self.t = 0
177 self.stats = {}
179 @property
180 def residual(self):
181 """Make sure the pySDC problem residual and this residual are the same"""
182 if hasattr(self, 'prob'):
183 return self.prob.residual
184 else:
185 return self._residual
187 @residual.setter
188 def residual(self, value):
189 """Make sure the pySDC problem residual and this residual are the same"""
190 if hasattr(self, 'prob'):
191 if self.useMPIController:
192 self.controller.S.levels[0].prob.residual = value
193 else:
194 for S in self.controller.MS:
195 S.levels[0].prob.residual = value
196 else:
197 self._residual = value
199 @property
200 def step(self):
201 """Get the first step on the controller"""
202 if self.useMPIController:
203 return self.controller.S
204 else:
205 return self.controller.MS[0]
207 @property
208 def level(self):
209 """Get the finest pySDC level"""
210 return self.step.levels[0]
212 @wrapper_apply
213 def apply(self, x_out, x_in):
214 """
215 Apply the time discretization to advance one whole time step.
217 Args:
218 x_out (:class:`Function`): the output field to be computed.
219 x_in (:class:`Function`): the input field.
220 """
221 self.x0_pySDC.functionspace.assign(x_in)
222 assert np.isclose(
223 self.level.params.dt * self.n_steps, float(self.dt)
224 ), 'Step sizes have diverged between pySDC and Gusto'
226 if self.dt_next is not None:
227 assert (
228 self.timestepper is not None
229 ), 'You need to set self.timestepper to the timestepper in order to facilitate adaptive step size selection here!'
230 self.timestepper.dt = fd.Constant(self.dt_next * self.n_steps)
231 self.t = self.timestepper.t
233 uend, _stats = self.controller.run(u0=self.x0_pySDC, t0=float(self.t), Tend=float(self.t + self.dt))
235 # update time variables
236 if not np.isclose(self.level.params.dt * self.n_steps, float(self.dt)):
237 self.dt_next = self.level.params.dt
239 self.t = get_sorted(_stats, type='_time', recomputed=False, comm=self.controller_communicator)[-1][1]
241 # update time of the Gusto stepper.
242 # After this step, the Gusto stepper updates its time again to arrive at the correct time
243 if self.timestepper is not None:
244 self.timestepper.t = fd.Constant(float(self.t) - float(self.dt))
246 self.dt = fd.Constant(self.level.params.dt * self.n_steps)
248 # update stats and output
249 self.stats = {**self.stats, **_stats}
250 x_out.assign(uend.functionspace)