Fixes and refactor for RKN

Author: brownbaerchen

pySDC/implementations/sweeper_classes/Runge_Kutta.py

@@ -267,26 +267,28 @@ def compute_end_point(self):
         """
         In this Runge-Kutta implementation, the solution to the step is always stored in the last node
         """
-        if self.level.f[1] is None:
-            self.level.uend = self.level.u[0]
+        lvl = self.level
+
+        if lvl.f[1] is None:
+            lvl.uend = lvl.prob.dtype_u(lvl.u[0])
             if type(self.coll) == ButcherTableauEmbedded:
-                self.u_secondary = self.level.u[0].copy()
+                self.u_secondary = lvl.prob.dtype_u(lvl.u[0])
         elif self.coll.globally_stiffly_accurate:
-            self.level.uend = self.level.u[-1]
+            lvl.uend = lvl.prob.dtype_u(lvl.u[-1])
             if type(self.coll) == ButcherTableauEmbedded:
-                self.u_secondary = self.level.u[0].copy()
-                for w2, k in zip(self.coll.weights[1], self.level.f[1:]):
-                    self.u_secondary += self.level.dt * w2 * k
+                self.u_secondary = lvl.prob.dtype_u(lvl.u[0])
+                for w2, k in zip(self.coll.weights[1], lvl.f[1:]):
+                    self.u_secondary += lvl.dt * w2 * k
         else:
-            self.level.uend = self.level.u[0].copy()
+            lvl.uend = lvl.prob.dtype_u(lvl.u[0])
             if type(self.coll) == ButcherTableau:
-                for w, k in zip(self.coll.weights, self.level.f[1:]):
-                    self.level.uend += self.level.dt * w * k
+                for w, k in zip(self.coll.weights, lvl.f[1:]):
+                    lvl.uend += lvl.dt * w * k
             elif type(self.coll) == ButcherTableauEmbedded:
-                self.u_secondary = self.level.u[0].copy()
-                for w1, w2, k in zip(self.coll.weights[0], self.coll.weights[1], self.level.f[1:]):
-                    self.level.uend += self.level.dt * w1 * k
-                    self.u_secondary += self.level.dt * w2 * k
+                self.u_secondary = lvl.prob.dtype_u(lvl.u[0])
+                for w1, w2, k in zip(self.coll.weights[0], self.coll.weights[1], lvl.f[1:]):
+                    lvl.uend += lvl.dt * w1 * k
+                    self.u_secondary += lvl.dt * w2 * k
 
     @property
     def level(self):
@@ -444,32 +446,34 @@ def compute_end_point(self):
         """
         In this Runge-Kutta implementation, the solution to the step is always stored in the last node
         """
-        if self.level.f[1] is None:
-            self.level.uend = self.level.u[0]
+        lvl = self.level
+
+        if lvl.f[1] is None:
+            lvl.uend = lvl.prob.dtype_u(lvl.u[0])
             if type(self.coll) == ButcherTableauEmbedded:
-                self.u_secondary = self.level.u[0].copy()
+                self.u_secondary = lvl.prob.dtype_u(lvl.u[0])
         elif self.coll.globally_stiffly_accurate and self.coll_explicit.globally_stiffly_accurate:
-            self.level.uend = self.level.u[-1]
+            lvl.uend = lvl.u[-1]
             if type(self.coll) == ButcherTableauEmbedded:
-                self.u_secondary = self.level.u[0].copy()
-                for w2, w2E, k in zip(self.coll.weights[1], self.coll_explicit.weights[1], self.level.f[1:]):
-                    self.u_secondary += self.level.dt * (w2 * k.impl + w2E * k.expl)
+                self.u_secondary = lvl.prob.dtype_u(lvl.u[0])
+                for w2, w2E, k in zip(self.coll.weights[1], self.coll_explicit.weights[1], lvl.f[1:]):
+                    self.u_secondary += lvl.dt * (w2 * k.impl + w2E * k.expl)
         else:
-            self.level.uend = self.level.u[0].copy()
+            lvl.uend = lvl.prob.dtype_u(lvl.u[0])
             if type(self.coll) == ButcherTableau:
-                for w, wE, k in zip(self.coll.weights, self.coll_explicit.weights, self.level.f[1:]):
-                    self.level.uend += self.level.dt * (w * k.impl + wE * k.expl)
+                for w, wE, k in zip(self.coll.weights, self.coll_explicit.weights, lvl.f[1:]):
+                    lvl.uend += lvl.dt * (w * k.impl + wE * k.expl)
             elif type(self.coll) == ButcherTableauEmbedded:
-                self.u_secondary = self.level.u[0].copy()
+                self.u_secondary = lvl.u[0].copy()
                 for w1, w2, w1E, w2E, k in zip(
                     self.coll.weights[0],
                     self.coll.weights[1],
                     self.coll_explicit.weights[0],
                     self.coll_explicit.weights[1],
-                    self.level.f[1:],
+                    lvl.f[1:],
                 ):
-                    self.level.uend += self.level.dt * (w1 * k.impl + w1E * k.expl)
-                    self.u_secondary += self.level.dt * (w2 * k.impl + w2E * k.expl)
+                    lvl.uend += lvl.dt * (w1 * k.impl + w1E * k.expl)
+                    self.u_secondary += lvl.dt * (w2 * k.impl + w2E * k.expl)
 
 
 class ForwardEuler(RungeKutta):

pySDC/implementations/sweeper_classes/Runge_Kutta_Nystrom.py

@@ -4,11 +4,74 @@
 from pySDC.core.sweeper import Sweeper, _Pars
 from pySDC.core.errors import ParameterError
 from pySDC.implementations.datatype_classes.particles import particles, fields, acceleration
-from pySDC.implementations.sweeper_classes.Runge_Kutta import ButcherTableau
-from copy import deepcopy
 from pySDC.implementations.sweeper_classes.Runge_Kutta import RungeKutta
 
 
+class ButcherTableauNoCollUpdate(object):
+    """Version of Butcher Tableau that does not need a collocation update because the weights are put in the last line of Q"""
+
+    def __init__(self, weights, nodes, matrix):
+        """
+        Initialization routine to get a quadrature matrix out of a Butcher tableau
+
+        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) != 1:
+            raise ParameterError(f'Incompatible dimension of weights! Need 1, got {len(weights.shape)}')
+        elif len(weights) != matrix.shape[0]:
+            raise ParameterError(f'Incompatible number of weights! Need {matrix.shape[0]}, got {len(weights)}')
+
+        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)}')
+
+        self.globally_stiffly_accurate = np.allclose(matrix[-1], weights)
+
+        self.tleft = 0.0
+        self.tright = 1.0
+        self.num_solution_stages = 0 if self.globally_stiffly_accurate else 1
+        self.num_nodes = matrix.shape[0] + self.num_solution_stages
+        self.weights = weights
+
+        if self.globally_stiffly_accurate:
+            # For globally stiffly accurate methods, the last row of the Butcher tableau is the same as the weights.
+            self.nodes = np.append([0], nodes)
+            self.Qmat = np.zeros([self.num_nodes + 1, self.num_nodes + 1])
+            self.Qmat[1:, 1:] = matrix
+        else:
+            self.nodes = np.append(np.append([0], nodes), [1])
+            self.Qmat = np.zeros([self.num_nodes + 1, self.num_nodes + 1])
+            self.Qmat[1:-1, 1:-1] = matrix
+            self.Qmat[-1, 1:-1] = weights  # this is for computing the solution to the step from the previous stages
+
+        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 - self.num_solution_stages))
+
+
 class RungeKuttaNystrom(RungeKutta):
     """
     Runge-Kutta scheme that fits the interface of a sweeper.
@@ -34,56 +97,31 @@ class RungeKuttaNystrom(RungeKutta):
     All of these variables are either determined by the RK rule, or are not part of an RK scheme.
 
     Attribues:
-        butcher_tableau (ButcherTableau): Butcher tableau for the Runge-Kutta scheme that you want
+        butcher_tableau (ButcherTableauNoCollUpdate): Butcher tableau for the Runge-Kutta scheme that you want
     """
 
+    ButcherTableauClass = ButcherTableauNoCollUpdate
+    weights_bar = None
+    matrix_bar = None
+
     def __init__(self, params):
         """
         Initialization routine for the custom sweeper
 
         Args:
             params: parameters for the sweeper
         """
-        # set up logger
-        self.logger = logging.getLogger('sweeper')
-
-        essential_keys = ['butcher_tableau']
-        for key in essential_keys:
-            if key not in params:
-                msg = 'need %s to instantiate step, only got %s' % (key, str(params.keys()))
-                self.logger.error(msg)
-                raise ParameterError(msg)
-
-        # check if some parameters are set which only apply to actual sweepers
-        for key in ['initial_guess', 'collocation_class', 'num_nodes']:
-            if key in params:
-                self.logger.warning(f'"{key}" will be ignored by Runge-Kutta sweeper')
-
-        # set parameters to their actual values
-        params['initial_guess'] = 'zero'
-        params['collocation_class'] = type(params['butcher_tableau'])
-        params['num_nodes'] = params['butcher_tableau'].num_nodes
-
-        # disable residual computation by default
-        params['skip_residual_computation'] = params.get(
-            'skip_residual_computation', ('IT_CHECK', 'IT_FINE', 'IT_COARSE', 'IT_UP', 'IT_DOWN')
-        )
-
-        self.params = _Pars(params)
-
-        self.coll = params['butcher_tableau']
-        self.coll_bar = params['butcher_tableau_bar']
-
-        # This will be set as soon as the sweeper is instantiated at the level
-        self.__level = None
-
-        self.parallelizable = False
-        self.QI = self.coll.Qmat
+        super().__init__(params)
+        self.coll_bar = self.get_Butcher_tableau_bar()
         self.Qx = self.coll_bar.Qmat
 
+    @classmethod
+    def get_Butcher_tableau_bar(cls):
+        return cls.ButcherTableauClass(cls.weights_bar, cls.nodes, cls.matrix_bar)
+
     def get_full_f(self, f):
         """
-        Test the right hand side funtion is the correct type
+        Test the right hand side function is the correct type
 
         Args:
             f (dtype_f): Right hand side at a single node
@@ -118,7 +156,7 @@ def update_nodes(self):
 
         for m in range(0, M):
             # build rhs, consisting of the known values from above and new values from previous nodes (at k+1)
-            rhs = deepcopy(L.u[0])
+            rhs = P.dtype_u(L.u[0])
             rhs.pos += L.dt * self.coll.nodes[m + 1] * L.u[0].vel
 
             for j in range(1, m + 1):
@@ -147,7 +185,7 @@ def update_nodes(self):
             if self.coll.implicit:
                 # That is why it only works for the Velocity-Verlet scheme
                 L.f[0] = P.eval_f(L.u[0], L.time)
-                L.f[m + 1] = deepcopy(L.f[0])
+                L.f[m + 1] = P.dtype_f(L.f[0])
             else:
                 if m != self.coll.num_nodes - 1:
                     L.f[m + 1] = P.eval_f(L.u[m + 1], L.time + L.dt * self.coll.nodes[m])
@@ -173,23 +211,18 @@ class RKN(RungeKuttaNystrom):
     Chapter: II.14 Numerical methods for Second order differential equations
     """
 
-    def __init__(self, params):
-        nodes = np.array([0.0, 0.5, 0.5, 1])
-        weights = np.array([1.0, 2.0, 2.0, 1.0]) / 6.0
-        matrix = np.zeros([4, 4])
-        matrix[1, 0] = 0.5
-        matrix[2, 1] = 0.5
-        matrix[3, 2] = 1.0
-
-        weights_bar = np.array([1.0, 1.0, 1.0, 0]) / 6.0
-        matrix_bar = np.zeros([4, 4])
-        matrix_bar[1, 0] = 1 / 8
-        matrix_bar[2, 0] = 1 / 8
-        matrix_bar[3, 2] = 1 / 2
-        params['butcher_tableau'] = ButcherTableau(weights, nodes, matrix)
-        params['butcher_tableau_bar'] = ButcherTableau(weights_bar, nodes, matrix_bar)
+    nodes = np.array([0.0, 0.5, 0.5, 1])
+    weights = np.array([1.0, 2.0, 2.0, 1.0]) / 6.0
+    matrix = np.zeros([4, 4])
+    matrix[1, 0] = 0.5
+    matrix[2, 1] = 0.5
+    matrix[3, 2] = 1.0
 
-        super(RKN, self).__init__(params)
+    weights_bar = np.array([1.0, 1.0, 1.0, 0]) / 6.0
+    matrix_bar = np.zeros([4, 4])
+    matrix_bar[1, 0] = 1 / 8
+    matrix_bar[2, 0] = 1 / 8
+    matrix_bar[3, 2] = 1 / 2
 
 
 class Velocity_Verlet(RungeKuttaNystrom):
@@ -198,15 +231,9 @@ class Velocity_Verlet(RungeKuttaNystrom):
     https://de.wikipedia.org/wiki/Verlet-Algorithmus
     """
 
-    def __init__(self, params):
-        nodes = np.array([1.0, 1.0])
-        weights = np.array([1 / 2, 0])
-        matrix = np.zeros([2, 2])
-        matrix[1, 1] = 1
-        weights_bar = np.array([1 / 2, 0])
-        matrix_bar = np.zeros([2, 2])
-        params['butcher_tableau'] = ButcherTableau(weights, nodes, matrix)
-        params['butcher_tableau_bar'] = ButcherTableau(weights_bar, nodes, matrix_bar)
-        params['Velocity_verlet'] = True
-
-        super(Velocity_Verlet, self).__init__(params)
+    nodes = np.array([1.0, 1.0])
+    weights = np.array([1 / 2, 0])
+    matrix = np.zeros([2, 2])
+    matrix[1, 1] = 1
+    weights_bar = np.array([1 / 2, 0])
+    matrix_bar = np.zeros([2, 2])

pySDC/tests/test_sweepers/test_Runge_Kutta_sweeper.py

@@ -120,12 +120,16 @@ def single_run(sweeper_name, dt, lambdas, use_RK_sweeper=True, Tend=None, useGPU
         sweeper = controller.MS[0].levels[0].sweep
         sweeper.QI = rk_sweeper.get_Q_matrix()
         sweeper.coll = rk_sweeper.get_Butcher_tableau()
+        _compute_end_point = type(sweeper).compute_end_point
         type(sweeper).compute_end_point = rk_sweeper.compute_end_point
 
     prob = controller.MS[0].levels[0].prob
     ic = prob.u_exact(0)
     u_end, stats = controller.run(ic, 0.0, 5 * dt if Tend is None else Tend)
 
+    if not use_RK_sweeper:
+        type(sweeper).compute_end_point = _compute_end_point
+
     return stats, ic, controller