Merge pull request #281 from brownbaerchen/interpolate_restarts

Interpolate after restarts

Author: pancetta

pySDC/core/Controller.py

@@ -32,13 +32,14 @@ class controller(object):
     Base abstract controller class
     """
 
-    def __init__(self, controller_params, description):
+    def __init__(self, controller_params, description, useMPI=None):
         """
         Initialization routine for the base controller
 
         Args:
             controller_params (dict): parameter set for the controller and the steps
         """
+        self.useMPI = useMPI
 
         # check if we have a hook on this list. If not, use default class.
         self.__hooks = []
@@ -288,7 +289,7 @@ def add_convergence_controller(self, convergence_controller, description, params
             None
         '''
         # check if we passed any sort of special params
-        params = {} if params is None else params
+        params = {**({} if params is None else params), 'useMPI': self.useMPI}
 
         # check if we already have the convergence controller or if we want to have it multiple times
         if convergence_controller not in [type(me) for me in self.convergence_controllers] or allow_double:

pySDC/core/ConvergenceController.py

@@ -6,6 +6,7 @@
 class Pars(FrozenClass):
     def __init__(self, params):
         self.control_order = 0  # integer that determines the order in which the convergence controllers are called
+        self.useMPI = None  # depends on the controller
 
         for k, v in params.items():
             setattr(self, k, v)

pySDC/helpers/plot_helper.py

@@ -25,7 +25,7 @@ def figsize(textwidth, scale, ratio):
     return fig_size
 
 
-def figsize_by_journal(journal, scale, ratio):  # pragma no cover
+def figsize_by_journal(journal, scale, ratio):  # pragma: no cover
     """
     Get figsize for specific journal. If you supply a text height, we will rescale the figure to fit on the page instead
     of the parameters supplied.

pySDC/implementations/controller_classes/controller_MPI.py

@@ -25,7 +25,7 @@ def __init__(self, controller_params, description, comm):
         """
 
         # call parent's initialization routine
-        super(controller_MPI, self).__init__(controller_params, description)
+        super().__init__(controller_params, description, useMPI=True)
 
         # create single step per processor
         self.S = step(description)

pySDC/implementations/controller_classes/controller_nonMPI.py

@@ -30,7 +30,7 @@ def __init__(self, num_procs, controller_params, description):
             raise ControllerError('predict flag is ignored, use predict_type instead')
 
         # call parent's initialization routine
-        super(controller_nonMPI, self).__init__(controller_params, description)
+        super().__init__(controller_params, description, useMPI=False)
 
         self.MS = [stepclass.step(description)]
 

pySDC/implementations/convergence_controller_classes/adaptivity.py

@@ -6,7 +6,6 @@
 from pySDC.implementations.convergence_controller_classes.basic_restarting import (
     BasicRestartingNonMPI,
 )
-from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
 from pySDC.implementations.hooks.log_step_size import LogStepSize
 
 
@@ -175,7 +174,7 @@ def dependencies(self, controller, description, **kwargs):
         super(Adaptivity, self).dependencies(controller, description)
 
         controller.add_convergence_controller(
-            EstimateEmbeddedError.get_implementation("nonMPI" if type(controller) == controller_nonMPI else "MPI"),
+            EstimateEmbeddedError.get_implementation("nonMPI" if not self.params.useMPI else "MPI"),
             description=description,
         )
 

pySDC/implementations/convergence_controller_classes/check_convergence.py

@@ -45,10 +45,9 @@ def dependencies(self, controller, description, **kwargs):
             from pySDC.implementations.convergence_controller_classes.estimate_embedded_error import (
                 EstimateEmbeddedError,
             )
-            from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
 
             controller.add_convergence_controller(
-                EstimateEmbeddedError.get_implementation("nonMPI" if type(controller) == controller_nonMPI else "MPI"),
+                EstimateEmbeddedError.get_implementation("nonMPI" if not self.params.useMPI else "MPI"),
                 description=description,
             )
 

pySDC/implementations/convergence_controller_classes/estimate_contraction_factor.py

@@ -1,7 +1,6 @@
 import numpy as np
 
 from pySDC.core.ConvergenceController import ConvergenceController
-from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
 from pySDC.implementations.convergence_controller_classes.estimate_embedded_error import EstimateEmbeddedError
 
 
@@ -36,7 +35,7 @@ def dependencies(self, controller, description, **kwargs):
             None
         """
         controller.add_convergence_controller(
-            EstimateEmbeddedError.get_implementation("nonMPI" if type(controller) == controller_nonMPI else "MPI"),
+            EstimateEmbeddedError.get_implementation("nonMPI" if not self.params.useMPI else "MPI"),
             description=description,
         )
 

pySDC/implementations/convergence_controller_classes/hotrod.py

@@ -1,7 +1,6 @@
 import numpy as np
 
 from pySDC.core.ConvergenceController import ConvergenceController
-from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
 from pySDC.implementations.convergence_controller_classes.estimate_extrapolation_error import (
     EstimateExtrapolationErrorNonMPI,
 )
@@ -49,7 +48,7 @@ def dependencies(self, controller, description, **kwargs):
         Returns:
             None
         """
-        if type(controller) == controller_nonMPI:
+        if not self.params.useMPI:
             from pySDC.implementations.convergence_controller_classes.estimate_embedded_error import (
                 EstimateEmbeddedErrorNonMPI,
             )

pySDC/implementations/convergence_controller_classes/interpolate_between_restarts.py

@@ -0,0 +1,92 @@
+import numpy as np
+from pySDC.core.ConvergenceController import ConvergenceController, Status
+from pySDC.core.Lagrange import LagrangeApproximation
+from pySDC.core.Collocation import CollBase
+
+
+class InterpolateBetweenRestarts(ConvergenceController):
+    """
+    Interpolate the solution and right hand side to the new set of collocation nodes after a restart.
+    The idea is that when you adjust the step size between restarts, you already know what the new quadrature method
+    is going to be and possibly interpolating the current iterate to these results in a better initial guess than
+    spreading the initial conditions or whatever you usually like to do.
+    """
+
+    def setup(self, controller, params, description, **kwargs):
+        """
+        Store the initial guess used in the sweeper when no restart has happened
+
+        Args:
+            controller (pySDC.Controller.controller): The controller
+            params (dict): Parameters for the convergence controller
+            description (dict): The description object used to instantiate the controller
+        """
+        defaults = {
+            'control_order': 50,
+        }
+        return {**defaults, **super().setup(controller, params, description, **kwargs)}
+
+    def setup_status_variables(self, controller, **kwargs):
+        """
+        Add variables to the sweeper containing the interpolated solution and right hand side.
+
+        Args:
+            controller (pySDC.Controller.controller): The controller
+        """
+        self.status = Status(['u_inter', 'f_inter', 'perform_interpolation'])
+
+        self.status.u_inter = []
+        self.status.f_inter = []
+        self.status.perform_interpolation = False
+
+    def post_spread_processing(self, controller, step, **kwargs):
+        """
+        Spread the interpolated values to the collocation nodes. This overrides whatever the sweeper uses for prediction.
+
+        Args:
+            controller (pySDC.Controller.controller): The controller
+            step (pySDC.Step.step): The current step
+        """
+        if self.status.perform_interpolation:
+            for i in range(len(step.levels)):
+                level = step.levels[i]
+                for m in range(len(level.u)):
+                    level.u[m][:] = self.status.u_inter[i][m][:]
+                    level.f[m][:] = self.status.f_inter[i][m][:]
+
+            # reset the status variables
+            self.status.perform_interpolation = False
+            self.status.u_inter = []
+            self.status.f_inter = []
+
+    def post_iteration_processing(self, controller, step, **kwargs):
+        """
+        Interpolate the solution and right hand sides and store them in the sweeper, where they will be distributed
+        accordingly in the prediction step.
+
+        This function is called after every iteration instead of just after the step because we might choose to stop
+        iterating as soon as we have decided to restart. If we let the step continue to iterate, this is not the most
+        efficient implementation and you may choose to write a different convergence controller.
+
+        The interpolation is based on Thibaut's magic.
+
+        Args:
+            controller (pySDC.Controller): The controller
+            step (pySDC.Step.step): The current step
+        """
+        if step.status.restart and all([level.status.dt_new for level in step.levels]):
+            for level in step.levels:
+                nodes_old = level.sweep.coll.nodes.copy()
+                nodes_new = level.sweep.coll.nodes.copy() * level.status.dt_new / level.params.dt
+
+                interpolator = LagrangeApproximation(points=np.append(0, nodes_old))
+                self.status.u_inter += [(interpolator.getInterpolationMatrix(np.append(0, nodes_new)) @ level.u[:])[:]]
+                self.status.f_inter += [(interpolator.getInterpolationMatrix(np.append(0, nodes_new)) @ level.f[:])[:]]
+
+                self.status.perform_interpolation = True
+
+                self.log(
+                    f'Interpolating before restart from dt={level.params.dt:.2e} to dt={level.status.dt_new:.2e}', step
+                )
+        else:
+            self.status.perform_interpolation = False