Coverage for pySDC/core/convergence_controller.py: 84%
111 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1import logging
2from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING
3from pySDC.helpers.pysdc_helper import FrozenClass
5if TYPE_CHECKING:
6 from pySDC.core.controller import Controller
7 from pySDC.core.step import Step
10# short helper class to add params as attributes
11class Pars(FrozenClass):
12 """
13 Frozen parameters of a convergence controller as attributes, with defaults for `control_order` and `useMPI`.
14 """
16 def __init__(self, params: Dict[str, Any]) -> None:
17 """
18 Args:
19 params (dict): Parameters to store as attributes, overriding the defaults
20 """
21 self.control_order: int = 0 # integer that determines the order in which the convergence controllers are called
22 self.useMPI: Optional[bool] = None # depends on the controller
24 for k, v in params.items():
25 setattr(self, k, v)
27 self._freeze()
30# short helper class to store status variables
31class Status(FrozenClass):
32 """
33 Frozen container for named variables of a convergence controller.
34 Initialize status variables with None, since at the time of instantiation of the convergence controllers, not all
35 relevant information about the controller are known.
36 """
38 def __init__(self, status_variabes: List[str]) -> None:
39 """
40 Args:
41 status_variabes (list of str): Names of the status variables, all initialized to None
42 """
43 [setattr(self, key, None) for key in status_variabes]
45 self._freeze()
48class ConvergenceController(object):
49 """
50 Abstract base class for convergence controllers, which plug into the controller to steer iterations and step sizes.
51 """
53 def __init__(
54 self, controller: 'Controller', params: Dict[str, Any], description: Dict[str, Any], **kwargs: Any
55 ) -> None:
56 """
57 Initialization routine
59 Args:
60 controller (pySDC.Controller): The controller
61 params (dict): The params passed for this specific convergence controller
62 description (dict): The description object used to instantiate the controller
63 """
64 self.controller: 'Controller' = controller
65 self.params: Pars = Pars(self.setup(controller, params, description))
66 params_ok, msg = self.check_parameters(controller, params, description)
67 assert params_ok, f'{type(self).__name__} -- {msg}'
68 self.dependencies(controller, description)
69 self.logger: logging.Logger = logging.getLogger(f"{type(self).__name__}")
71 if self.params.useMPI:
72 self.prepare_MPI_datatypes()
74 def prepare_MPI_logical_operations(self) -> None:
75 """
76 Prepare MPI logical operations so we don't need to import mpi4py all the time
77 """
78 from mpi4py import MPI
80 self.MPI_LAND = MPI.LAND
81 self.MPI_LOR = MPI.LOR
83 def prepare_MPI_datatypes(self) -> None:
84 """
85 Prepare MPI datatypes so we don't need to import mpi4py all the time
86 """
87 from mpi4py import MPI
89 self.MPI_INT = MPI.INT
90 self.MPI_DOUBLE = MPI.DOUBLE
91 self.MPI_BOOL = MPI.BOOL
93 def log(self, msg: str, S: 'Step', level: int = 15, **kwargs: Any) -> None:
94 """
95 Shortcut that has a default level for the logger. 15 is above debug but below info.
97 Args:
98 msg (str): Message you want to log
99 S (pySDC.step): The current step
100 level (int): the level passed to the logger
102 Returns:
103 None
104 """
105 self.logger.log(level, f'Process {S.status.slot:2d} on time {S.time:.6f} - {msg}')
106 return None
108 def debug(self, msg: str, S: 'Step', **kwargs: Any) -> None:
109 """
110 Shortcut to pass messages at debug level to the logger.
112 Args:
113 msg (str): Message you want to log
114 S (pySDC.step): The current step
116 Returns:
117 None
118 """
119 self.log(msg=msg, S=S, level=10, **kwargs)
120 return None
122 def setup(
123 self, controller: 'Controller', params: Dict[str, Any], description: Dict[str, Any], **kwargs: Any
124 ) -> Dict[str, Any]:
125 """
126 Setup various variables that only need to be set once in the beginning.
127 If the convergence controller is added automatically, you can give it params by adding it manually.
128 It will be instantiated only once with the manually supplied parameters overriding automatically added
129 parameters.
131 This function scans the convergence controllers supplied to the description object for instances of itself.
132 This corresponds to the convergence controller being added manually by the user. If something is found, this
133 function will then return a composite dictionary from the `params` passed to this function as well as the
134 `params` passed manually, with priority to manually added parameters. If you added the convergence controller
135 manually, that is of course the same and nothing happens. If, on the other hand, the convergence controller was
136 added automatically, the `params` passed here will come from whatever added it and you can now override
137 parameters by adding the convergence controller manually.
138 This relies on children classes to return a composite dictionary from their defaults and from the result of this
139 function, so you should write
141 .. code-block:: python
143 return {**defaults, **super().setup(controller, params, description, **kwargs)}
145 when overloading this method in a child class, with `defaults` a dictionary containing default parameters.
147 Args:
148 controller (pySDC.Controller): The controller
149 params (dict): The params passed for this specific convergence controller
150 description (dict): The description object used to instantiate the controller
152 Returns:
153 (dict): The updated params dictionary after setup
154 """
155 # allow to change parameters by adding the convergence controller manually
156 return {**params, **description.get('convergence_controllers', {}).get(type(self), {})}
158 def dependencies(self, controller: 'Controller', description: Dict[str, Any], **kwargs: Any) -> None:
159 """
160 Load dependencies on other convergence controllers here.
162 Args:
163 controller (pySDC.Controller): The controller
164 description (dict): The description object used to instantiate the controller
166 Returns:
167 None
168 """
169 pass
171 def check_parameters(
172 self, controller: 'Controller', params: Dict[str, Any], description: Dict[str, Any], **kwargs: Any
173 ) -> Tuple[bool, str]:
174 """
175 Check whether parameters are compatible with whatever assumptions went into the step size functions etc.
177 Args:
178 controller (pySDC.Controller): The controller
179 params (dict): The params passed for this specific convergence controller
180 description (dict): The description object used to instantiate the controller
182 Returns:
183 bool: Whether the parameters are compatible
184 str: The error message
185 """
186 return True, ""
188 def check_iteration_status(self, controller: 'Controller', S: 'Step', **kwargs: Any) -> None:
189 """
190 Determine whether to keep iterating or not in this function.
192 Args:
193 controller (pySDC.Controller): The controller
194 S (pySDC.Step): The current step
196 Returns:
197 None
198 """
199 pass
201 def get_new_step_size(self, controller: 'Controller', S: 'Step', **kwargs: Any) -> None:
202 """
203 This function allows to set a step size with arbitrary criteria.
204 Make sure to give an order to the convergence controller by setting the `control_order` variable in the params.
205 This variable is an integer and you can see what the current order is by using
206 `controller.print_convergence_controllers()`.
208 Args:
209 controller (pySDC.Controller): The controller
210 S (pySDC.Step): The current step
212 Returns:
213 None
214 """
215 pass
217 def determine_restart(self, controller: 'Controller', S: 'Step', **kwargs: Any) -> None:
218 """
219 Determine for each step separately if it wants to be restarted for whatever reason.
221 Args:
222 controller (pySDC.Controller): The controller
223 S (pySDC.Step): The current step
225 Returns:
226 None
227 """
228 pass
230 def reset_status_variables(self, controller: 'Controller', **kwargs: Any) -> None:
231 """
232 Reset status variables.
233 This is called in the `restart_block` function.
234 Args:
235 controller (pySDC.Controller): The controller
237 Returns:
238 None
239 """
240 return None
242 def setup_status_variables(self, controller: 'Controller', **kwargs: Any) -> None:
243 """
244 Setup status variables.
245 This is not done at the time of instantiation, since the controller is not fully instantiated at that time and
246 hence not all information are available. Instead, this function is called after the controller has been fully
247 instantiated.
249 Args:
250 controller (pySDC.Controller): The controller
252 Returns:
253 None
254 """
255 return None
257 def reset_buffers_nonMPI(self, controller: 'Controller', **kwargs: Any) -> None:
258 """
259 Buffers refer to variables used across multiple steps that are stored in the convergence controller classes to
260 imitate communication in non MPI versions. These have to be reset in order to replicate availability of
261 variables in MPI versions.
263 For instance, if step 0 sets self.buffers.x = 1 from self.buffers.x = 0, when the same MPI rank uses the
264 variable with step 1, it will still carry the value of self.buffers.x = 1, equivalent to a send from the rank
265 computing step 0 to the rank computing step 1.
267 However, you can only receive what somebody sent and in order to make sure that is true for the non MPI
268 versions, we reset after each iteration so you cannot use this function to communicate backwards from the last
269 step to the first one for instance.
271 This function is called both at the end of instantiating the controller, as well as after each iteration.
273 Args:
274 controller (pySDC.Controller): The controller
276 Returns:
277 None
278 """
279 pass
281 def pre_iteration_processing(self, controller: 'Controller', S: 'Step', **kwargs: Any) -> None:
282 """
283 Do whatever you want to before each iteration here.
285 Args:
286 controller (pySDC.Controller): The controller
287 S (pySDC.Step): The current step
289 Returns:
290 None
291 """
292 pass
294 def post_iteration_processing(self, controller: 'Controller', S: 'Step', **kwargs: Any) -> None:
295 """
296 Do whatever you want to after each iteration here.
298 Args:
299 controller (pySDC.Controller): The controller
300 S (pySDC.Step): The current step
302 Returns:
303 None
304 """
305 pass
307 def post_iteration_processing_block(self, controller: 'Controller', **kwargs: Any) -> None:
308 """
309 Do whatever you want to after each iteration, once for the whole block rather than once per step.
311 Use this instead of `post_iteration_processing` for state that belongs to the block rather than
312 to a step. A convergence controller is instantiated once per *controller*, so anything stored on
313 `self` spans a whole block in `controller_nonMPI` but a single rank under MPI, while the per-step
314 hooks fire once per step. A recursion driven from a per-step hook therefore advances a different
315 number of times depending only on how the run was decomposed. State that genuinely belongs to a
316 step should go on `S.status` via `add_status_variable_to_step` instead.
318 Note this fires once per iteration **per rank**, not once per block: under MPI every rank calls
319 it, so a quantity that has to agree across the block still needs a reduction in here. See
320 `AdaptiveAlpha` for how that looks.
322 Note also what that does *not* buy you. Under a pipelined method an early step converges and
323 drops out, so ranks take part in different numbers of iterations by design; a recursion driven
324 from here therefore still advances a different number of times on different ranks, and that is
325 the pipelining rather than a defect. Where every step iterates together -- ParaDiag, which
326 cannot converge one step at a time because the transform needs them all -- this is exact.
328 Args:
329 controller (pySDC.Controller): The controller
331 Returns:
332 None
333 """
334 pass
336 def post_step_processing(self, controller: 'Controller', S: 'Step', **kwargs: Any) -> None:
337 """
338 Do whatever you want to after each step here.
340 Args:
341 controller (pySDC.Controller): The controller
342 S (pySDC.Step): The current step
344 Returns:
345 None
346 """
347 pass
349 def post_run_processing(self, controller: 'Controller', S: 'Step', **kwargs: Any) -> None:
350 """
351 Do whatever you want to after the run here.
353 Args:
354 controller (pySDC.Controller): The controller
355 S (pySDC.Step): The current step
357 Returns:
358 None
359 """
360 pass
362 def prepare_next_block(
363 self, controller: 'Controller', S: 'Step', size: int, time: Any, Tend: float, **kwargs: Any
364 ) -> None:
365 """
366 Prepare stuff like spreading step sizes or whatever.
368 Args:
369 controller (pySDC.Controller): The controller
370 S (pySDC.Step): The current step
371 size (int): The number of ranks
372 time (float): The current time will be list in nonMPI controller implementation
373 Tend (float): The final time
375 Returns:
376 None
377 """
378 pass
380 def convergence_control(self, controller: 'Controller', S: 'Step', **kwargs: Any) -> None:
381 """
382 Call all the functions related to convergence control.
383 This is called in `it_check` in the controller after every iteration just after `post_iteration_processing`.
384 Args:
385 controller (pySDC.Controller): The controller
386 S (pySDC.Step): The current step
388 Returns:
389 None
390 """
392 self.get_new_step_size(controller, S, **kwargs)
393 self.determine_restart(controller, S, **kwargs)
394 self.check_iteration_status(controller, S, **kwargs)
396 return None
398 def communicate_convergence(self, controller: 'Controller', S: 'Step', **kwargs: Any) -> None:
399 """
400 Share convergence status across the block, however the block happens to be spread out.
402 Called once per step after every step has decided for itself, which is the only point at
403 which a reduction over the block is meaningful. `CheckConvergence` is the one that does
404 something here.
406 Args:
407 controller (pySDC.Controller): The controller
408 S (pySDC.Step): The current step
410 Returns:
411 None
412 """
413 pass
415 def post_spread_processing(self, controller: 'Controller', S: 'Step', **kwargs: Any) -> None:
416 """
417 This function is called at the end of the `SPREAD` stage in the controller
419 Args:
420 controller (pySDC.Controller): The controller
421 S (pySDC.Step): The current step
422 """
423 pass
425 def send(self, comm, dest, data, blocking=False, **kwargs):
426 """
427 Send data to a different rank
429 Args:
430 comm (mpi4py.MPI.Intracomm): Communicator
431 dest (int): The target rank
432 data: Data to be sent
433 blocking (bool): Whether the communication is blocking or not
435 Returns:
436 request handle of the communication
437 """
438 kwargs['tag'] = kwargs.get('tag', abs(self.params.control_order))
440 # log what's happening for debug purposes
441 self.logger.debug(f'Step {comm.rank} {"" if blocking else "i"}sends to step {dest} with tag {kwargs["tag"]}')
443 if blocking:
444 req = comm.send(data, dest=dest, **kwargs)
445 else:
446 req = comm.isend(data, dest=dest, **kwargs)
448 return req
450 def recv(self, comm, source, **kwargs):
451 """
452 Receive some data
454 Args:
455 comm (mpi4py.MPI.Intracomm): Communicator
456 source (int): Where to look for receiving
458 Returns:
459 whatever has been received
460 """
461 kwargs['tag'] = kwargs.get('tag', abs(self.params.control_order))
463 # log what's happening for debug purposes
464 self.logger.debug(f'Step {comm.rank} receives from step {source} with tag {kwargs["tag"]}')
466 data = comm.recv(source=source, **kwargs)
468 return data
470 def Send(self, comm, dest, buffer, blocking=False, **kwargs):
471 """
472 Send data to a different rank
474 Args:
475 comm (mpi4py.MPI.Intracomm): Communicator
476 dest (int): The target rank
477 buffer: Buffer for the data
478 blocking (bool): Whether the communication is blocking or not
480 Returns:
481 request handle of the communication
482 """
483 kwargs['tag'] = kwargs.get('tag', abs(self.params.control_order))
485 # log what's happening for debug purposes
486 self.logger.debug(f'Step {comm.rank} {"" if blocking else "i"}Sends to step {dest} with tag {kwargs["tag"]}')
488 if blocking:
489 req = comm.Send(buffer, dest=dest, **kwargs)
490 else:
491 req = comm.Isend(buffer, dest=dest, **kwargs)
493 return req
495 def Recv(self, comm, source, buffer, **kwargs):
496 """
497 Receive some data
499 Args:
500 comm (mpi4py.MPI.Intracomm): Communicator
501 source (int): Where to look for receiving
502 buffer: Buffer to receive the data into
504 Returns:
505 None, the data is received into `buffer`
506 """
507 kwargs['tag'] = kwargs.get('tag', abs(self.params.control_order))
509 # log what's happening for debug purposes
510 self.logger.debug(f'Step {comm.rank} Receives from step {source} with tag {kwargs["tag"]}')
512 data = comm.Recv(buffer, source=source, **kwargs)
514 return data
516 def add_status_variable_to_step(self, key, value=None):
517 """
518 Add a status variable to the steps of the controller and optionally set its value on all of them.
520 The variable is registered on the class of the step status, see `FrozenClass.add_attr`, so it exists on every
521 step.
523 Args:
524 key (str): Name of the status variable
525 value: Value to set on all steps, which is skipped if None
527 Returns:
528 None
529 """
530 steps = self.controller.steps
532 steps[0].status.add_attr(key)
534 if value is not None:
535 self.set_step_status_variable(key, value)
537 def set_step_status_variable(self, key, value):
538 """
539 Set a status variable on all steps of the controller.
541 Args:
542 key (str): Name of the status variable, which must have been added already
543 value: Value to set
545 Returns:
546 None
547 """
548 steps = self.controller.steps
550 for S in steps:
551 S.status.__dict__[key] = value
553 def add_status_variable_to_level(self, key, value=None):
554 """
555 Add a status variable to the levels of the controller and optionally set its value on all of them.
557 The variable is registered on the class of the level status, see `FrozenClass.add_attr`, so it exists on every
558 level.
560 Args:
561 key (str): Name of the status variable
562 value: Value to set on all levels of all steps, which is skipped if None
564 Returns:
565 None
566 """
567 steps = self.controller.steps
569 steps[0].levels[0].status.add_attr(key)
571 if value is not None:
572 self.set_level_status_variable(key, value)
574 def set_level_status_variable(self, key, value):
575 """
576 Set a status variable on all levels of all steps of the controller.
578 Args:
579 key (str): Name of the status variable, which must have been added already
580 value: Value to set
582 Returns:
583 None
584 """
585 steps = self.controller.steps
587 for S in steps:
588 for L in S.levels:
589 L.status.__dict__[key] = value