Coverage for pySDC/implementations/sweeper_classes/generic_implicit_MPI.py: 97%
105 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1from mpi4py import MPI
3from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit
4from pySDC.core.sweeper import Sweeper, ParameterError
5import logging
8class SweeperMPI(Sweeper):
9 """
10 MPI based sweeper where each rank administers one collocation node. Adapt sweepers to MPI by use of multiple inheritance.
11 See for example the `generic_implicit_MPI` sweeper, which has a class definition:
13 .. code-block:: python
15 class generic_implicit_MPI(SweeperMPI, generic_implicit):
17 this means in inherits both from `SweeperMPI` and `generic_implicit`. The hierarchy works such that functions are first
18 called from `SweeperMPI` and then from `generic_implicit`. For instance, in the `__init__` function, the `SweeperMPI`
19 class adds a communicator and nothing else. The `generic_implicit` implicit class adds a preconditioner and so on.
20 It's a bit confusing because `self.params` is overwritten in the second call to the `__init__` of the core `sweeper`
21 class, but the `SweeperMPI` class adds parameters to the `params` dictionary, which will again be added in
22 `generic_implicit`.
23 """
25 def __init__(self, params, level):
26 """
27 Initialization routine for the sweeper
29 Args:
30 params: parameters for the sweeper
31 level (pySDC.Level.level): the level that uses this sweeper
32 """
33 self.logger = logging.getLogger('sweeper')
35 if 'comm' not in params.keys():
36 params['comm'] = MPI.COMM_WORLD
37 self.logger.debug('Using MPI.COMM_WORLD for the communicator because none was supplied in the params.')
38 super().__init__(params, level)
40 if self.params.comm.size != self.coll.num_nodes:
41 raise NotImplementedError(
42 f'The communicator in the {type(self).__name__} sweeper needs to have one rank for each node as of now! That means we need {self.coll.num_nodes} nodes, but got {self.params.comm.size} processes.'
43 )
45 # each rank uses only its diagonal entry of the preconditioners, so anything off the diagonal would be dropped
46 if not self.parallelizable:
47 raise ParameterError(
48 f'{type(self).__name__} needs diagonal preconditioners, but got QI={self.params.QI!r}'
49 + (f' and QE={self.params.QE!r}' if hasattr(self.params, 'QE') else '')
50 + '. Use a diagonal QI such as "MIN-SR-S", "MIN-SR-NS", "IEpar" or "MIN-SR-FLEX", and QE="PIC".'
51 )
53 @property
54 def comm(self):
55 """
56 The communicator across the collocation nodes, taken from the parameters.
58 Returns:
59 mpi4py.MPI.Intracomm: the communicator
60 """
61 return self.params.comm
63 @property
64 def rank(self):
65 """
66 The rank in the node communicator, which is also the index of the collocation node this rank administers.
68 Returns:
69 int: the rank
70 """
71 return self.comm.rank
73 def compute_end_point(self):
74 """
75 Compute u at the right point of the interval
77 The value uend computed here is a full evaluation of the Picard formulation unless do_full_update==False
79 Returns:
80 None
81 """
83 L = self.level
84 P = L.prob
86 # check if Mth node is equal to right point and do_coll_update is false, perform a simple copy
87 if self.coll.right_is_node and not self.params.do_coll_update:
88 # a copy is sufficient
89 root = self.comm.Get_size() - 1
90 if self.comm.rank == root:
91 L.uend = P.dtype_u(L.u[-1])
92 else:
93 L.uend = P.dtype_u(L.u[0])
94 self.comm.Bcast(L.uend, root=root)
95 else:
96 raise NotImplementedError('require last node to be identical with right interval boundary')
98 return None
100 def compute_residual(self, stage=None):
101 """
102 Computation of the residual using the collocation matrix Q
104 Args:
105 stage (str): The current stage of the step the level belongs to
106 """
108 L = self.level
110 # Check if we want to skip the residual computation to gain performance
111 # Keep in mind that skipping any residual computation is likely to give incorrect outputs of the residual!
112 if stage in self.params.skip_residual_computation:
113 L.status.residual = 0.0 if L.status.residual is None else L.status.residual
114 return None
116 # compute the residual for each node
118 # build QF(u)
119 res = self.integrate(last_only=L.params.residual_type[:4] == 'last')
120 res += L.u[0] - L.u[self.rank + 1]
121 # add tau if associated
122 if L.tau[self.rank] is not None:
123 res += L.tau[self.rank]
124 # use abs function from data type here
125 res_norm = abs(res)
127 # find maximal residual over the nodes
128 if L.params.residual_type == 'full_abs':
129 L.status.residual = self.comm.allreduce(res_norm, op=MPI.MAX)
130 elif L.params.residual_type == 'last_abs':
131 L.status.residual = self.comm.bcast(res_norm, root=self.comm.size - 1)
132 elif L.params.residual_type == 'full_rel':
133 L.status.residual = self.comm.allreduce(res_norm / abs(L.u[0]), op=MPI.MAX)
134 elif L.params.residual_type == 'last_rel':
135 L.status.residual = self.comm.bcast(res_norm / abs(L.u[0]), root=self.comm.size - 1)
136 else:
137 raise NotImplementedError(f'residual type \"{L.params.residual_type}\" not implemented!')
139 # indicate that the residual has seen the new values
140 L.status.updated = False
142 return None
144 def predict(self):
145 """
146 Predictor to fill values at nodes before first sweep
148 Default prediction for the sweepers, only copies the values to all collocation nodes
149 and evaluates the RHS of the ODE there
150 """
152 L = self.level
153 P = L.prob
155 # evaluate RHS at left point
156 L.f[0] = P.eval_f(L.u[0], L.time)
158 m = self.rank
160 if self.params.initial_guess == 'spread':
161 # copy u[0] to all collocation nodes, evaluate RHS
162 L.u[m + 1] = P.dtype_u(L.u[0])
163 L.f[m + 1] = P.eval_f(L.u[m + 1], L.time + L.dt * self.coll.nodes[m])
164 elif self.params.initial_guess == 'copy':
165 # copy u[0] and RHS evaluation to all collocation nodes
166 L.u[m + 1] = P.dtype_u(L.u[0])
167 L.f[m + 1] = P.dtype_f(L.f[0])
168 elif self.params.initial_guess == 'zero':
169 # zeros solution for u and RHS
170 L.u[m + 1] = P.dtype_u(init=P.init, val=0.0)
171 L.f[m + 1] = P.dtype_f(init=P.init, val=0.0)
172 else:
173 raise ParameterError(f'initial_guess option {self.params.initial_guess} not implemented')
175 # indicate that this level is now ready for sweeps
176 L.status.unlocked = True
177 L.status.updated = True
179 def communicate_tau_correction_for_full_interval(self):
180 """
181 Broadcast the tau correction of the full interval, stored in ``L.tau[-1]`` on the last rank, to all ranks.
183 Returns:
184 None
185 """
186 L = self.level
187 P = L.prob
188 if self.rank < self.comm.size - 1:
189 L.tau[-1] = P.u_init
190 self.comm.Bcast(L.tau[-1], root=self.comm.size - 1)
193class generic_implicit_MPI(SweeperMPI, generic_implicit):
194 """
195 Generic implicit sweeper parallelized across the nodes.
196 Please supply a communicator as `comm` to the parameters!
198 Attributes:
199 rank (int): MPI rank
200 """
202 def integrate(self, last_only=False):
203 """
204 Integrates the right-hand side
206 Args:
207 last_only (bool): Integrate only the last node for the residual or all of them
209 Returns:
210 list of dtype_u: containing the integral as values
211 """
212 L = self.level
213 P = L.prob
215 me = P.dtype_u(P.init, val=0.0)
216 for m in [self.coll.num_nodes - 1] if last_only else range(self.coll.num_nodes):
217 recvBuf = me if m == self.rank else None
218 self.comm.Reduce(
219 L.dt * self.coll.Qmat[m + 1, self.rank + 1] * L.f[self.rank + 1], recvBuf, root=m, op=MPI.SUM
220 )
222 return me
224 def update_nodes(self):
225 """
226 Update the u- and f-values at the collocation nodes -> corresponds to a single sweep over all nodes
228 Returns:
229 None
230 """
232 L = self.level
233 P = L.prob
235 # only if the level has been touched before
236 assert L.status.unlocked
238 # gather all terms which are known already (e.g. from the previous iteration)
239 # this corresponds to u0 + QF(u^k) - QdF(u^k) + tau
241 # get QF(u^k)
242 rhs = self.integrate()
244 rhs -= L.dt * self.QI[self.rank + 1, self.rank + 1] * L.f[self.rank + 1]
246 # add initial value
247 rhs += L.u[0]
248 # add tau if associated
249 if L.tau[self.rank] is not None:
250 rhs += L.tau[self.rank]
252 # build rhs, consisting of the known values from above and new values from previous nodes (at k+1)
254 # implicit solve with prefactor stemming from the diagonal of Qd
255 L.u[self.rank + 1] = P.solve_system(
256 rhs,
257 L.dt * self.QI[self.rank + 1, self.rank + 1],
258 L.u[self.rank + 1],
259 L.time + L.dt * self.coll.nodes[self.rank],
260 )
261 # update function values
262 L.f[self.rank + 1] = P.eval_f(L.u[self.rank + 1], L.time + L.dt * self.coll.nodes[self.rank])
264 # indicate presence of new values at this level
265 L.status.updated = True
267 return None
269 def compute_end_point(self):
270 """
271 Compute u at the right point of the interval
273 The value uend computed here is a full evaluation of the Picard formulation unless do_full_update==False
275 Returns:
276 None
277 """
279 L = self.level
280 P = L.prob
282 # check if Mth node is equal to right point and do_coll_update is false, perform a simple copy
283 if self.coll.right_is_node and not self.params.do_coll_update:
284 super().compute_end_point()
285 else:
286 L.uend = P.dtype_u(L.u[0])
287 self.comm.Allreduce(L.dt * self.coll.weights[self.rank] * L.f[self.rank + 1], L.uend, op=MPI.SUM)
288 L.uend += L.u[0]
290 # add up tau correction of the full interval (last entry)
291 if L.tau[self.rank] is not None:
292 self.communicate_tau_correction_for_full_interval()
293 L.uend += L.tau[-1]
294 return None