Coverage for pySDC/implementations/sweeper_classes/delta_form_MPI.py: 92%
61 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-25 20:28 +0000
1r"""
2Node-parallel delta-form sweeper.
4Kept in its own module because importing it requires ``mpi4py``, which
5:mod:`~pySDC.implementations.sweeper_classes.delta_form` does not.
7The MPI sweeper assigns one collocation node per rank and therefore uses only the **diagonal** of
8:math:`Q^\Delta`. The delta form collapses accordingly: with
9:math:`\varepsilon_r = u_0 + \tau_r + \Delta t (Q f^k)_r - u^k_r` for the rank's own node,
11.. math::
12 \delta_r = \varepsilon_r + \Delta t Q^\Delta_{rr}\,\big(f(u^k_r + \delta_r) - f(u^k_r)\big),
14with no accumulation over other nodes. The node-local piece is identical to the serial case, so
15:meth:`DeltaFormMixin._solve_correction` is reused unchanged and all three strategies
16(``solve_system_delta``, ``linear_implicit``, substitution fallback) work here too.
18"""
20from mpi4py import MPI
22from pySDC.implementations.sweeper_classes.delta_form import DeltaFormMixin
23from pySDC.implementations.sweeper_classes.generic_implicit_MPI import generic_implicit_MPI
26class delta_implicit_MPI(DeltaFormMixin, generic_implicit_MPI):
27 """Delta-form counterpart of :class:`generic_implicit_MPI`. One collocation node per rank."""
29 def _set_work_scale(self, values):
30 """
31 Agree the divisor across the node communicator.
33 Each rank holds one collocation node, so left alone every rank would scale its corrections by
34 its own node's residual and quantise differently from a serial run doing the same work. The
35 scale is conceptually the residual's magnitude, which is a property of the whole sweep, so it
36 is reduced. One scalar per sweep.
38 Parameters
39 ----------
40 values : list
41 This rank's residual, in backend units.
42 """
43 super()._set_work_scale(values)
44 if self._work_dtype is not None and self._scales_corrections():
45 self._work_scale = self.comm.allreduce(self._work_scale, op=MPI.MAX)
47 def _residual_nodes(self):
48 r"""
49 Compute :math:`\varepsilon_r` at this rank's node, in backend precision.
51 Returned as a one-element list, so the multi-level machinery -- which is written for a list
52 of nodes -- reads the same here as it does serially. A residual handed down by a transfer is
53 preferred over rebuilding one, exactly as in the serial sweeper.
55 Returns
56 -------
57 list
58 This rank's collocation residual, as a single-element list.
59 """
60 if self.eps_in is not None:
61 return self.eps_in
63 lvl = self.level
64 eps = self.integrate()
65 eps += lvl.u[0]
66 eps -= lvl.u[self.rank + 1]
67 if lvl.tau[self.rank] is not None:
68 eps += lvl.tau[self.rank]
69 return [eps]
71 def _sweep_nodes(self):
72 """
73 Perform one delta-form sweep for this rank's collocation node.
75 Returns
76 -------
77 None
78 """
79 lvl = self.level
80 prob = lvl.prob
81 assert lvl.status.unlocked
82 self._delta_setup()
84 rank = self.rank
85 t_node = lvl.time + lvl.dt * self.coll.nodes[rank]
86 alpha = lvl.dt * self.QI[rank + 1, rank + 1]
88 u_old = prob.dtype_u(lvl.u[rank + 1])
89 f_old = prob.dtype_f(lvl.f[rank + 1])
91 residual = self._residual_nodes()
92 self._set_work_scale(residual)
93 delta = self._solve_correction(self._to_work(prob, residual[0]), alpha, u_old, f_old, t_node)
95 lvl.u[rank + 1] = u_old + self._to_backend(prob, self._to_work(prob, delta))
96 lvl.f[rank + 1] = prob.eval_f(lvl.u[rank + 1], t_node)
98 # The sweep itself needs no increment -- QI is diagonal here, so nothing accumulates across
99 # nodes. A level that inherited its residual does need one, to advance that residual, so it
100 # is evaluated exactly when there is something to advance.
101 if getattr(self, 'eps_in', None) is not None:
102 self._f_increment(prob, lvl.f[rank + 1], f_old, u_old, delta, t_node)
104 lvl.status.updated = True
105 return None
107 def advance_residual(self, eps, deltas, dfs):
108 r"""
109 Advance this rank's residual by an update.
111 The :math:`Q\,\Delta f` term couples all collocation nodes, so it is a reduction here
112 rather than a sum, in the same shape :meth:`generic_implicit_MPI.integrate` uses.
114 Parameters
115 ----------
116 eps : list
117 This rank's residual, as a single-element list.
118 deltas : list
119 The update applied to this rank's nodal value.
120 dfs : list
121 The resulting right-hand side increment at this rank's node.
123 Returns
124 -------
125 list
126 The advanced residual.
127 """
128 lvl = self.level
129 integral = lvl.prob.dtype_u(lvl.prob.init, val=0.0)
130 for m in range(self.coll.num_nodes):
131 recvBuf = integral if m == self.rank else None
132 self.comm.Reduce(lvl.dt * self.coll.Qmat[m + 1, self.rank + 1] * dfs[0], recvBuf, root=m, op=MPI.SUM)
133 return [self._to_work(lvl.prob, eps[0] - deltas[0] + integral)]
135 def compute_residual(self, stage=''):
136 """
137 Report the tracked residual, reduced over the node communicator.
139 The serial version takes a maximum over the nodes it holds; here each rank holds one, so the
140 same quantity is an ``allreduce``. Shaped after
141 :meth:`generic_implicit_MPI.compute_residual`.
143 Parameters
144 ----------
145 stage : str
146 The stage of the step this level belongs to.
148 Returns
149 -------
150 None
151 """
152 if self.eps_in is None:
153 return generic_implicit_MPI.compute_residual(self, stage=stage)
155 lvl = self.level
156 if stage in self.params.skip_residual_computation:
157 lvl.status.residual = 0.0 if lvl.status.residual is None else lvl.status.residual
158 return None
160 norm = abs(self.eps_in[0])
161 kind = lvl.params.residual_type
162 if kind.endswith('rel'):
163 norm = norm / abs(lvl.u[0])
164 if kind.startswith('full'):
165 lvl.status.residual = self.comm.allreduce(norm, op=MPI.MAX)
166 elif kind.startswith('last'):
167 lvl.status.residual = self.comm.bcast(norm, root=self.comm.size - 1)
168 else:
169 raise NotImplementedError(f'residual type "{kind}" not implemented!')
170 lvl.status.updated = False
171 return None