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