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