Refactor of RK implementation (#499)

* Globally stiffly accurate ARK methods

* Refactor and fix

* Fixes and refactor for RKN

pySDC/implementations/convergence_controller_classes/estimate_embedded_error.py

@@ -93,10 +93,10 @@ def estimate_embedded_error_serial(self, L):
             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])
+            L.sweep.compute_end_point()
+            return abs(L.uend - L.sweep.u_secondary)
         elif self.params.sweeper_type == "SDC":
-            # order rises by one between sweeps, making this so ridiculously easy
+            # order rises by one between sweeps
             return abs(L.uold[-1] - L.u[-1])
         elif self.params.sweeper_type == 'MPI':
             comm = L.sweep.comm

pySDC/implementations/sweeper_classes/Runge_Kutta.py

@@ -17,44 +17,16 @@ def __init__(self, weights, nodes, matrix):
             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.check_method(weights, nodes, matrix)
 
         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.num_nodes = matrix.shape[0]
         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.nodes = np.append([0], nodes)
+        self.Qmat = np.zeros([self.num_nodes + 1, self.num_nodes + 1])
+        self.Qmat[1:, 1:] = matrix
 
         self.left_is_node = True
         self.right_is_node = self.nodes[-1] == self.tright
@@ -67,66 +39,58 @@ def __init__(self, weights, nodes, matrix):
         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))
+        self.implicit = any(matrix[i, i] != 0 for i in range(self.num_nodes))
 
-
-class ButcherTableauEmbedded(object):
-    def __init__(self, weights, nodes, matrix):
+    def check_method(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 that the method is entered in the correct format
         """
-        # 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_solution_stages = 2
-        self.num_nodes = matrix.shape[0] + self.num_solution_stages
-        self.tleft = 0.0
-        self.tright = 1.0
+        self.check_weights(weights, nodes, matrix)
 
-        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
+    def check_weights(self, weights, nodes, matrix):
+        """
+        Check that the weights of the method are entered in the correct format
+        """
+        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)}')
 
-        self.left_is_node = True
-        self.right_is_node = self.nodes[-1] == self.tright
+    @property
+    def globally_stiffly_accurate(self):
+        return np.allclose(self.Qmat[-1, 1:], self.weights)
 
-        # 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 ButcherTableauEmbedded(ButcherTableau):
+
+    def check_weights(self, weights, nodes, matrix):
+        """
+        Check that the weights of the method are entered in the correct format
+        """
+        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])}')
+
+    @property
+    def globally_stiffly_accurate(self):
+        return np.allclose(self.Qmat[-1, 1:], self.weights[0])
 
 
 class RungeKutta(Sweeper):
@@ -292,8 +256,7 @@ def update_nodes(self):
                 lvl.u[m + 1][:] = rhs[:]
 
             # update function values (we don't usually need to evaluate the RHS at the solution of the step)
-            if m < M - self.coll.num_solution_stages or self.params.eval_rhs_at_right_boundary:
-                lvl.f[m + 1] = prob.eval_f(lvl.u[m + 1], lvl.time + lvl.dt * self.coll.nodes[m + 1])
+            lvl.f[m + 1] = prob.eval_f(lvl.u[m + 1], lvl.time + lvl.dt * self.coll.nodes[m + 1])
 
         # indicate presence of new values at this level
         lvl.status.updated = True
@@ -304,7 +267,28 @@ def compute_end_point(self):
         """
         In this Runge-Kutta implementation, the solution to the step is always stored in the last node
         """
-        self.level.uend = self.level.u[-1]
+        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 = lvl.prob.dtype_u(lvl.u[0])
+        elif self.coll.globally_stiffly_accurate:
+            lvl.uend = lvl.prob.dtype_u(lvl.u[-1])
+            if type(self.coll) == ButcherTableauEmbedded:
+                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:
+            lvl.uend = lvl.prob.dtype_u(lvl.u[0])
+            if type(self.coll) == ButcherTableau:
+                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 = 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):
@@ -356,6 +340,7 @@ class RungeKuttaIMEX(RungeKutta):
     """
 
     matrix_explicit = None
+    weights_explicit = None
     ButcherTableauClass_explicit = ButcherTableau
 
     def __init__(self, params):
@@ -366,6 +351,7 @@ def __init__(self, params):
             params: parameters for the sweeper
         """
         super().__init__(params)
+        type(self).weights_explicit = self.weights if self.weights_explicit is None else self.weights_explicit
         self.coll_explicit = self.get_Butcher_tableau_explicit()
         self.QE = self.coll_explicit.Qmat
 
@@ -388,7 +374,7 @@ def predict(self):
 
     @classmethod
     def get_Butcher_tableau_explicit(cls):
-        return cls.ButcherTableauClass_explicit(cls.weights, cls.nodes, cls.matrix_explicit)
+        return cls.ButcherTableauClass_explicit(cls.weights_explicit, cls.nodes, cls.matrix_explicit)
 
     def integrate(self):
         """
@@ -448,15 +434,47 @@ def update_nodes(self):
             else:
                 lvl.u[m + 1][:] = rhs[:]
 
-            # update function values (we don't usually need to evaluate the RHS at the solution of the step)
-            if m < M - self.coll.num_solution_stages or self.params.eval_rhs_at_right_boundary:
-                lvl.f[m + 1] = prob.eval_f(lvl.u[m + 1], lvl.time + lvl.dt * self.coll.nodes[m + 1])
+            # update function values
+            lvl.f[m + 1] = prob.eval_f(lvl.u[m + 1], lvl.time + lvl.dt * self.coll.nodes[m + 1])
 
         # indicate presence of new values at this level
         lvl.status.updated = True
 
         return None
 
+    def compute_end_point(self):
+        """
+        In this Runge-Kutta implementation, the solution to the step is always stored in the last node
+        """
+        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 = lvl.prob.dtype_u(lvl.u[0])
+        elif self.coll.globally_stiffly_accurate and self.coll_explicit.globally_stiffly_accurate:
+            lvl.uend = lvl.u[-1]
+            if type(self.coll) == ButcherTableauEmbedded:
+                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:
+            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, lvl.f[1:]):
+                    lvl.uend += lvl.dt * (w * k.impl + wE * k.expl)
+            elif type(self.coll) == ButcherTableauEmbedded:
+                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],
+                    lvl.f[1:],
+                ):
+                    lvl.uend += lvl.dt * (w1 * k.impl + w1E * k.expl)
+                    self.u_secondary += lvl.dt * (w2 * k.impl + w2E * k.expl)
+
 
 class ForwardEuler(RungeKutta):
     """
@@ -480,6 +498,14 @@ class BackwardEuler(RungeKutta):
     nodes, weights, matrix = generator.genCoeffs()
 
 
+class IMEXEuler(RungeKuttaIMEX):
+    nodes = BackwardEuler.nodes
+    weights = BackwardEuler.weights
+
+    matrix = BackwardEuler.matrix
+    matrix_explicit = ForwardEuler.matrix
+
+
 class CrankNicolson(RungeKutta):
     """
     Implicit Runge-Kutta method of second order, A-stable.
@@ -521,8 +547,13 @@ class Heun_Euler(RungeKutta):
     Second order explicit embedded Runge-Kutta method.
     """
 
+    ButcherTableauClass = ButcherTableauEmbedded
+
     generator = RK_SCHEMES["HEUN"]()
-    nodes, weights, matrix = generator.genCoeffs()
+    nodes, _weights, matrix = generator.genCoeffs()
+    weights = np.zeros((2, len(_weights)))
+    weights[0] = _weights
+    weights[1] = matrix[-1]
 
     @classmethod
     def get_update_order(cls):
@@ -697,3 +728,41 @@ class ARK548L2SA(RungeKuttaIMEX):
     @classmethod
     def get_update_order(cls):
         return 5
+
+
+class ARK324L2SAERK(RungeKutta):
+    generator = RK_SCHEMES["ARK324L2SAERK"]()
+    nodes, weights, matrix = generator.genCoeffs(embedded=True)
+    ButcherTableauClass = ButcherTableauEmbedded
+
+    @classmethod
+    def get_update_order(cls):
+        return 3
+
+
+class ARK324L2SAESDIRK(ARK324L2SAERK):
+    generator = RK_SCHEMES["ARK324L2SAESDIRK"]()
+    matrix = generator.Q
+
+
+class ARK32(RungeKuttaIMEX):
+    ButcherTableauClass = ButcherTableauEmbedded
+    ButcherTableauClass_explicit = ButcherTableauEmbedded
+
+    nodes = ARK324L2SAESDIRK.nodes
+    weights = ARK324L2SAESDIRK.weights
+
+    matrix = ARK324L2SAESDIRK.matrix
+    matrix_explicit = ARK324L2SAERK.matrix
+
+    @classmethod
+    def get_update_order(cls):
+        return 3
+
+
+class ARK2(RungeKuttaIMEX):
+    generator_IMP = RK_SCHEMES["ARK222EDIRK"]()
+    generator_EXP = RK_SCHEMES["ARK222ERK"]()
+
+    nodes, weights, matrix = generator_IMP.genCoeffs()
+    _, weights_explicit, matrix_explicit = generator_EXP.genCoeffs()