Coverage for pySDC/implementations/sweeper_classes/delta_form_MPI.py: 77%
62 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +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:`r_m = u_0 + \tau_m + \Delta t (Q f^k)_m - u^k_m` for the rank's own node :math:`m`,
11.. math::
12 \delta_m = r_m + \Delta t Q^\Delta_{mm}\,\big(f(u^k_m + \delta_m) - f(u^k_m)\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 """
28 Delta-form counterpart of :class:`~pySDC.implementations.sweeper_classes.generic_implicit_MPI.generic_implicit_MPI`.
29 One collocation node per rank.
30 """
32 def _set_work_scale(self, values):
33 """
34 Agree the divisor across the node communicator.
36 Each rank holds one collocation node, so left alone every rank would scale its corrections by
37 its own node's residual and quantise differently from a serial run doing the same work. The
38 scale is conceptually the residual's magnitude, which is a property of the whole sweep, so it
39 is reduced. One scalar per sweep.
41 Parameters
42 ----------
43 values : list
44 This rank's residual, in backend units.
45 """
46 super()._set_work_scale(values)
47 if self._work_dtype is not None and self._scales_corrections():
48 self._work_scale = self.comm.allreduce(self._work_scale, op=MPI.MAX)
50 def _residual_nodes(self):
51 r"""
52 Compute :math:`r_m` at this rank's node, in backend precision.
54 Returned as a one-element list, so the multi-level machinery -- which is written for a list
55 of nodes -- reads the same here as it does serially. A residual handed down by a transfer is
56 preferred over rebuilding one, exactly as in the serial sweeper.
58 Returns
59 -------
60 list
61 This rank's collocation residual, as a single-element list.
62 """
63 if self.eps_in is not None:
64 return self.eps_in
66 lvl = self.level
67 eps = self.integrate()
68 eps += lvl.u[0]
69 eps -= lvl.u[self.rank + 1]
70 if lvl.tau[self.rank] is not None:
71 eps += lvl.tau[self.rank]
72 return [eps]
74 def _sweep_nodes(self):
75 """
76 Perform one delta-form sweep for this rank's collocation node.
78 Returns
79 -------
80 None
81 """
82 lvl = self.level
83 prob = lvl.prob
84 assert lvl.status.unlocked
85 self._delta_setup()
87 rank = self.rank
88 t_node = lvl.time + lvl.dt * self.coll.nodes[rank]
89 alpha = lvl.dt * self.QI[rank + 1, rank + 1]
91 u_old = prob.dtype_u(lvl.u[rank + 1])
92 f_old = prob.dtype_f(lvl.f[rank + 1])
94 residual = self._residual_nodes()
95 self._set_work_scale(residual)
96 # a copy, as in the serial sweeper: the solve modifies its right-hand side in place, and the
97 # residual may be the one a transfer handed down, which is advanced after the sweep
98 rhs_corr = self._to_work(prob, residual[0])
99 delta = self._solve_correction(type(rhs_corr)(rhs_corr), alpha, u_old, f_old, t_node)
101 lvl.u[rank + 1] = u_old + self._to_backend(prob, self._to_work(prob, delta))
102 lvl.f[rank + 1] = prob.eval_f(lvl.u[rank + 1], t_node)
104 # The sweep itself needs no increment -- QI is diagonal here, so nothing accumulates across
105 # nodes. A level that inherited its residual does need one, to advance that residual, so it
106 # is evaluated exactly when there is something to advance.
107 if getattr(self, 'eps_in', None) is not None:
108 self._f_increment(prob, lvl.f[rank + 1], f_old, u_old, delta, t_node)
110 lvl.status.updated = True
111 return None
113 def advance_residual(self, eps, deltas, dfs):
114 r"""
115 Advance this rank's residual by an update.
117 The :math:`Q\,\Delta f` term couples all collocation nodes, so it is a reduction here
118 rather than a sum, in the same shape :meth:`generic_implicit_MPI.integrate` uses.
120 Parameters
121 ----------
122 eps : list
123 This rank's residual, as a single-element list.
124 deltas : list
125 The update applied to this rank's nodal value.
126 dfs : list
127 The resulting right-hand side increment at this rank's node.
129 Returns
130 -------
131 list
132 The advanced residual.
133 """
134 lvl = self.level
135 integral = lvl.prob.dtype_u(lvl.prob.init, val=0.0)
136 for m in range(self.coll.num_nodes):
137 recvBuf = integral if m == self.rank else None
138 self.comm.Reduce(lvl.dt * self.coll.Qmat[m + 1, self.rank + 1] * dfs[0], recvBuf, root=m, op=MPI.SUM)
139 return [self._to_work(lvl.prob, eps[0] - deltas[0] + integral)]
141 def compute_residual(self, stage=''):
142 """
143 Report the tracked residual, reduced over the node communicator.
145 The serial version takes a maximum over the nodes it holds; here each rank holds one, so the
146 same quantity is an ``allreduce``. Shaped after
147 :meth:`generic_implicit_MPI.compute_residual`.
149 Parameters
150 ----------
151 stage : str
152 The stage of the step this level belongs to.
154 Returns
155 -------
156 None
157 """
158 if self.eps_in is None:
159 return generic_implicit_MPI.compute_residual(self, stage=stage)
161 lvl = self.level
162 if stage in self.params.skip_residual_computation:
163 lvl.status.residual = 0.0 if lvl.status.residual is None else lvl.status.residual
164 return None
166 norm = abs(self.eps_in[0])
167 kind = lvl.params.residual_type
168 if kind.endswith('rel'):
169 norm = norm / abs(lvl.u[0])
170 if kind.startswith('full'):
171 lvl.status.residual = self.comm.allreduce(norm, op=MPI.MAX)
172 elif kind.startswith('last'):
173 lvl.status.residual = self.comm.bcast(norm, root=self.comm.size - 1)
174 else:
175 raise NotImplementedError(f'residual type "{kind}" not implemented!')
176 lvl.status.updated = False
177 return None