Coverage for pySDC/implementations/controller_classes/controller_nonMPI.py: 99%
296 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1import itertools
2import copy as cp
3import numpy as np
4import dill
6from pySDC.core.controller import Controller
7from pySDC.core.step import Step
8from pySDC.core.errors import ControllerError, CommunicationError
9from pySDC.implementations.convergence_controller_classes.basic_restarting import BasicRestarting
12class controller_nonMPI(Controller):
13 """
15 Controller running SDC, MLSDC and PFASST with the time steps of a block emulated serially in one process.
17 PFASST controller, running serialized version of PFASST in blocks (MG-style)
19 """
21 def __init__(self, num_procs, controller_params, description):
22 """
23 Initialization routine for PFASST controller
25 Args:
26 num_procs: number of parallel time steps (still serial, though), can be 1
27 controller_params: parameter set for the controller and the steps
28 description: all the parameters to set up the rest (levels, problems, transfer, ...)
29 """
31 if 'predict' in controller_params:
32 raise ControllerError('predict flag is ignored, use predict_type instead')
34 # call parent's initialization routine
35 super().__init__(controller_params, description, useMPI=False)
37 self.MS: list[Step] = [Step(description)]
39 # try to initialize via dill.copy (much faster for many time-steps)
40 try:
41 for _ in range(num_procs - 1):
42 self.MS.append(dill.copy(self.MS[0]))
43 # if this fails (e.g. due to un-picklable data in the steps), initialize separately
44 except (dill.PicklingError, TypeError, ValueError) as error:
45 self.logger.warning(f'Need to initialize steps separately due to pickling error: {error}')
46 for _ in range(num_procs - 1):
47 self.MS.append(Step(description))
49 self.base_convergence_controllers += [BasicRestarting.get_implementation(useMPI=False)]
50 for convergence_controller in self.base_convergence_controllers:
51 self.add_convergence_controller(convergence_controller, description)
53 if self.params.dump_setup:
54 self.dump_setup(step=self.MS[0], controller_params=controller_params, description=description)
56 if num_procs > 1 and len(self.MS[0].levels) > 1:
57 for S in self.MS:
58 for L in S.levels:
59 if not L.sweep.coll.right_is_node:
60 raise ControllerError("For PFASST to work, we assume uend^k = u_M^k")
62 if all(len(S.levels) == len(self.MS[0].levels) for S in self.MS):
63 self.nlevels = len(self.MS[0].levels)
64 else:
65 raise ControllerError('all steps need to have the same number of levels')
67 if self.nlevels == 0:
68 raise ControllerError('need at least one level')
70 # The stages read `nsweeps` off a step they own, the same way `controller_MPI` reads it off
71 # the one step a rank has, so this only has to establish that any step will do.
72 for nl in range(self.nlevels):
73 if not all(S.levels[nl].params.nsweeps == self.MS[0].levels[nl].params.nsweeps for S in self.MS):
74 raise ControllerError('all steps need to agree on the number of sweeps per level')
76 # `it_coarse` sweeps the coarsest level exactly once. Single-level Gauss-like MSSDC routes
77 # through it too, so reject multiple sweeps there as well instead of silently ignoring them.
78 # `mssdc_jac` only decides the routing when there is more than one step: a single step is
79 # plain SDC and always goes through `it_fine`, which honours nsweeps.
80 if self.MS[0].levels[-1].params.nsweeps > 1 and (
81 self.nlevels > 1 or (num_procs > 1 and not self.params.mssdc_jac)
82 ):
83 raise ControllerError('this controller cannot do multiple sweeps on coarsest level')
85 self.check_variable_coefficients(num_procs)
87 if self.nlevels == 1 and self.params.predict_type is not None:
88 self.logger.warning(
89 'you have specified a predictor type but only a single level.. predictor will be ignored'
90 )
92 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
93 C.reset_buffers_nonMPI(self)
94 C.setup_status_variables(self, MS=self.MS)
96 self.stages = self.get_stages()
98 def run(self, u0, t0, Tend):
99 """
100 Main driver for running the serial version of SDC, MSSDC, MLSDC and PFASST (virtual parallelism)
102 Args:
103 u0: initial values
104 t0: starting time
105 Tend: ending time
107 Returns:
108 end values on the finest level
109 stats object containing statistics for each step, each level and each iteration
110 """
112 # some initializations and reset of statistics
113 uend = None
114 num_procs = len(self.MS)
115 for hook in self.hooks:
116 hook.reset_stats()
118 # start from the step size in the description, not the one adaptivity left from a previous run
119 for S in self.MS:
120 for L in S.levels:
121 L.params.dt = L.params.dt_initial
123 # initial ordering of the steps: 0,1,...,Np-1
124 slots = list(range(num_procs))
126 # initialize time variables of each step
127 time = [t0 + sum(self.MS[j].dt for j in range(p)) for p in slots]
129 # determine which steps are still active (time < Tend)
130 active = [self.step_is_active(time[p], block_start=time[0], t0=t0, Tend=Tend, dt=self.MS[p].dt) for p in slots]
132 if not any(active):
133 raise ControllerError('Nothing to do, check t0, dt and Tend.')
135 # compress slots according to active steps, i.e. remove all steps which have times above Tend
136 active_slots = list(itertools.compress(slots, active))
138 # initialize block of steps with u0
139 self.restart_block(active_slots, time, u0)
141 for hook in self.hooks:
142 hook.post_setup(step=None, level_number=None)
144 # call pre-run hook
145 for S in self.MS:
146 for hook in self.hooks:
147 hook.pre_run(step=S, level_number=0)
149 # main loop: as long as at least one step is still active (time < Tend), do something
150 while any(active):
151 MS_active = [self.MS[p] for p in active_slots]
152 done = False
153 while not done:
154 done = self.pfasst(MS_active)
156 restarts = [S.status.restart for S in MS_active]
157 restart_at = np.where(restarts)[0][0] if True in restarts else len(MS_active)
158 if True in restarts: # restart part of the block
159 # initial condition to next block is initial condition of step that needs restarting
160 uend = self.MS[restart_at].levels[0].u[0]
161 time[active_slots[0]] = time[restart_at]
162 self.logger.info(f'Starting next block with initial conditions from step {restart_at}')
164 else: # move on to next block
165 # initial condition for next block is last solution of current block
166 uend = self.MS[active_slots[-1]].levels[0].uend
167 time[active_slots[0]] = time[active_slots[-1]] + self.MS[active_slots[-1]].dt
169 for S in MS_active[:restart_at]:
170 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
171 C.post_step_processing(self, S, MS=MS_active)
173 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
174 [C.prepare_next_block(self, S, len(active_slots), time, Tend, MS=MS_active) for S in self.MS]
176 # setup the times of the steps for the next block
177 for i in range(1, len(active_slots)):
178 time[active_slots[i]] = time[active_slots[i] - 1] + self.MS[active_slots[i] - 1].dt
180 # determine new set of active steps and compress slots accordingly
181 active = [
182 self.step_is_active(time[p], block_start=time[0], t0=t0, Tend=Tend, dt=self.MS[p].dt) for p in slots
183 ]
184 active_slots = list(itertools.compress(slots, active))
186 # restart active steps (reset all values and pass uend to u0)
187 self.restart_block(active_slots, time, uend)
189 # call post-run hook
190 for S in self.MS:
191 for hook in self.hooks:
192 hook.post_run(step=S, level_number=0)
194 for S in self.MS:
195 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
196 C.post_run_processing(self, S, MS=MS_active)
198 return uend, self.return_stats()
200 def restart_block(self, active_slots, time, u0):
201 """
202 Helper routine to reset/restart block of (active) steps
204 Args:
205 active_slots: list of active steps
206 time: list of new times
207 u0: initial value to distribute across the steps
209 """
211 # loop over active slots (not directly, since we need the previous entry as well)
212 for j in range(len(active_slots)):
213 # get slot number
214 p = active_slots[j]
216 # store current slot number for diagnostics
217 self.MS[p].status.slot = p
218 # store link to previous step
219 self.MS[p].prev = self.MS[active_slots[j - 1]]
220 # resets step
221 self.MS[p].reset_step()
222 # determine whether I am the first and/or last in line
223 self.MS[p].status.first = active_slots.index(p) == 0
224 self.MS[p].status.last = active_slots.index(p) == len(active_slots) - 1
225 # initialize step with u0
226 self.MS[p].init_step(u0)
227 # reset some values
228 self.MS[p].status.done = False
229 self.MS[p].status.prev_done = False
230 self.MS[p].status.iter = 0
231 self.MS[p].status.stage = 'SPREAD'
232 self.MS[p].status.force_done = False
233 self.MS[p].status.time_size = len(active_slots)
235 for l in self.MS[p].levels:
236 l.tag = None
237 l.status.sweep = 1
239 for p in active_slots:
240 for lvl in self.MS[p].levels:
241 lvl.status.time = time[p]
243 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
244 C.reset_status_variables(self, active_slots=active_slots)
246 def send_full(self, S, level=None, add_to_stats=False):
247 """
248 Function to perform the send, including bookkeeping and logging
250 Args:
251 S: the current step
252 level: the level number
253 add_to_stats: a flag to end recording data in the hooks (defaults to False)
255 Note:
256 Computing the end point is this function's job, not the caller's, exactly as in
257 `controller_MPI.send_full`. It happens whether or not anyone is listening, because the
258 last step has no successor but its `uend` is still the block's result.
259 """
260 for hook in self.hooks:
261 hook.pre_comm(step=S, level_number=level)
263 # sending here means computing uend ("one-sided communication")
264 S.levels[level].sweep.compute_end_point()
266 if not S.status.last:
267 self.logger.debug(
268 'Process %2i provides data on level %2i with tag %s' % (S.status.slot, level, S.status.iter)
269 )
270 S.levels[level].tag = cp.deepcopy((level, S.status.iter, S.status.slot))
272 for hook in self.hooks:
273 hook.post_comm(step=S, level_number=level, add_to_stats=add_to_stats)
275 def recv_full(self, S, level=None, add_to_stats=False):
276 """
277 Function to perform the recv, including bookkeeping and logging
279 Args:
280 S: the current step
281 level: the level number
282 add_to_stats: a flag to end recording data in the hooks (defaults to False)
283 """
285 def recv(target, source, tag=None):
286 """
287 Receive function
289 Args:
290 target: level which will receive the values
291 source: level which initiated the send
292 tag: identifier to check if this message is really for me
293 """
295 if tag is not None and source.tag != tag:
296 raise CommunicationError('source and target tag are not the same, got %s and %s' % (source.tag, tag))
297 # simply do a deepcopy of the values uend to become the new u0 at the target
298 target.u[0] = target.prob.dtype_u(source.uend)
299 # re-evaluate f on left interval boundary
300 target.f[0] = target.prob.eval_f(target.u[0], target.time)
302 for hook in self.hooks:
303 hook.pre_comm(step=S, level_number=level)
304 if not S.status.prev_done and not S.status.first:
305 self.logger.debug(
306 'Process %2i receives from %2i on level %2i with tag %s'
307 % (S.status.slot, S.prev.status.slot, level, S.status.iter)
308 )
309 recv(S.levels[level], S.prev.levels[level], tag=(level, S.status.iter, S.prev.status.slot))
310 for hook in self.hooks:
311 hook.post_comm(step=S, level_number=level, add_to_stats=add_to_stats)
313 def pfasst(self, local_MS_active):
314 """
315 Main function including the stages of SDC, MLSDC and PFASST (the "controller")
317 For the workflow of this controller, check out one of our PFASST talks or the pySDC paper
319 This method changes self.MS directly by accessing active steps through local_MS_active. Nothing is returned.
321 Args:
322 local_MS_active (list): all active steps
323 """
325 # if all stages are the same (or DONE), continue, otherwise abort
326 stages = [S.status.stage for S in local_MS_active if S.status.stage != 'DONE']
327 if stages[1:] == stages[:-1]:
328 stage = stages[0]
329 else:
330 raise ControllerError('not all stages are equal')
332 self.logger.debug(stage)
334 MS_running = [S for S in local_MS_active if S.status.stage != 'DONE']
336 self.stages.get(stage, self.default)(MS_running)
338 return all(S.status.done for S in local_MS_active)
340 def get_stages(self):
341 """
342 The stages this controller can be in, and what to run in each.
344 A subclass that iterates differently replaces the iteration stages here and says which one
345 to enter in `next_iteration_stage`; everything around the iteration is the same for any
346 algorithm this controller runs.
348 Returns:
349 dict: stage name -> the method that runs it
350 """
351 return {
352 'SPREAD': self.spread,
353 'PREDICT': self.predict,
354 'IT_CHECK': self.it_check,
355 'IT_FINE': self.it_fine,
356 'IT_DOWN': self.it_down,
357 'IT_COARSE': self.it_coarse,
358 'IT_UP': self.it_up,
359 }
361 def next_iteration_stage(self, S):
362 """
363 The stage that starts one iteration of the algorithm.
365 Args:
366 S (pySDC.Step.step): The current step
368 Returns:
369 str: name of the stage to enter
370 """
371 if len(S.levels) > 1: # MLSDC or PFASST
372 return 'IT_DOWN'
373 elif S.status.time_size == 1 or self.params.mssdc_jac: # SDC or parallel MSSDC (Jacobi-like)
374 return 'IT_FINE'
375 else:
376 return 'IT_COARSE' # serial MSSDC (Gauss-like)
378 def spread(self, local_MS_running):
379 """
380 Spreading phase
382 Args:
383 local_MS_running (list): list of currently running steps
384 """
386 for S in local_MS_running:
387 # first stage: spread values
388 for hook in self.hooks:
389 hook.pre_step(step=S, level_number=0)
391 # call predictor from sweeper
392 S.levels[0].sweep.predict()
394 self.compute_residual_after_spread(S)
396 # update stage
397 if len(S.levels) > 1: # MLSDC or PFASST with predict
398 S.status.stage = 'PREDICT'
399 else:
400 S.status.stage = 'IT_CHECK'
402 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
403 C.post_spread_processing(self, S, MS=local_MS_running)
405 def predict(self, local_MS_running):
406 """
407 Predictor phase
409 Args:
410 local_MS_running (list): list of currently running steps
411 """
413 for S in local_MS_running:
414 for hook in self.hooks:
415 hook.pre_predict(step=S, level_number=0)
417 if self.params.predict_type is None:
418 pass
420 elif self.params.predict_type == 'fine_only':
421 # do a fine sweep only
422 for S in local_MS_running:
423 S.levels[0].sweep.update_nodes()
425 elif self.params.predict_type == 'pfasst_burnin':
426 # loop over all steps
427 for S in local_MS_running:
428 # restrict to coarsest level
429 for l in range(1, len(S.levels)):
430 S.transfer(source=S.levels[l - 1], target=S.levels[l])
432 # loop over all steps
433 for q in range(len(local_MS_running)):
434 # loop over last steps: [1,2,3,4], [2,3,4], [3,4], [4]
435 for p in range(q, len(local_MS_running)):
436 S = local_MS_running[p]
438 # do the sweep with new values
439 S.levels[-1].sweep.update_nodes()
441 # send updated values on coarsest level
442 self.send_full(S, level=len(S.levels) - 1)
444 # loop over last steps: [2,3,4], [3,4], [4]
445 for p in range(q + 1, len(local_MS_running)):
446 S = local_MS_running[p]
447 # receive values sent during previous sweep
448 self.recv_full(S, level=len(S.levels) - 1, add_to_stats=(p == len(local_MS_running) - 1))
450 # loop over all steps
451 for S in local_MS_running:
452 # interpolate back to finest level
453 for l in range(len(S.levels) - 1, 0, -1):
454 S.transfer(source=S.levels[l], target=S.levels[l - 1])
456 # send updated values forward
457 self.send_full(S, level=0)
458 # receive values
459 self.recv_full(S, level=0)
461 # end this with a fine sweep
462 for S in local_MS_running:
463 S.levels[0].sweep.update_nodes()
465 elif self.params.predict_type == 'fmg':
466 # TODO: implement FMG predictor
467 raise NotImplementedError('FMG predictor is not yet implemented')
469 else:
470 raise ControllerError('Wrong predictor type, got %s' % self.params.predict_type)
472 for S in local_MS_running:
473 for hook in self.hooks:
474 hook.post_predict(step=S, level_number=0)
476 for S in local_MS_running:
477 # update stage
478 S.status.stage = 'IT_CHECK'
480 def it_check(self, local_MS_running):
481 """
482 Key routine to check for convergence/termination
484 Args:
485 local_MS_running (list): list of currently running steps
486 """
488 self.prepare_convergence_check(local_MS_running)
490 for S in local_MS_running:
491 if S.status.iter > 0:
492 for hook in self.hooks:
493 hook.post_iteration(step=S, level_number=0)
495 # decide if the step is done, needs to be restarted and other things convergence related
496 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
497 C.post_iteration_processing(self, S, MS=local_MS_running)
498 C.convergence_control(self, S, MS=local_MS_running)
500 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
501 C.post_iteration_processing_block(self, MS=local_MS_running)
503 for S in local_MS_running:
504 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
505 C.communicate_convergence(self, S, MS=local_MS_running)
507 if not S.status.done:
508 # increment iteration count here (and only here)
509 S.status.iter += 1
510 for hook in self.hooks:
511 hook.pre_iteration(step=S, level_number=0)
512 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
513 C.pre_iteration_processing(self, S, MS=local_MS_running)
515 S.status.stage = self.next_iteration_stage(S)
516 else:
517 for hook in self.hooks:
518 hook.post_step(step=S, level_number=0)
519 S.status.stage = 'DONE'
521 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
522 C.reset_buffers_nonMPI(self)
524 def compute_residual_after_spread(self, S):
525 """
526 Make the residual current right after the initial guess, for algorithms that need it there.
528 This controller computes the residual at the top of `it_check` instead, so there is nothing
529 to do; an algorithm whose residual is not available at that point overrides this. Called
530 before `post_spread_processing`, because convergence controllers there may still change the
531 initial iterate.
533 Args:
534 S (pySDC.Step.step): The current step
535 """
536 pass
538 def prepare_convergence_check(self, local_MS_running):
539 """
540 Make current what `it_check` is about to read: the end point and the residual.
542 For a sweep-based algorithm both come out of the exchange with the neighbours, so this is
543 the send and receive plus the residual it enables. An algorithm that gets its residual from
544 the iteration instead overrides this, and still owes the end point.
546 Args:
547 local_MS_running (list): list of currently running steps
548 """
549 for S in local_MS_running:
550 # send updated values forward
551 self.send_full(S, level=0)
552 # receive values
553 self.recv_full(S, level=0)
554 # compute current residual
555 S.levels[0].sweep.compute_residual(stage='IT_CHECK')
557 def it_fine(self, local_MS_running: list[Step]):
558 """
559 Fine sweeps
561 Args:
562 local_MS_running (list): list of currently running steps
563 """
565 for S in local_MS_running:
566 S.levels[0].status.sweep = 0
568 nsweeps = local_MS_running[0].levels[0].params.nsweeps
570 for k in range(nsweeps):
571 for S in local_MS_running:
572 S.levels[0].status.sweep += 1
574 for S in local_MS_running:
575 # send updated values forward
576 self.send_full(S, level=0)
577 # receive values
578 self.recv_full(S, level=0, add_to_stats=(k == nsweeps - 1))
580 for S in local_MS_running:
581 # standard sweep workflow: update nodes, log progress (the residual follows once, in IT_CHECK)
582 for hook in self.hooks:
583 hook.pre_sweep(step=S, level_number=0)
585 # update QDelta coefficients if variable preconditioner, indexed by the sweep count across iterations
586 S.levels[0].sweep.updateVariableCoeffs((S.status.iter - 1) * nsweeps + k + 1)
587 S.levels[0].sweep.update_nodes()
589 for hook in self.hooks:
590 hook.post_sweep(step=S, level_number=0)
592 for S in local_MS_running:
593 # update stage
594 S.status.stage = 'IT_CHECK'
596 def it_down(self, local_MS_running):
597 """
598 Go down the hierarchy from finest to coarsest level
600 Args:
601 local_MS_running (list): list of currently running steps
602 """
604 for S in local_MS_running:
605 S.transfer(source=S.levels[0], target=S.levels[1])
607 for l in range(1, self.nlevels - 1):
608 # sweep on middle levels (not on finest, not on coarsest, though)
610 nsweeps = local_MS_running[0].levels[l].params.nsweeps
612 for _ in range(nsweeps):
613 for S in local_MS_running:
614 # send updated values forward
615 self.send_full(S, level=l)
616 # receive values
617 self.recv_full(S, level=l)
619 for S in local_MS_running:
620 for hook in self.hooks:
621 hook.pre_sweep(step=S, level_number=l)
622 S.levels[l].sweep.update_nodes()
623 for hook in self.hooks:
624 hook.post_sweep(step=S, level_number=l)
626 for S in local_MS_running:
627 # transfer further down the hierarchy
628 S.transfer(source=S.levels[l], target=S.levels[l + 1])
630 for S in local_MS_running:
631 # update stage
632 S.status.stage = 'IT_COARSE'
634 def it_coarse(self, local_MS_running):
635 """
636 Coarse sweep
638 Args:
639 local_MS_running (list): list of currently running steps
640 """
642 for S in local_MS_running:
643 # receive from previous step (if not first)
644 self.recv_full(S, level=len(S.levels) - 1)
646 # do the sweep
647 for hook in self.hooks:
648 hook.pre_sweep(step=S, level_number=len(S.levels) - 1)
649 S.levels[-1].sweep.update_nodes()
650 for hook in self.hooks:
651 hook.post_sweep(step=S, level_number=len(S.levels) - 1)
653 # send to succ step
654 self.send_full(S, level=len(S.levels) - 1, add_to_stats=True)
656 # update stage
657 if len(S.levels) > 1: # MLSDC or PFASST
658 S.status.stage = 'IT_UP'
659 else: # MSSDC
660 S.status.stage = 'IT_CHECK'
662 def it_up(self, local_MS_running):
663 """
664 Prolong corrections up to finest level (parallel)
666 Args:
667 local_MS_running (list): list of currently running steps
668 """
670 for l in range(self.nlevels - 1, 0, -1):
671 for S in local_MS_running:
672 # prolong values
673 S.transfer(source=S.levels[l], target=S.levels[l - 1])
675 # on middle levels: do communication and sweep as usual
676 if l - 1 > 0:
677 nsweeps = local_MS_running[0].levels[l - 1].params.nsweeps
679 for k in range(nsweeps):
680 for S in local_MS_running:
681 # send updated values forward
682 self.send_full(S, level=l - 1)
683 # receive values
684 self.recv_full(S, level=l - 1, add_to_stats=(k == nsweeps - 1))
686 for S in local_MS_running:
687 for hook in self.hooks:
688 hook.pre_sweep(step=S, level_number=l - 1)
689 S.levels[l - 1].sweep.update_nodes()
690 for hook in self.hooks:
691 hook.post_sweep(step=S, level_number=l - 1)
693 for S in local_MS_running:
694 # update stage
695 S.status.stage = 'IT_FINE'
697 def default(self, local_MS_running):
698 """
699 Default routine to catch wrong status
701 Args:
702 local_MS_running (list): list of currently running steps
703 """
704 raise ControllerError('Unknown stage, got %s' % local_MS_running[0].status.stage) # TODO