Coverage for pySDC/implementations/convergence_controller_classes/spread_step_sizes.py: 100%
61 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1import numpy as np
2from pySDC.core.convergence_controller import ConvergenceController
5class SpreadStepSizesBlockwise(ConvergenceController):
6 """
7 Give all steps of the next block the same step size, taken from the last step or from the first restarted one.
8 By block we refer to a composite collocation problem, which is solved in pipelined SDC parallel-in-time.
10 Also, we overrule the step size control here, if we get close to the final time and we would take too large of a
11 step otherwise.
12 """
14 def setup(self, controller, params, description, **kwargs):
15 """
16 Define parameters here
18 Args:
19 controller (pySDC.Controller): The controller
20 params (dict): The params passed for this specific convergence controller
21 description (dict): The description object used to instantiate the controller
23 Returns:
24 (dict): The updated params dictionary
25 """
26 defaults = {
27 "control_order": +100,
28 "spread_from_first_restarted": True,
29 "overwrite_to_reach_Tend": True,
30 }
32 return {**defaults, **super().setup(controller, params, description, **kwargs)}
34 @classmethod
35 def get_implementation(cls, useMPI, **kwargs):
36 """
37 Get MPI or non-MPI version
39 Args:
40 useMPI (bool): The implementation that you want
42 Returns:
43 cls: The child class implementing the desired flavor
44 """
45 if useMPI:
46 return SpreadStepSizesBlockwiseMPI
47 else:
48 return SpreadStepSizesBlockwiseNonMPI
50 def get_step_from_which_to_spread(self, restarts, new_steps, size, S):
51 """
52 Return the index of the step from which to spread the step size to all steps in the next block.
54 Args:
55 restarts (list): List of booleans for each step, showing if it wants to be restarted
56 new_steps (list): List of the new step sizes on the finest level of each step
57 size (int): Size of the time communicator
58 S (pySDC.Step.step): The current step
60 Returns:
61 int: The index of the step from which we want to spread the step size
62 """
63 if True in restarts:
64 restart_at = np.where(restarts)[0][0]
65 if self.params.spread_from_first_restarted:
66 spread_from_step = restart_at
67 else:
68 # we want to spread the smallest step size out of the steps that want to be restarted
69 spread_from_step = restart_at + np.argmin(new_steps[restart_at:])
70 self.debug(
71 f'Detected restart from step {restart_at}. Spreading step size from step {spread_from_step}: {new_steps[restart_at]:.2e}.',
72 S,
73 )
74 else:
75 restart_at = size - 1
76 spread_from_step = restart_at
77 self.debug(f'Spreading step size from last step: {new_steps[restart_at]:.2e}.', S)
79 return spread_from_step, restart_at
82class SpreadStepSizesBlockwiseNonMPI(SpreadStepSizesBlockwise):
83 """
84 Spread one step size to all steps of the next block for the non-MPI controller, reading the steps directly.
85 """
87 def get_step_from_which_to_spread(self, MS, S):
88 """
89 Return the index of the step from which to spread the step size to all steps in the next block.
91 Args:
92 MS (list): Active steps
93 S (pySDC.Step.step): The current step
95 Returns:
96 int: The index of the step from which we want to spread the step size
97 """
98 restarts = [me.status.restart for me in MS]
99 new_steps = [me.levels[0].status.dt_new if me.levels[0].status.dt_new else 1e9 for me in MS]
100 return super().get_step_from_which_to_spread(restarts, new_steps, len(MS), S)
102 def prepare_next_block(self, controller, S, size, time, Tend, MS, **kwargs):
103 """
104 Spread the step size of the last step with no restarted predecessors to all steps and limit the step size based
105 on Tend
107 Args:
108 controller (pySDC.Controller): The controller
109 S (pySDC.step): The current step
110 size (int): The number of ranks
111 time (list): List containing the time of all the steps handled by the controller (or float in MPI implementation)
112 Tend (float): Final time of the simulation
113 MS (list): Active steps
115 Returns:
116 None
117 """
118 # inactive steps don't need to participate
119 if S not in MS:
120 return None
122 spread_from_step, restart_at = self.get_step_from_which_to_spread(MS, S)
124 # Compute the maximum allowed step size based on Tend.
125 dt_all = [0.0] + [me.dt for me in MS if not me.status.first]
126 dt_max = (
127 (Tend - time[restart_at] - dt_all[restart_at]) / size if self.params.overwrite_to_reach_Tend else np.inf
128 )
130 # record the step sizes to restart with from all the levels of the step
131 new_steps = [None] * len(S.levels)
132 for i in range(len(MS[spread_from_step].levels)):
133 l = MS[spread_from_step].levels[i]
134 # overrule the step size control to reach Tend if needed
135 new_steps[i] = min(
136 [
137 l.status.dt_new if l.status.dt_new is not None else l.params.dt,
138 max([dt_max, l.params.dt_initial]),
139 ]
140 )
141 if (
142 new_steps[i] < (l.status.dt_new if l.status.dt_new is not None else l.params.dt)
143 and i == 0
144 and l.status.dt_new is not None
145 ):
146 self.log(
147 f"Overwriting stepsize control to reach Tend: {Tend:.2e}! New step size: {new_steps[i]:.2e}", S
148 )
150 # spread the step sizes to all levels
151 for i in range(len(S.levels)):
152 S.levels[i].params.dt = new_steps[i]
154 return None
157class SpreadStepSizesBlockwiseMPI(SpreadStepSizesBlockwise):
158 """
159 Spread one step size to all steps of the next block for the MPI controller, with collectives across the time ranks.
160 """
162 def get_step_from_which_to_spread(self, comm, S):
163 """
164 Return the index of the step from which to spread the step size to all steps in the next block.
166 Args:
167 comm (mpi4py.MPI.Intracomm): Communicator
168 S (pySDC.Step.step): The current step
170 Returns:
171 int: The index of the step from which we want to spread the step size
172 """
173 restarts = comm.allgather(S.status.restart)
174 new_steps = [me if me is not None else 1e9 for me in comm.allgather(S.levels[0].status.dt_new)]
176 return super().get_step_from_which_to_spread(restarts, new_steps, comm.size, S)
178 def prepare_next_block(self, controller, S, size, time, Tend, comm, **kwargs):
179 """
180 Spread the step size of the last step with no restarted predecessors to all steps and limit the step size based
181 on Tend
183 Args:
184 controller (pySDC.Controller): The controller
185 S (pySDC.step): The current step
186 size (int): The number of ranks
187 time (float): Time of the first step
188 Tend (float): Final time of the simulation
189 comm (mpi4py.MPI.Intracomm): Communicator
191 Returns:
192 None
193 """
194 spread_from_step, restart_at = self.get_step_from_which_to_spread(comm, S)
196 # Compute the maximum allowed step size based on Tend.
197 dt_max = comm.bcast((Tend - time) / size, root=restart_at) if self.params.overwrite_to_reach_Tend else np.inf
199 # record the step sizes to restart with from all the levels of the step
200 new_steps = [None] * len(S.levels)
201 if S.status.slot == spread_from_step:
202 for i in range(len(S.levels)):
203 l = S.levels[i]
204 # overrule the step size control to reach Tend if needed
205 new_steps[i] = min(
206 [
207 l.status.dt_new if l.status.dt_new is not None else l.params.dt,
208 max([dt_max, l.params.dt_initial]),
209 ]
210 )
212 if (
213 new_steps[i] < l.status.dt_new
214 if l.status.dt_new is not None
215 else l.params.dt and l.status.dt_new is not None
216 ):
217 self.log(
218 f"Overwriting stepsize control to reach Tend: {Tend:.2e}! New step size: {new_steps[i]:.2e}", S
219 )
220 new_steps = comm.bcast(new_steps, root=spread_from_step)
222 # spread the step sizes to all levels
223 for i in range(len(S.levels)):
224 S.levels[i].params.dt = new_steps[i]
226 return None