Coverage for pySDC/implementations/controller_classes/controller_ParaDiag_nonMPI.py: 100%
86 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1import numpy as np
3from pySDC.helpers.ParaDiagHelper import get_G_inv_matrix
4from pySDC.implementations.controller_classes.ParaDiag import ParaDiag
5from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
8class controller_ParaDiag_nonMPI(ParaDiag, controller_nonMPI):
9 """
11 ParaDiag controller with the time steps of a block emulated serially in one process.
13 This is `controller_nonMPI` with a different iteration: where PFASST sweeps and cascades through
14 the levels, ParaDiag diagonalizes across the steps. Everything around the iteration -- blocks,
15 windowing, restarts, convergence -- is the driver it inherits, which is why the dispatcher it
16 uses is still called `pfasst`.
18 This controller uses the increment formulation. That is to say, we setup the residual of the all at once problem,
19 put it on the right hand side, invert the ParaDiag preconditioner on the left-hand side to compute the increment
20 and then add the increment onto the solution. For this reason, we need to replace the solution values in the steps
21 with the residual values before the solves and then put the solution plus increment back into the steps. This is a
22 bit counter to what you expect when you access the `u` variable in the levels, but it is mathematically advantageous.
23 """
25 def __init__(self, num_procs, controller_params, description):
26 """
27 Initialization routine for ParaDiag controller
29 Args:
30 num_procs: number of parallel time steps (still serial, though), can be 1
31 controller_params: parameter set for the controller and the steps
32 description: all the parameters to set up the rest (levels, problems, transfer, ...)
33 """
34 self.prepare_ParaDiag_params(controller_params, description)
36 self.sweeper_params = description['sweeper_params']
37 self._G_inv_alpha = self.resolve_alpha(controller_params['alpha'], 0)
39 # the steps are copies of the first one, so give that one its own G^-1 before they are built
40 description['sweeper_params']['G_inv'] = get_G_inv_matrix(
41 0, num_procs, self._G_inv_alpha, description['sweeper_params']
42 )
44 super().__init__(num_procs, controller_params, description)
46 self.n_steps = num_procs
48 # every step but the first still has the first one's G^-1
49 self.set_G_inv(self._G_inv_alpha)
51 # ------------------------------------------------------------------ the ParaDiag iteration
53 def apply_matrix(self, mat, quantity):
54 """
55 Apply a matrix on the step level. Needs to be square. Puts the result back into the controller.
57 Args:
58 mat: square LxL matrix with L number of steps
59 quantity (str): The quantity the matrix is applied to, `'residual'` or `'increment'`
60 """
61 L = len(self.MS)
62 assert np.allclose(mat.shape, L)
63 assert len(mat.shape) == 2
65 level = self.MS[0].levels[0]
66 M = level.sweep.params.num_nodes
67 prob = level.prob
69 # buffer for storing the result
70 res = [
71 None,
72 ] * L
74 if quantity == 'residual':
75 me = [S.levels[0].residual for S in self.MS]
76 elif quantity == 'increment':
77 me = [S.levels[0].increment for S in self.MS]
78 else:
79 raise NotImplementedError
81 # compute matrix-vector product
82 for i in range(mat.shape[0]):
83 res[i] = [prob.u_init for _ in range(M)]
84 for j in range(mat.shape[1]):
85 for m in range(M):
86 res[i][m] += mat[i, j] * me[j][m]
88 # put the result in the "output"
89 for i in range(mat.shape[0]):
90 for m in range(M):
91 me[i][m] = res[i][m]
93 def compute_all_at_once_residual(self, local_MS_running):
94 """
95 This requires to communicate the solutions at the end of the steps to be the initial conditions for the next
96 steps. Afterwards, the residual can be computed locally on the steps.
98 Args:
99 local_MS_running (list): list of currently running steps
100 """
102 for S in local_MS_running:
103 # communicate initial conditions
104 S.levels[0].sweep.compute_end_point()
106 for hook in self.hooks:
107 hook.pre_comm(step=S, level_number=0)
109 if not S.status.first:
110 S.levels[0].u[0] = S.prev.levels[0].uend
112 for hook in self.hooks:
113 hook.post_comm(step=S, level_number=0, add_to_stats=True)
115 # compute residuals locally
116 S.levels[0].sweep.compute_residual()
118 def set_G_inv(self, alpha):
119 """
120 Give every step the G^-1 that belongs to where it sits in the block.
122 Args:
123 alpha (float): the alpha this G^-1 is built from
124 """
125 L = len(self.MS)
126 for l, S in enumerate(self.MS):
127 S.levels[0].sweep.set_G_inv(get_G_inv_matrix(l, L, alpha, self.sweeper_params))
129 def update_G_inv(self, k=0):
130 """
131 Rebuild G^-1 on every step if alpha changed with the iteration.
133 Args:
134 k (int): 0-based ParaDiag iteration index
135 """
136 alpha = self.get_alpha(k)
137 if alpha == self._G_inv_alpha:
138 return
139 self._G_inv_alpha = alpha
140 self.set_G_inv(alpha)
142 def update_solution(self, local_MS_running):
143 """
144 Since we solve for the increment, we need to update the solution between iterations by adding the increment.
146 Args:
147 local_MS_running (list): list of currently running steps
148 """
149 for S in local_MS_running:
150 for m in range(S.levels[0].sweep.coll.num_nodes):
151 S.levels[0].u[m + 1] += S.levels[0].increment[m]
153 def prepare_Jacobians(self, local_MS_running):
154 """
155 Average the solution over the steps of the block, node by node, for constructing average Jacobians.
157 Does nothing unless `average_jacobian` is set. Stores the list of averages as `u_avg` on the finest level of
158 every
159 running step, which all share the same list.
161 Args:
162 local_MS_running (list): list of currently running steps
163 """
164 # get solutions for constructing average Jacobians
165 if self.params.average_jacobian:
166 level = local_MS_running[0].levels[0]
167 M = level.sweep.coll.num_nodes
169 u_avg = [level.prob.dtype_u(level.prob.init, val=0) for _ in range(M)]
171 # communicate average solution
172 for S in local_MS_running:
173 for m in range(M):
174 u_avg[m] += S.levels[0].u[m + 1] / self.n_steps
176 # store the averaged solution in the steps
177 for S in local_MS_running:
178 S.levels[0].u_avg = u_avg
180 def it_ParaDiag(self, local_MS_running):
181 """
182 Do a single ParaDiag iteration. Does the following steps
183 - (1) Compute the residual of the all-at-once / composite collocation problem
184 - (2) Compute an FFT in time to diagonalize the preconditioner
185 - (3) Solve the collocation problems locally on the steps for the increment
186 - (4) Compute iFFT in time to go back to the original base
187 - (5) Update the solution by adding increment
189 Note that this is the only place where we compute the all-at-once residual because it requires communication and
190 swaps the solution values for the residuals. So after the residual tolerance is reached, one more ParaDiag
191 iteration will be done.
193 Args:
194 local_MS_running (list): list of currently running steps
195 """
197 for S in local_MS_running:
198 for hook in self.hooks:
199 hook.pre_sweep(step=S, level_number=0)
201 # `it_check` has already incremented the counter, so the first sweep is k = 0
202 k = max(local_MS_running[0].status.iter - 1, 0)
203 self.update_G_inv(k)
205 # communicate average residual for setting up Jacobians for non-linear problems
206 self.prepare_Jacobians(local_MS_running)
208 # compute the all-at-once residual to use as right hand side
209 self.compute_all_at_once_residual(local_MS_running)
211 # weighted FFT of the residual in time
212 self.FFT_in_time(quantity='residual', k=k)
214 # perform local solves of "collocation problems" on the steps (can be done in parallel)
215 for S in local_MS_running:
216 assert len(S.levels) == 1, 'Multi-level SDC not implemented in ParaDiag'
217 S.levels[0].sweep.update_nodes()
219 # inverse FFT of the increment in time
220 self.iFFT_in_time(quantity='increment', k=k)
222 # get the next iterate by adding increment to previous iterate
223 self.update_solution(local_MS_running)
225 for S in local_MS_running:
226 for hook in self.hooks:
227 hook.post_sweep(step=S, level_number=0)
229 # update stage
230 for S in local_MS_running:
231 S.status.stage = 'IT_CHECK'