Merge pull request #164 from brownbaerchen/rk-sweeper

Added Cash-Karp method to do embedded Runge-Kutta 5(4)

Author: pancetta

pySDC/implementations/convergence_controller_classes/adaptivity.py

@@ -67,7 +67,7 @@ def get_new_step_size(self, controller, S):
         '''
         raise NotImplementedError('Please implement a rule for updating the step size!')
 
-    def compute_optimatal_step_size(self, beta, dt, e_tol, e_est, order):
+    def compute_optimal_step_size(self, beta, dt, e_tol, e_est, order):
         '''
         Compute the optimal step size for the current step based on the order of the scheme.
         This function can be called from `get_new_step_size` for various implementations of adaptivity, but notably not
@@ -192,8 +192,8 @@ def get_new_step_size(self, controller, S):
             # compute next step size
             order = S.status.iter  # embedded error estimate is same order as time marching
             e_est = self.get_local_error_estimate(controller, S)
-            L.status.dt_new = self.compute_optimatal_step_size(self.params.beta, L.params.dt, self.params.e_tol, e_est,
-                                                               order)
+            L.status.dt_new = self.compute_optimal_step_size(self.params.beta, L.params.dt, self.params.e_tol, e_est,
+                                                             order)
             self.log(f'Adjusting step size from {L.params.dt:.2e} to {L.status.dt_new:.2e}', S)
 
         return None
@@ -212,6 +212,55 @@ def get_local_error_estimate(self, controller, S, **kwargs):
         return S.levels[0].status.error_embedded_estimate
 
 
+class AdaptivityRK(Adaptivity):
+    '''
+    Adaptivity for Runge-Kutta methods. Basically, we need to change the order in the step size update
+    '''
+    def check_parameters(self, controller, params, description):
+        '''
+        Check whether parameters are compatible with whatever assumptions went into the step size functions etc.
+        For adaptivity, we need to know the order of the scheme.
+
+        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:
+            bool: Whether the parameters are compatible
+            str: The error message
+        '''
+        if 'update_order' not in params.keys():
+            return False, 'Adaptivity needs an order for the update rule! Please set some up in \
+description[\'convergence_control_params\'][\'update_order\']!'
+
+        return super(AdaptivityRK, self).check_parameters(controller, params, description)
+
+    def get_new_step_size(self, controller, S):
+        '''
+        Determine a step size for the next step from an embedded estimate of the local error of the current step.
+
+        Args:
+            controller (pySDC.Controller): The controller
+            S (pySDC.Step): The current step
+
+        Returns:
+            None
+        '''
+        # check if we performed the desired amount of sweeps
+        if S.status.iter == S.params.maxiter:
+            L = S.levels[0]
+
+            # compute next step size
+            order = self.params.update_order
+            e_est = self.get_local_error_estimate(controller, S)
+            L.status.dt_new = self.compute_optimal_step_size(self.params.beta, L.params.dt, self.params.e_tol, e_est,
+                                                             order)
+            self.log(f'Adjusting step size from {L.params.dt:.2e} to {L.status.dt_new:.2e}', S)
+
+        return None
+
+
 class AdaptivityResidual(AdaptivityBase):
     '''
     Do adaptivity based on residual.

pySDC/implementations/convergence_controller_classes/estimate_embedded_error.py

@@ -3,6 +3,8 @@
 from pySDC.core.ConvergenceController import ConvergenceController, Pars
 from pySDC.implementations.convergence_controller_classes.store_uold import StoreUOld
 
+from pySDC.implementations.sweeper_classes.Runge_Kutta import RungeKutta
+
 
 class EstimateEmbeddedError(ConvergenceController):
     '''
@@ -14,7 +16,7 @@ class EstimateEmbeddedError(ConvergenceController):
 
     def setup(self, controller, params, description):
         '''
-        Add a default value for control order to the parameters
+        Add a default value for control order to the parameters and check if we are using a Runge-Kutta sweeper
 
         Args:
             controller (pySDC.Controller): The controller
@@ -24,11 +26,12 @@ def setup(self, controller, params, description):
         Returns:
             dict: Updated parameters
         '''
-        return {'control_order': -80, **params}
+        sweeper_type = 'RK' if RungeKutta in description['sweeper_class'].__bases__ else 'SDC'
+        return {'control_order': -80, 'sweeper_type': sweeper_type, **params}
 
     def dependencies(self, controller, description):
         '''
-        Load the convergence controller that stores the solution of the last sweep
+        Load the convergence controller that stores the solution of the last sweep unless we are doing Runge-Kutta
 
         Args:
             controller (pySDC.Controller): The controller
@@ -37,9 +40,32 @@ def dependencies(self, controller, description):
         Returns:
             None
         '''
-        controller.add_convergence_controller(StoreUOld, description=description)
+        if RungeKutta not in description['sweeper_class'].__bases__:
+            controller.add_convergence_controller(StoreUOld, description=description)
         return None
 
+    def estimate_embedded_error_serial(self, L):
+        '''
+        Estimate the serial embedded error, which may need to be modified for a parallel estimate.
+
+        Depending on the type of sweeper, the lower order solution is stored in a different place.
+
+        Args:
+            L (pySDC.level): The level
+
+        Returns:
+            dtype_u: The embedded error estimate
+        '''
+        if self.params.sweeper_type == 'RK':
+            # lower order solution is stored in the second to last entry of L.u
+            return abs(L.u[-2] - L.u[-1])
+        elif self.params.sweeper_type == 'SDC':
+            # order rises by one between sweeps, making this so ridiculously easy
+            return abs(L.uold[-1] - L.u[-1])
+        else:
+            raise NotImplementedError(f'Don\'t know how to estimate embedded error for sweeper type \
+{self.params.sweeper_type}')
+
 
 class EstimateEmbeddedErrorNonMPI(EstimateEmbeddedError):
 
@@ -84,10 +110,9 @@ def post_iteration_processing(self, controller, S):
             raise NotImplementedError('Embedded error estimate only works for serial multi-level or parallel single \
 level')
 
-        if S.status.iter > 1:
+        if S.status.iter > 1 or self.params.sweeper_type == 'RK':
             for L in S.levels:
-                # order rises by one between sweeps, making this so ridiculously easy
-                temp = abs(L.uold[-1] - L.u[-1])
+                temp = self.estimate_embedded_error_serial(L)
                 L.status.error_embedded_estimate = max([abs(temp - self.buffers.e_em_last), np.finfo(float).eps])
 
             self.buffers.e_em_last = temp * 1.

pySDC/implementations/sweeper_classes/Runge_Kutta.py

@@ -9,7 +9,7 @@
 class ButcherTableau(object):
     def __init__(self, weights, nodes, matrix):
         """
-        Initialization routine for an collocation object
+        Initialization routine to get a quadrature matrix out of a Butcher tableau
 
         Args:
             weights (numpy.ndarray): Butcher tableau weights
@@ -61,6 +61,64 @@ def __init__(self, weights, nodes, matrix):
         self.implicit = any([matrix[i, i] != 0 for i in range(self.num_nodes - 1)])
 
 
+class ButcherTableauEmbedded(object):
+    def __init__(self, weights, nodes, matrix):
+        """
+        Initialization routine to get a quadrature matrix out of a Butcher tableau for embedded RK methods.
+
+        Be aware that the method that generates the final solution should be in the first row of the weights matrix.
+
+        Args:
+            weights (numpy.ndarray): Butcher tableau weights
+            nodes (numpy.ndarray): Butcher tableau nodes
+            matrix (numpy.ndarray): Butcher tableau entries
+        """
+        # check if the arguments have the correct form
+        if type(matrix) != np.ndarray:
+            raise ParameterError('Runge-Kutta matrix needs to be supplied as  a numpy array!')
+        elif len(np.unique(matrix.shape)) != 1 or len(matrix.shape) != 2:
+            raise ParameterError('Runge-Kutta matrix needs to be a square 2D numpy array!')
+
+        if type(weights) != np.ndarray:
+            raise ParameterError('Weights need to be supplied as a numpy array!')
+        elif len(weights.shape) != 2:
+            raise ParameterError(f'Incompatible dimension of weights! Need 2, got {len(weights.shape)}')
+        elif len(weights[0]) != matrix.shape[0]:
+            raise ParameterError(f'Incompatible number of weights! Need {matrix.shape[0]}, got {len(weights[0])}')
+
+        if type(nodes) != np.ndarray:
+            raise ParameterError('Nodes need to be supplied as a numpy array!')
+        elif len(nodes.shape) != 1:
+            raise ParameterError(f'Incompatible dimension of nodes! Need 1, got {len(nodes.shape)}')
+        elif len(nodes) != matrix.shape[0]:
+            raise ParameterError(f'Incompatible number of nodes! Need {matrix.shape[0]}, got {len(nodes)}')
+
+        # Set number of nodes, left and right interval boundaries
+        self.num_nodes = matrix.shape[0] + 2
+        self.tleft = 0.
+        self.tright = 1.
+
+        self.nodes = np.append(np.append([0], nodes), [1, 1])
+        self.weights = weights
+        self.Qmat = np.zeros([self.num_nodes + 1, self.num_nodes + 1])
+        self.Qmat[1:-2, 1:-2] = matrix
+        self.Qmat[-1, 1:-2] = weights[0]  # this is for computing the higher order solution
+        self.Qmat[-2, 1:-2] = weights[1]  # this is for computing the lower order solution
+
+        self.left_is_node = True
+        self.right_is_node = self.nodes[-1] == self.tright
+
+        # compute distances between the nodes
+        if self.num_nodes > 1:
+            self.delta_m = self.nodes[1:] - self.nodes[:-1]
+        else:
+            self.delta_m = np.zeros(1)
+        self.delta_m[0] = self.nodes[0] - self.tleft
+
+        # check if the RK scheme is implicit
+        self.implicit = any([matrix[i, i] != 0 for i in range(self.num_nodes - 2)])
+
+
 class RungeKutta(generic_implicit):
     """
     Runge-Kutta scheme that fits the interface of a sweeper.
@@ -215,3 +273,34 @@ def __init__(self, params):
         matrix[3, 2] = 1.
         params['butcher_tableau'] = ButcherTableau(weights, nodes, matrix)
         super(RK4, self).__init__(params)
+
+
+class Heun_Euler(RungeKutta):
+    '''
+    Second order explicit embedded Runge-Kutta
+    '''
+    def __init__(self, params):
+        nodes = np.array([0, 1])
+        weights = np.array([[0.5, 0.5], [1, 0]])
+        matrix = np.zeros((2, 2))
+        matrix[1, 0] = 1
+        params['butcher_tableau'] = ButcherTableauEmbedded(weights, nodes, matrix)
+        super(Heun_Euler, self).__init__(params)
+
+
+class Cash_Karp(RungeKutta):
+    '''
+    Fifth order explicit embedded Runge-Kutta
+    '''
+    def __init__(self, params):
+        nodes = np.array([0, 0.2, 0.3, 0.6, 1., 7. / 8.])
+        weights = np.array([[37. / 378., 0., 250. / 621., 125. / 594., 0., 512. / 1771.],
+                            [2825. / 27648., 0., 18575. / 48384., 13525. / 55296., 277. / 14336., 1. / 4.]])
+        matrix = np.zeros((6, 6))
+        matrix[1, 0] = 1. / 5.
+        matrix[2, :2] = [3. / 40., 9. / 40.]
+        matrix[3, :3] = [0.3, -0.9, 1.2]
+        matrix[4, :4] = [-11. / 54., 5. / 2., -70. / 27., 35. / 27.]
+        matrix[5, :5] = [1631. / 55296., 175. / 512., 575. / 13824., 44275. / 110592., 253. / 4096.]
+        params['butcher_tableau'] = ButcherTableauEmbedded(weights, nodes, matrix)
+        super(Cash_Karp, self).__init__(params)

pySDC/playgrounds/Runge_Kutta/Runge-Kutta-Methods.ipynb

@@ -141,16 +141,20 @@ def plot(res, ax, k):
     color = plt.rcParams['axes.prop_cycle'].by_key()['color'][k - 2]
 
     for i in range(len(keys)):
+        if all([me == 0. for me in res[keys[i]]]):
+            continue
         order = get_accuracy_order(res, key=keys[i])
         if keys[i] == 'e_embedded':
             label = rf'$k={{{np.mean(order):.2f}}}$'
-            assert np.isclose(np.mean(order), k, atol=3e-1), f'Expected embedded error estimate to have order {k} \
+            assert np.isclose(np.mean(order), k, atol=4e-1), f'Expected embedded error estimate to have order {k} \
 but got {np.mean(order):.2f}'
 
         elif keys[i] == 'e_extrapolated':
             label = None
             assert np.isclose(np.mean(order), k + 1, rtol=3e-1), f'Expected extrapolation error estimate to have order \
 {k+1} but got {np.mean(order):.2f}'
+        else:
+            label = None
         ax.loglog(res['dt'], res[keys[i]], color=color, ls=ls[i], label=label)
 
     ax.set_xlabel(r'$\Delta t$')
@@ -198,11 +202,12 @@ def plot_orders(ax, ks, serial, Tend_fixed=None, custom_description=None, prob=r
         plot_order(res, ax, k)
 
 
-def plot_all_errors(ax, ks, serial, Tend_fixed=None, custom_description=None, prob=run_piline):
+def plot_all_errors(ax, ks, serial, Tend_fixed=None, custom_description=None, prob=run_piline, dt_list=None,
+                    custom_controller_params=None):
     for i in range(len(ks)):
         k = ks[i]
         res = multiple_runs(k=k, serial=serial, Tend_fixed=Tend_fixed, custom_description=custom_description,
-                            prob=prob)
+                            prob=prob, dt_list=dt_list, custom_controller_params=custom_controller_params)
 
         # visualize results
         plot(res, ax, k)