Coverage for pySDC/implementations/controller_classes/ParaDiag.py: 94%
53 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-15 06:23 +0000
1import logging
2from typing import Any, Dict
4import numpy as np
7class ParaDiag:
8 """
9 What ParaDiag is, independently of how the block is spread over processes.
11 Mixed into a controller, this replaces its iteration and leaves everything around it alone:
13 class controller_ParaDiag_nonMPI(ParaDiag, controller_nonMPI)
14 class controller_ParaDiag_MPI(ParaDiag, controller_MPI)
16 It deliberately has no `__init__` and no controller of its own to inherit from, so that it can
17 be mixed into either transport's controller without the two having to agree on a constructor
18 signature. A concrete ParaDiag controller calls `prepare_ParaDiag_params` on the two
19 dictionaries and then hands them to its controller's initialisation.
21 What stays with the concrete classes is everything whose implementation depends on where the
22 other steps are: `apply_matrix`, `prepare_Jacobians`, `compute_all_at_once_residual`,
23 `update_G_inv` and the block driver.
24 """
26 @staticmethod
27 def prepare_ParaDiag_params(controller_params: Dict[str, Any], description: Dict[str, Any]) -> None:
28 """
29 Check and complete the parameters ParaDiag needs, in place.
31 Call this *before* the controller's own initialisation: it only reads and writes the two
32 dictionaries, and must have run by the time the steps are built.
34 Args:
35 controller_params (dict): parameter set for the controller and the steps
36 description (dict): all the parameters to set up the rest (levels, problems, ...)
37 """
38 from pySDC.implementations.sweeper_classes.ParaDiagSweepers import QDiagonalization
40 if QDiagonalization in description['sweeper_class'].__mro__:
41 description['sweeper_params']['ignore_ic'] = True
42 description['sweeper_params']['update_f_evals'] = False
43 else:
44 logging.getLogger('controller').warning(
45 f'Warning: Your sweeper class {description["sweeper_class"]} is not derived from {QDiagonalization}. You probably want to use another sweeper class.'
46 )
48 if not controller_params.get('all_to_done', True):
49 raise NotImplementedError('ParaDiag only implemented with option `all_to_done=True`')
50 if 'alpha' not in controller_params.keys():
51 from pySDC.core.errors import ParameterError
53 raise ParameterError('Please supply alpha as a parameter to the ParaDiag controller!')
54 controller_params['average_jacobian'] = controller_params.get('average_jacobian', True)
56 controller_params['all_to_done'] = True
58 # ------------------------------------------------------------------ the iteration
60 def get_stages(self) -> Dict[str, Any]:
61 """
62 ParaDiag has one iteration stage, and no predictor because it has no coarse level.
64 Returns:
65 dict: stage name -> the method that runs it
66 """
67 return {
68 'SPREAD': self.spread,
69 'IT_CHECK': self.it_check,
70 'IT_PARADIAG': self.it_ParaDiag,
71 }
73 def next_iteration_stage(self, S: Any) -> str:
74 """
75 Args:
76 S (pySDC.Step.step): The current step
78 Returns:
79 str: name of the stage to enter
80 """
81 return 'IT_PARADIAG'
83 def compute_residual_after_spread(self, S: Any) -> None:
84 """
85 ParaDiag's residual is the one of the composite collocation problem, which `it_ParaDiag`
86 computes as part of the iteration. The convergence check runs before the first iteration,
87 so the initial guess needs its residual here.
89 Args:
90 S (pySDC.Step.step): The current step
91 """
92 S.levels[0].sweep.compute_residual()
94 def step_is_active(self, time: float, block_start: float, Tend: float) -> bool:
95 """
96 ParaDiag diagonalizes across the whole block, so it cannot drop a step out of one. A block
97 that starts before `Tend` is run whole, past `Tend` if need be.
99 Args:
100 time (float): when this step starts
101 block_start (float): when the first step of this step's block starts
102 Tend (float): ending time
104 Returns:
105 bool: whether this step takes part
106 """
107 active = block_start < Tend - 10 * np.finfo(float).eps
109 if active and time >= Tend - 10 * np.finfo(float).eps:
110 self.logger.warning(
111 'Warning: This controller will solve past your desired end time until the end of its block!'
112 )
114 return active
116 def prepare_convergence_check(self, *args: Any, **kwargs: Any) -> None:
117 """
118 The residual is already current -- `it_ParaDiag` computed it, and recomputing it the way a
119 sweep-based algorithm does would need initial conditions that have not been communicated
120 yet. The end point is not, because nothing sent it anywhere, so compute it here: `it_check`
121 is what publishes `uend` when a step is done.
123 Takes whatever its controller passes -- a block of steps or a communicator -- and needs
124 none of it, because the steps to do this for are the ones this controller owns.
125 """
126 for S in self.steps:
127 S.levels[0].sweep.compute_end_point()
129 # ------------------------------------------------------------------ alpha and the transform
131 @staticmethod
132 def resolve_alpha(alpha: Any, k: int = 0) -> float:
133 """
134 Read the alpha for iteration `k` out of whatever the user supplied.
136 `alpha` may be a single number, a sequence indexed by iteration (the last entry is reused once
137 it runs out), or a callable taking the iteration index. Making it iteration dependent lets the
138 outer iteration start with a well-conditioned alpha and tighten it later.
140 Static because the steps need an alpha before the controller has parameters to read it from.
142 Args:
143 alpha: the alpha parameter as supplied by the user
144 k (int): iteration index
146 Returns:
147 float: alpha to use for this iteration
148 """
149 if callable(alpha):
150 return float(alpha(k))
151 if hasattr(alpha, '__len__'):
152 return float(alpha[min(k, len(alpha) - 1)])
153 return float(alpha)
155 def get_alpha(self, k: int = 0) -> float:
156 """
157 Get the ParaDiag alpha parameter for iteration `k`.
159 Args:
160 k (int): iteration index
162 Returns:
163 float: alpha to use for this iteration
164 """
165 return self.resolve_alpha(self.params.alpha, k)
167 def get_FFT_matrices(self, k: int = 0) -> Any:
168 """
169 Get the weighted FFT and iFFT matrices for iteration `k`, rebuilding them only when alpha
170 actually changes.
172 Args:
173 k (int): iteration index
175 Returns:
176 tuple: the forward and backward weighted FFT matrices
177 """
178 alpha = self.get_alpha(k)
179 if getattr(self, '_cached_alpha', None) != alpha:
180 from pySDC.helpers.ParaDiagHelper import get_weighted_FFT_matrix, get_weighted_iFFT_matrix
182 self._FFT_matrix = get_weighted_FFT_matrix(self.n_steps, alpha)
183 self._iFFT_matrix = get_weighted_iFFT_matrix(self.n_steps, alpha)
184 self._cached_alpha = alpha
185 return self._FFT_matrix, self._iFFT_matrix
187 def FFT_in_time(self, quantity: Any, k: int = 0) -> None:
188 """
189 Compute weighted forward FFT in time. The weighting is determined by the alpha parameter in ParaDiag
191 Note: The implementation via matrix-vector multiplication may be inefficient and less stable compared to an FFT
192 with transposes!
194 Args:
195 quantity (str): the level attribute to transform
196 k (int): iteration index, for an iteration dependent alpha
197 """
198 self.apply_matrix(self.get_FFT_matrices(k)[0], quantity)
200 def iFFT_in_time(self, quantity: Any, k: int = 0) -> None:
201 """
202 Compute weighted backward FFT in time. The weighting is determined by the alpha parameter in ParaDiag
204 Args:
205 quantity (str): the level attribute to transform
206 k (int): iteration index, for an iteration dependent alpha
207 """
208 self.apply_matrix(self.get_FFT_matrices(k)[1], quantity)
210 # ------------------------------------------------------------------ what the transport supplies
212 def apply_matrix(self, mat: Any, quantity: str) -> None:
213 """
214 Apply a square matrix across the steps, in place.
216 How this is done depends entirely on where the other steps are, so the concrete controllers
217 implement it.
219 Args:
220 mat: square matrix with as many rows as there are steps
221 quantity (str): 'residual' or 'increment', the level attribute to transform
222 """
223 raise NotImplementedError('ParaDiag controllers have to implement apply_matrix')
225 def update_G_inv(self, k: int = 0) -> None:
226 """
227 Rebuild G^-1 on the local step(s) when alpha changes with the iteration.
229 G^-1 depends on alpha, so an iteration dependent alpha means the sweeper's diagonalization has
230 to be recomputed. Subclasses implement this because only they know which steps they own.
232 Args:
233 k (int): iteration index
234 """
235 raise NotImplementedError('ParaDiag controllers have to implement update_G_inv')