Coverage for pySDC/implementations/controller_classes/controller_nonMPI.py: 99%
297 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-25 20:28 +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 PFASST controller, running serialized version of PFASST in blocks (MG-style)
17 """
19 def __init__(self, num_procs, controller_params, description):
20 """
21 Initialization routine for PFASST controller
23 Args:
24 num_procs: number of parallel time steps (still serial, though), can be 1
25 controller_params: parameter set for the controller and the steps
26 description: all the parameters to set up the rest (levels, problems, transfer, ...)
27 """
29 if 'predict' in controller_params:
30 raise ControllerError('predict flag is ignored, use predict_type instead')
32 # call parent's initialization routine
33 super().__init__(controller_params, description, useMPI=False)
35 self.MS: list[Step] = [Step(description)]
37 # try to initialize via dill.copy (much faster for many time-steps)
38 try:
39 for _ in range(num_procs - 1):
40 self.MS.append(dill.copy(self.MS[0]))
41 # if this fails (e.g. due to un-picklable data in the steps), initialize separately
42 except (dill.PicklingError, TypeError, ValueError) as error:
43 self.logger.warning(f'Need to initialize steps separately due to pickling error: {error}')
44 for _ in range(num_procs - 1):
45 self.MS.append(Step(description))
47 self.base_convergence_controllers += [BasicRestarting.get_implementation(useMPI=False)]
48 for convergence_controller in self.base_convergence_controllers:
49 self.add_convergence_controller(convergence_controller, description)
51 if self.params.dump_setup:
52 self.dump_setup(step=self.MS[0], controller_params=controller_params, description=description)
54 if num_procs > 1 and len(self.MS[0].levels) > 1:
55 for S in self.MS:
56 for L in S.levels:
57 if not L.sweep.coll.right_is_node:
58 raise ControllerError("For PFASST to work, we assume uend^k = u_M^k")
60 if all(len(S.levels) == len(self.MS[0].levels) for S in self.MS):
61 self.nlevels = len(self.MS[0].levels)
62 else:
63 raise ControllerError('all steps need to have the same number of levels')
65 if self.nlevels == 0:
66 raise ControllerError('need at least one level')
68 # The stages read `nsweeps` off a step they own, the same way `controller_MPI` reads it off
69 # the one step a rank has, so this only has to establish that any step will do.
70 for nl in range(self.nlevels):
71 if not all(S.levels[nl].params.nsweeps == self.MS[0].levels[nl].params.nsweeps for S in self.MS):
72 raise ControllerError('all steps need to agree on the number of sweeps per level')
74 # `it_coarse` sweeps the coarsest level exactly once. Single-level Gauss-like MSSDC routes
75 # through it too, so reject multiple sweeps there as well instead of silently ignoring them.
76 # `mssdc_jac` only decides the routing when there is more than one step: a single step is
77 # plain SDC and always goes through `it_fine`, which honours nsweeps.
78 if self.MS[0].levels[-1].params.nsweeps > 1 and (
79 self.nlevels > 1 or (num_procs > 1 and not self.params.mssdc_jac)
80 ):
81 raise ControllerError('this controller cannot do multiple sweeps on coarsest level')
83 self.check_variable_coefficients(num_procs)
85 if self.nlevels == 1 and self.params.predict_type is not None:
86 self.logger.warning(
87 'you have specified a predictor type but only a single level.. predictor will be ignored'
88 )
90 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
91 C.reset_buffers_nonMPI(self)
92 C.setup_status_variables(self, MS=self.MS)
94 self.stages = self.get_stages()
96 def run(self, u0, t0, Tend):
97 """
98 Main driver for running the serial version of SDC, MSSDC, MLSDC and PFASST (virtual parallelism)
100 Args:
101 u0: initial values
102 t0: starting time
103 Tend: ending time
105 Returns:
106 end values on the finest level
107 stats object containing statistics for each step, each level and each iteration
108 """
110 # some initializations and reset of statistics
111 uend = None
112 num_procs = len(self.MS)
113 for hook in self.hooks:
114 hook.reset_stats()
116 # initial ordering of the steps: 0,1,...,Np-1
117 slots = list(range(num_procs))
119 # initialize time variables of each step
120 time = [t0 + sum(self.MS[j].dt for j in range(p)) for p in slots]
122 # determine which steps are still active (time < Tend)
123 active = [self.step_is_active(time[p], time[0], Tend) for p in slots]
125 if not any(active):
126 raise ControllerError('Nothing to do, check t0, dt and Tend.')
128 # compress slots according to active steps, i.e. remove all steps which have times above Tend
129 active_slots = list(itertools.compress(slots, active))
131 # initialize block of steps with u0
132 self.restart_block(active_slots, time, u0)
134 for hook in self.hooks:
135 hook.post_setup(step=None, level_number=None)
137 # call pre-run hook
138 for S in self.MS:
139 for hook in self.hooks:
140 hook.pre_run(step=S, level_number=0)
142 # main loop: as long as at least one step is still active (time < Tend), do something
143 while any(active):
144 MS_active = [self.MS[p] for p in active_slots]
145 done = False
146 while not done:
147 done = self.pfasst(MS_active)
149 restarts = [S.status.restart for S in MS_active]
150 restart_at = np.where(restarts)[0][0] if True in restarts else len(MS_active)
151 if True in restarts: # restart part of the block
152 # initial condition to next block is initial condition of step that needs restarting
153 uend = self.MS[restart_at].levels[0].u[0]
154 time[active_slots[0]] = time[restart_at]
155 self.logger.info(f'Starting next block with initial conditions from step {restart_at}')
157 else: # move on to next block
158 # initial condition for next block is last solution of current block
159 uend = self.MS[active_slots[-1]].levels[0].uend
160 time[active_slots[0]] = time[active_slots[-1]] + self.MS[active_slots[-1]].dt
162 for S in MS_active[:restart_at]:
163 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
164 C.post_step_processing(self, S, MS=MS_active)
166 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
167 [C.prepare_next_block(self, S, len(active_slots), time, Tend, MS=MS_active) for S in self.MS]
169 # setup the times of the steps for the next block
170 for i in range(1, len(active_slots)):
171 time[active_slots[i]] = time[active_slots[i] - 1] + self.MS[active_slots[i] - 1].dt
173 # determine new set of active steps and compress slots accordingly
174 active = [self.step_is_active(time[p], time[0], Tend) for p in slots]
175 active_slots = list(itertools.compress(slots, active))
177 # restart active steps (reset all values and pass uend to u0)
178 self.restart_block(active_slots, time, uend)
180 # call post-run hook
181 for S in self.MS:
182 for hook in self.hooks:
183 hook.post_run(step=S, level_number=0)
185 for S in self.MS:
186 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
187 C.post_run_processing(self, S, MS=MS_active)
189 return uend, self.return_stats()
191 def restart_block(self, active_slots, time, u0):
192 """
193 Helper routine to reset/restart block of (active) steps
195 Args:
196 active_slots: list of active steps
197 time: list of new times
198 u0: initial value to distribute across the steps
200 """
202 # loop over active slots (not directly, since we need the previous entry as well)
203 for j in range(len(active_slots)):
204 # get slot number
205 p = active_slots[j]
207 # store current slot number for diagnostics
208 self.MS[p].status.slot = p
209 # store link to previous step
210 self.MS[p].prev = self.MS[active_slots[j - 1]]
211 # resets step
212 self.MS[p].reset_step()
213 # determine whether I am the first and/or last in line
214 self.MS[p].status.first = active_slots.index(p) == 0
215 self.MS[p].status.last = active_slots.index(p) == len(active_slots) - 1
216 # initialize step with u0
217 self.MS[p].init_step(u0)
218 # reset some values
219 self.MS[p].status.done = False
220 self.MS[p].status.prev_done = False
221 self.MS[p].status.iter = 0
222 self.MS[p].status.stage = 'SPREAD'
223 self.MS[p].status.force_done = False
224 self.MS[p].status.time_size = len(active_slots)
226 for l in self.MS[p].levels:
227 l.tag = None
228 l.status.sweep = 1
230 for p in active_slots:
231 for lvl in self.MS[p].levels:
232 lvl.status.time = time[p]
234 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
235 C.reset_status_variables(self, active_slots=active_slots)
237 def send_full(self, S, level=None, add_to_stats=False):
238 """
239 Function to perform the send, including bookkeeping and logging
241 Args:
242 S: the current step
243 level: the level number
244 add_to_stats: a flag to end recording data in the hooks (defaults to False)
246 Note:
247 Computing the end point is this function's job, not the caller's, exactly as in
248 `controller_MPI.send_full`. It happens whether or not anyone is listening, because the
249 last step has no successor but its `uend` is still the block's result.
250 """
251 for hook in self.hooks:
252 hook.pre_comm(step=S, level_number=level)
254 # sending here means computing uend ("one-sided communication")
255 S.levels[level].sweep.compute_end_point()
257 if not S.status.last:
258 self.logger.debug(
259 'Process %2i provides data on level %2i with tag %s' % (S.status.slot, level, S.status.iter)
260 )
261 S.levels[level].tag = cp.deepcopy((level, S.status.iter, S.status.slot))
263 for hook in self.hooks:
264 hook.post_comm(step=S, level_number=level, add_to_stats=add_to_stats)
266 def recv_full(self, S, level=None, add_to_stats=False):
267 """
268 Function to perform the recv, including bookkeeping and logging
270 Args:
271 S: the current step
272 level: the level number
273 add_to_stats: a flag to end recording data in the hooks (defaults to False)
274 """
276 def recv(target, source, tag=None):
277 """
278 Receive function
280 Args:
281 target: level which will receive the values
282 source: level which initiated the send
283 tag: identifier to check if this message is really for me
284 """
286 if tag is not None and source.tag != tag:
287 raise CommunicationError('source and target tag are not the same, got %s and %s' % (source.tag, tag))
288 # simply do a deepcopy of the values uend to become the new u0 at the target
289 target.u[0] = target.prob.dtype_u(source.uend)
290 # re-evaluate f on left interval boundary
291 target.f[0] = target.prob.eval_f(target.u[0], target.time)
293 for hook in self.hooks:
294 hook.pre_comm(step=S, level_number=level)
295 if not S.status.prev_done and not S.status.first:
296 self.logger.debug(
297 'Process %2i receives from %2i on level %2i with tag %s'
298 % (S.status.slot, S.prev.status.slot, level, S.status.iter)
299 )
300 recv(S.levels[level], S.prev.levels[level], tag=(level, S.status.iter, S.prev.status.slot))
301 for hook in self.hooks:
302 hook.post_comm(step=S, level_number=level, add_to_stats=add_to_stats)
304 def pfasst(self, local_MS_active):
305 """
306 Main function including the stages of SDC, MLSDC and PFASST (the "controller")
308 For the workflow of this controller, check out one of our PFASST talks or the pySDC paper
310 This method changes self.MS directly by accessing active steps through local_MS_active. Nothing is returned.
312 Args:
313 local_MS_active (list): all active steps
314 """
316 # if all stages are the same (or DONE), continue, otherwise abort
317 stages = [S.status.stage for S in local_MS_active if S.status.stage != 'DONE']
318 if stages[1:] == stages[:-1]:
319 stage = stages[0]
320 else:
321 raise ControllerError('not all stages are equal')
323 self.logger.debug(stage)
325 MS_running = [S for S in local_MS_active if S.status.stage != 'DONE']
327 self.stages.get(stage, self.default)(MS_running)
329 return all(S.status.done for S in local_MS_active)
331 def get_stages(self):
332 """
333 The stages this controller can be in, and what to run in each.
335 A subclass that iterates differently replaces the iteration stages here and says which one
336 to enter in `next_iteration_stage`; everything around the iteration is the same for any
337 algorithm this controller runs.
339 Returns:
340 dict: stage name -> the method that runs it
341 """
342 return {
343 'SPREAD': self.spread,
344 'PREDICT': self.predict,
345 'IT_CHECK': self.it_check,
346 'IT_FINE': self.it_fine,
347 'IT_DOWN': self.it_down,
348 'IT_COARSE': self.it_coarse,
349 'IT_UP': self.it_up,
350 }
352 def next_iteration_stage(self, S):
353 """
354 The stage that starts one iteration of the algorithm.
356 Args:
357 S (pySDC.Step.step): The current step
359 Returns:
360 str: name of the stage to enter
361 """
362 if len(S.levels) > 1: # MLSDC or PFASST
363 return 'IT_DOWN'
364 elif S.status.time_size == 1 or self.params.mssdc_jac: # SDC or parallel MSSDC (Jacobi-like)
365 return 'IT_FINE'
366 else:
367 return 'IT_COARSE' # serial MSSDC (Gauss-like)
369 def spread(self, local_MS_running):
370 """
371 Spreading phase
373 Args:
374 local_MS_running (list): list of currently running steps
375 """
377 for S in local_MS_running:
378 # first stage: spread values
379 for hook in self.hooks:
380 hook.pre_step(step=S, level_number=0)
382 # call predictor from sweeper
383 S.levels[0].sweep.predict()
385 self.compute_residual_after_spread(S)
387 # update stage
388 if len(S.levels) > 1: # MLSDC or PFASST with predict
389 S.status.stage = 'PREDICT'
390 else:
391 S.status.stage = 'IT_CHECK'
393 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
394 C.post_spread_processing(self, S, MS=local_MS_running)
396 def predict(self, local_MS_running):
397 """
398 Predictor phase
400 Args:
401 local_MS_running (list): list of currently running steps
402 """
404 for S in local_MS_running:
405 for hook in self.hooks:
406 hook.pre_predict(step=S, level_number=0)
408 if self.params.predict_type is None:
409 pass
411 elif self.params.predict_type == 'fine_only':
412 # do a fine sweep only
413 for S in local_MS_running:
414 S.levels[0].sweep.update_nodes()
416 elif self.params.predict_type == 'pfasst_burnin':
417 # loop over all steps
418 for S in local_MS_running:
419 # restrict to coarsest level
420 for l in range(1, len(S.levels)):
421 S.transfer(source=S.levels[l - 1], target=S.levels[l])
423 # loop over all steps
424 for q in range(len(local_MS_running)):
425 # loop over last steps: [1,2,3,4], [2,3,4], [3,4], [4]
426 for p in range(q, len(local_MS_running)):
427 S = local_MS_running[p]
429 # do the sweep with new values
430 S.levels[-1].sweep.update_nodes()
432 # send updated values on coarsest level
433 self.send_full(S, level=len(S.levels) - 1)
435 # loop over last steps: [2,3,4], [3,4], [4]
436 for p in range(q + 1, len(local_MS_running)):
437 S = local_MS_running[p]
438 # receive values sent during previous sweep
439 self.recv_full(S, level=len(S.levels) - 1, add_to_stats=(p == len(local_MS_running) - 1))
441 # loop over all steps
442 for S in local_MS_running:
443 # interpolate back to finest level
444 for l in range(len(S.levels) - 1, 0, -1):
445 S.transfer(source=S.levels[l], target=S.levels[l - 1])
447 # send updated values forward
448 self.send_full(S, level=0)
449 # receive values
450 self.recv_full(S, level=0)
452 # end this with a fine sweep
453 for S in local_MS_running:
454 S.levels[0].sweep.update_nodes()
456 elif self.params.predict_type == 'fmg':
457 # TODO: implement FMG predictor
458 raise NotImplementedError('FMG predictor is not yet implemented')
460 else:
461 raise ControllerError('Wrong predictor type, got %s' % self.params.predict_type)
463 for S in local_MS_running:
464 for hook in self.hooks:
465 hook.post_predict(step=S, level_number=0)
467 for S in local_MS_running:
468 # update stage
469 S.status.stage = 'IT_CHECK'
471 def it_check(self, local_MS_running):
472 """
473 Key routine to check for convergence/termination
475 Args:
476 local_MS_running (list): list of currently running steps
477 """
479 self.prepare_convergence_check(local_MS_running)
481 for S in local_MS_running:
482 if S.status.iter > 0:
483 for hook in self.hooks:
484 hook.post_iteration(step=S, level_number=0)
486 # decide if the step is done, needs to be restarted and other things convergence related
487 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
488 C.post_iteration_processing(self, S, MS=local_MS_running)
489 C.convergence_control(self, S, MS=local_MS_running)
491 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
492 C.post_iteration_processing_block(self, MS=local_MS_running)
494 for S in local_MS_running:
495 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
496 C.communicate_convergence(self, S, MS=local_MS_running)
498 if not S.status.done:
499 # increment iteration count here (and only here)
500 S.status.iter += 1
501 for hook in self.hooks:
502 hook.pre_iteration(step=S, level_number=0)
503 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
504 C.pre_iteration_processing(self, S, MS=local_MS_running)
506 S.status.stage = self.next_iteration_stage(S)
507 else:
508 for hook in self.hooks:
509 hook.post_step(step=S, level_number=0)
510 S.status.stage = 'DONE'
512 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
513 C.reset_buffers_nonMPI(self)
515 def compute_residual_after_spread(self, S):
516 """
517 Make the residual current right after the initial guess, for algorithms that need it there.
519 This controller computes the residual at the top of `it_check` instead, so there is nothing
520 to do; an algorithm whose residual is not available at that point overrides this. Called
521 before `post_spread_processing`, because convergence controllers there may still change the
522 initial iterate.
524 Args:
525 S (pySDC.Step.step): The current step
526 """
527 pass
529 def prepare_convergence_check(self, local_MS_running):
530 """
531 Make current what `it_check` is about to read: the end point and the residual.
533 For a sweep-based algorithm both come out of the exchange with the neighbours, so this is
534 the send and receive plus the residual it enables. An algorithm that gets its residual from
535 the iteration instead overrides this, and still owes the end point.
537 Args:
538 local_MS_running (list): list of currently running steps
539 """
540 for S in local_MS_running:
541 # send updated values forward
542 self.send_full(S, level=0)
543 # receive values
544 self.recv_full(S, level=0)
545 # compute current residual
546 S.levels[0].sweep.compute_residual(stage='IT_CHECK')
548 def it_fine(self, local_MS_running: list[Step]):
549 """
550 Fine sweeps
552 Args:
553 local_MS_running (list): list of currently running steps
554 """
556 for S in local_MS_running:
557 S.levels[0].status.sweep = 0
559 nsweeps = local_MS_running[0].levels[0].params.nsweeps
561 for k in range(nsweeps):
562 for S in local_MS_running:
563 S.levels[0].status.sweep += 1
565 for S in local_MS_running:
566 # send updated values forward
567 self.send_full(S, level=0)
568 # receive values
569 self.recv_full(S, level=0, add_to_stats=(k == nsweeps - 1))
571 for S in local_MS_running:
572 # standard sweep workflow: update nodes, compute residual, log progress
573 for hook in self.hooks:
574 hook.pre_sweep(step=S, level_number=0)
576 S.levels[0].sweep.updateVariableCoeffs(k + 1) # update QDelta coefficients if variable preconditioner
577 S.levels[0].sweep.update_nodes()
578 S.levels[0].sweep.compute_residual(stage='IT_FINE')
580 for hook in self.hooks:
581 hook.post_sweep(step=S, level_number=0)
583 for S in local_MS_running:
584 # update stage
585 S.status.stage = 'IT_CHECK'
587 def it_down(self, local_MS_running):
588 """
589 Go down the hierarchy from finest to coarsest level
591 Args:
592 local_MS_running (list): list of currently running steps
593 """
595 for S in local_MS_running:
596 S.transfer(source=S.levels[0], target=S.levels[1])
598 for l in range(1, self.nlevels - 1):
599 # sweep on middle levels (not on finest, not on coarsest, though)
601 nsweeps = local_MS_running[0].levels[l].params.nsweeps
603 for _ in range(nsweeps):
604 for S in local_MS_running:
605 # send updated values forward
606 self.send_full(S, level=l)
607 # receive values
608 self.recv_full(S, level=l)
610 for S in local_MS_running:
611 for hook in self.hooks:
612 hook.pre_sweep(step=S, level_number=l)
613 S.levels[l].sweep.update_nodes()
614 S.levels[l].sweep.compute_residual(stage='IT_DOWN')
615 for hook in self.hooks:
616 hook.post_sweep(step=S, level_number=l)
618 for S in local_MS_running:
619 # transfer further down the hierarchy
620 S.transfer(source=S.levels[l], target=S.levels[l + 1])
622 for S in local_MS_running:
623 # update stage
624 S.status.stage = 'IT_COARSE'
626 def it_coarse(self, local_MS_running):
627 """
628 Coarse sweep
630 Args:
631 local_MS_running (list): list of currently running steps
632 """
634 for S in local_MS_running:
635 # receive from previous step (if not first)
636 self.recv_full(S, level=len(S.levels) - 1)
638 # do the sweep
639 for hook in self.hooks:
640 hook.pre_sweep(step=S, level_number=len(S.levels) - 1)
641 S.levels[-1].sweep.update_nodes()
642 S.levels[-1].sweep.compute_residual(stage='IT_COARSE')
643 for hook in self.hooks:
644 hook.post_sweep(step=S, level_number=len(S.levels) - 1)
646 # send to succ step
647 self.send_full(S, level=len(S.levels) - 1, add_to_stats=True)
649 # update stage
650 if len(S.levels) > 1: # MLSDC or PFASST
651 S.status.stage = 'IT_UP'
652 else: # MSSDC
653 S.status.stage = 'IT_CHECK'
655 def it_up(self, local_MS_running):
656 """
657 Prolong corrections up to finest level (parallel)
659 Args:
660 local_MS_running (list): list of currently running steps
661 """
663 for l in range(self.nlevels - 1, 0, -1):
664 for S in local_MS_running:
665 # prolong values
666 S.transfer(source=S.levels[l], target=S.levels[l - 1])
668 # on middle levels: do communication and sweep as usual
669 if l - 1 > 0:
670 nsweeps = local_MS_running[0].levels[l - 1].params.nsweeps
672 for k in range(nsweeps):
673 for S in local_MS_running:
674 # send updated values forward
675 self.send_full(S, level=l - 1)
676 # receive values
677 self.recv_full(S, level=l - 1, add_to_stats=(k == nsweeps - 1))
679 for S in local_MS_running:
680 for hook in self.hooks:
681 hook.pre_sweep(step=S, level_number=l - 1)
682 S.levels[l - 1].sweep.update_nodes()
683 S.levels[l - 1].sweep.compute_residual(stage='IT_UP')
684 for hook in self.hooks:
685 hook.post_sweep(step=S, level_number=l - 1)
687 for S in local_MS_running:
688 # update stage
689 S.status.stage = 'IT_FINE'
691 def default(self, local_MS_running):
692 """
693 Default routine to catch wrong status
695 Args:
696 local_MS_running (list): list of currently running steps
697 """
698 raise ControllerError('Unknown stage, got %s' % local_MS_running[0].status.stage) # TODO