Small cleanup in convergence controllers

pySDC/core/ConvergenceController.py

@@ -263,26 +263,7 @@ def prepare_next_block(self, controller, S, size, time, Tend, **kwargs):
             controller (pySDC.Controller): The controller
             S (pySDC.Step): The current step
             size (int): The number of ranks
-            time (float): The current time
-            Tend (float): The final time
-
-        Returns:
-            None
-        """
-        pass
-
-    def prepare_next_block_nonMPI(self, controller, MS, active_slots, time, Tend):
-        """
-        This is an extension to the function `prepare_next_block`, which is only called in the non MPI controller and
-        is needed because there is no chance to communicate backwards otherwise. While you should not do this in the
-        first place, the first step in the new block comes after the last step in the last block, such that it is still
-        in fact forwards communication, even though it looks backwards.
-
-        Args:
-            controller (pySDC.Controller): The controller
-            MS (list): All steps of the controller
-            active_slots (list): Index list of active steps
-            time (float): The current time
+            time (float): The current time will be list in nonMPI controller implementation
             Tend (float): The final time
 
         Returns:

pySDC/implementations/controller_classes/controller_MPI.py

@@ -62,7 +62,7 @@ def __init__(self, controller_params, description, comm):
 
         if num_levels == 1 and self.params.predict_type is not None:
             self.logger.warning(
-                'you have specified a predictor type but only a single level.. ' 'predictor will be ignored'
+                'you have specified a predictor type but only a single level.. predictor will be ignored'
             )
 
         for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
@@ -632,7 +632,7 @@ def it_check(self, comm, num_procs):
                 C.pre_iteration_processing(self, self.S, comm=comm)
 
             if self.params.use_iteration_estimator:
-                # store pervious iterate to compute difference later on
+                # store previous iterate to compute difference later on
                 self.S.levels[0].uold[1:] = self.S.levels[0].u[1:]
 
             if len(self.S.levels) > 1:  # MLSDC or PFASST

pySDC/implementations/controller_classes/controller_nonMPI.py

@@ -75,12 +75,12 @@ def __init__(self, num_procs, controller_params, description):
 
         if self.nlevels == 1 and self.params.predict_type is not None:
             self.logger.warning(
-                'you have specified a predictor type but only a single level.. ' 'predictor will be ignored'
+                'you have specified a predictor type but only a single level.. predictor will be ignored'
             )
 
         for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
             C.reset_buffers_nonMPI(self)
-            C.setup_status_variables(self)
+            C.setup_status_variables(self, MS=self.MS)
 
     def run(self, u0, t0, Tend):
         """
@@ -153,8 +153,7 @@ def run(self, u0, t0, Tend):
                     C.post_step_processing(self, S)
 
             for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
-                [C.prepare_next_block(self, S, len(active_slots), time, Tend) for S in self.MS]
-                C.prepare_next_block_nonMPI(self, self.MS, active_slots, time, Tend)
+                [C.prepare_next_block(self, S, len(active_slots), time, Tend, MS=MS_active) for S in self.MS]
 
             # setup the times of the steps for the next block
             for i in range(1, len(active_slots)):
@@ -351,7 +350,7 @@ def spread(self, local_MS_running):
                 S.status.stage = 'IT_CHECK'
 
             for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
-                C.post_spread_processing(self, S)
+                C.post_spread_processing(self, S, MS=local_MS_running)
 
     def predict(self, local_MS_running):
         """
@@ -496,8 +495,8 @@ def it_check(self, local_MS_running):
 
             # decide if the step is done, needs to be restarted and other things convergence related
             for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
-                C.post_iteration_processing(self, S)
-                C.convergence_control(self, S)
+                C.post_iteration_processing(self, S, MS=local_MS_running)
+                C.convergence_control(self, S, MS=local_MS_running)
 
         for S in local_MS_running:
             if not S.status.first:
@@ -521,7 +520,7 @@ def it_check(self, local_MS_running):
                 for hook in self.hooks:
                     hook.pre_iteration(step=S, level_number=0)
                 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
-                    C.pre_iteration_processing(self, S)
+                    C.pre_iteration_processing(self, S, MS=local_MS_running)
 
                 if len(S.levels) > 1:  # MLSDC or PFASST
                     S.status.stage = 'IT_DOWN'

pySDC/implementations/convergence_controller_classes/basic_restarting.py

@@ -1,7 +1,6 @@
 from pySDC.core.ConvergenceController import ConvergenceController, Pars
 from pySDC.implementations.convergence_controller_classes.spread_step_sizes import (
-    SpreadStepSizesBlockwiseNonMPI,
-    SpreadStepSizesBlockwiseMPI,
+    SpreadStepSizesBlockwise,
 )
 from pySDC.core.Errors import ConvergenceError
 
@@ -42,7 +41,7 @@ def __init__(self, controller, params, description, **kwargs):
             params (dict): Parameters for the convergence controller
             description (dict): The description object used to instantiate the controller
         """
-        super(BasicRestarting, self).__init__(controller, params, description)
+        super().__init__(controller, params, description)
         self.buffers = Pars({"restart": False, "max_restart_reached": False})
 
     def setup(self, controller, params, description, **kwargs):
@@ -68,6 +67,7 @@ def setup(self, controller, params, description, **kwargs):
             "control_order": 95,
             "max_restarts": 10,
             "crash_after_max_restarts": True,
+            "step_size_spreader": SpreadStepSizesBlockwise.get_implementation(useMPI=params['useMPI']),
         }
 
         from pySDC.implementations.hooks.log_restarts import LogRestarts
@@ -141,8 +141,8 @@ def prepare_next_block(self, controller, S, size, time, Tend, **kwargs):
 
         Args:
             controller (pySDC.Controller): The controller
-            MS (list): List of the steps of the controller
-            active_slots (list): Index list of active steps
+            S (pySDC.Step): The current step
+            size (int): The number of ranks
             time (list): List containing the time of all the steps
             Tend (float): Final time of the simulation
 
@@ -174,33 +174,6 @@ def reset_buffers_nonMPI(self, controller, **kwargs):
 
         return None
 
-    def setup(self, controller, params, description, **kwargs):
-        """
-        Define parameters here.
-
-        Default parameters are:
-         - control_order (int): The order relative to other convergence controllers
-         - max_restarts (int): Maximum number of restarts we allow each step before we just move on with whatever we
-                               have
-         - step_size_spreader (pySDC.ConvergenceController): A convergence controller that takes care of distributing
-                                                             the steps sizes between blocks
-
-        Args:
-            controller (pySDC.Controller): The controller
-            params (dict): The params passed for this specific convergence controller
-            description (dict): The description object used to instantiate the controller
-
-        Returns:
-            (dict): The updated params dictionary
-        """
-        defaults = {
-            "step_size_spreader": SpreadStepSizesBlockwiseNonMPI,
-        }
-        return {
-            **defaults,
-            **super(BasicRestartingNonMPI, self).setup(controller, params, description),
-        }
-
     def determine_restart(self, controller, S, **kwargs):
         """
         Restart all steps after the first one which wants to be restarted as well, but also check if we lost patience
@@ -245,49 +218,22 @@ def __init__(self, controller, params, description, **kwargs):
             params (dict): Parameters for the convergence controller
             description (dict): The description object used to instantiate the controller
         """
-        super(BasicRestartingMPI, self).__init__(controller, params, description)
+        super().__init__(controller, params, description)
         self.buffers = Pars({"restart": False, "max_restart_reached": False, 'restart_earlier': False})
 
-    def setup(self, controller, params, description, **kwargs):
-        """
-        Define parameters here.
-
-        Default parameters are:
-         - control_order (int): The order relative to other convergence controllers
-         - max_restarts (int): Maximum number of restarts we allow each step before we just move on with whatever we
-                               have
-         - step_size_spreader (pySDC.ConvergenceController): A convergence controller that takes care of distributing
-                                                             the steps sizes between blocks
-
-        Args:
-            controller (pySDC.Controller): The controller
-            params (dict): The params passed for this specific convergence controller
-            description (dict): The description object used to instantiate the controller
-
-        Returns:
-            (dict): The updated params dictionary
-        """
-        defaults = {
-            "step_size_spreader": SpreadStepSizesBlockwiseMPI,
-        }
-        return {
-            **defaults,
-            **super(BasicRestartingMPI, self).setup(controller, params, description),
-        }
-
-    def determine_restart(self, controller, S, **kwargs):
+    def determine_restart(self, controller, S, comm, **kwargs):
         """
         Restart all steps after the first one which wants to be restarted as well, but also check if we lost patience
         with the restarts and want to move on anyways.
 
         Args:
             controller (pySDC.Controller): The controller
             S (pySDC.Step): The current step
+            comm (mpi4py.MPI.Intracomm): Communicator
 
         Returns:
             None
         """
-        comm = kwargs['comm']
         assert S.status.slot == comm.rank
 
         if S.status.first:

pySDC/implementations/convergence_controller_classes/estimate_extrapolation_error.py

@@ -27,7 +27,7 @@ def __init__(self, controller, params, description, **kwargs):
         """
         self.prev = Status(["t", "u", "f", "dt"])  # store solutions etc. of previous steps here
         self.coeff = Status(["u", "f", "prefactor"])  # store coefficients for extrapolation here
-        super(EstimateExtrapolationErrorBase, self).__init__(controller, params, description)
+        super().__init__(controller, params, description)
         controller.add_hook(LogExtrapolationErrorEstimate)
 
     def setup(self, controller, params, description, **kwargs):
@@ -265,7 +265,7 @@ def setup(self, controller, params, description, **kwargs):
         Returns:
             dict: Updated parameters with default values
         """
-        default_params = super(EstimateExtrapolationErrorNonMPI, self).setup(controller, params, description)
+        default_params = super().setup(controller, params, description)
 
         non_mpi_defaults = {
             "no_storage": False,
@@ -283,7 +283,7 @@ def setup_status_variables(self, controller, **kwargs):
         Returns:
             None
         """
-        super(EstimateExtrapolationErrorNonMPI, self).setup_status_variables(controller, **kwargs)
+        super().setup_status_variables(controller, **kwargs)
 
         self.prev.t = np.array([None] * self.params.n)
         self.prev.dt = np.array([None] * self.params.n)
@@ -327,21 +327,23 @@ def post_iteration_processing(self, controller, S, **kwargs):
 
         return None
 
-    def prepare_next_block_nonMPI(self, controller, MS, active_slots, time, Tend, **kwargs):
+    def prepare_next_block(self, controller, S, size, time, Tend, MS, **kwargs):
         """
         If the no-memory-overhead version is used, we need to delete stuff that shouldn't be available. Otherwise, we
         need to store all the stuff that we can.
 
         Args:
             controller (pySDC.Controller): The controller
-            MS (list): All steps of the controller
-            active_slots (list): Index list of active steps
+            S (pySDC.step): The current step
+            size (int): Number of ranks
             time (float): The current time
             Tend (float): The final time
+            MS (list): Active steps
 
         Returns:
             None
         """
+
         # delete values that should not be available in the next step
         if self.params.no_storage:
             self.prev.t = np.array([None] * self.params.n)
@@ -351,13 +353,12 @@ def prepare_next_block_nonMPI(self, controller, MS, active_slots, time, Tend, **
 
         else:
             # decide where we need to restart to store everything up to that point
-            MS_active = [MS[i] for i in range(len(MS)) if i in active_slots]
-            restarts = [S.status.restart for S in MS_active]
-            restart_at = np.where(restarts)[0][0] if True in restarts else len(MS_active)
+            restarts = [S.status.restart for S in MS]
+            restart_at = np.where(restarts)[0][0] if True in restarts else len(MS)
 
             # store values in the current block that don't need restarting
-            if restart_at > 0:
-                [self.store_values(S) for S in MS_active[:restart_at]]
+            if restart_at > S.status.slot:
+                self.store_values(S)
 
         return None