Runge-Kutta sweeper for DAEs (#432)

* Added DIRK43_2 and EDIRK4 methods

* Runge-Kutta sweeper for DAEs

* Every problem needs a du_exact method when using RK-DAE sweepers

* Tests for RK methods

* Added label in documentation

* Typos + some corrections

* First requested changes

* Added statement: predict only with du_exact in first step

* Make function for implicit system to be solved as static method in FI-SDC and SI-SDC sweepers

* Removed deep copy

Author: lisawim

pySDC/implementations/sweeper_classes/Runge_Kutta.py

@@ -590,6 +590,38 @@ def get_update_order(cls):
         return 4
 
 
+class DIRK43_2(RungeKutta):
+    """
+    L-stable Diagonally Implicit RK method with four stages of order 3.
+    Taken from [here](https://en.wikipedia.org/wiki/List_of_Runge%E2%80%93Kutta_methods).
+    """
+
+    nodes = np.array([0.5, 2.0 / 3.0, 0.5, 1.0])
+    weights = np.array([3.0 / 2.0, -3.0 / 2.0, 0.5, 0.5])
+    matrix = np.zeros((4, 4))
+    matrix[0, 0] = 0.5
+    matrix[1, :2] = [1.0 / 6.0, 0.5]
+    matrix[2, :3] = [-0.5, 0.5, 0.5]
+    matrix[3, :] = [3.0 / 2.0, -3.0 / 2.0, 0.5, 0.5]
+    ButcherTableauClass = ButcherTableau
+
+
+class EDIRK4(RungeKutta):
+    """
+    Stiffly accurate, fourth-order EDIRK with four stages. Taken from
+    [here](https://ntrs.nasa.gov/citations/20160005923), second one in eq. (216).
+    """
+
+    nodes = np.array([0.0, 3.0 / 2.0, 7.0 / 5.0, 1.0])
+    weights = np.array([13.0, 84.0, -125.0, 70.0]) / 42.0
+    matrix = np.zeros((4, 4))
+    matrix[0, 0] = 0
+    matrix[1, :2] = [3.0 / 4.0, 3.0 / 4.0]
+    matrix[2, :3] = [447.0 / 675.0, -357.0 / 675.0, 855.0 / 675.0]
+    matrix[3, :] = [13.0 / 42.0, 84.0 / 42.0, -125.0 / 42.0, 70.0 / 42.0]
+    ButcherTableauClass = ButcherTableau
+
+
 class ESDIRK53(RungeKutta):
     """
     A-stable embedded RK pair of orders 5 and 3, ESDIRK5(3)6L[2]SA.

pySDC/projects/DAE/misc/ProblemDAE.py

@@ -36,14 +36,19 @@ def __init__(self, nvars, newton_tol):
         self.work_counters['newton'] = WorkCounter()
         self.work_counters['rhs'] = WorkCounter()
 
-    def solve_system(self, impl_sys, u0, t):
+    def solve_system(self, impl_sys, u_approx, factor, u0, t):
         r"""
         Solver for nonlinear implicit system (defined in sweeper).
 
         Parameters
         ----------
         impl_sys : callable
             Implicit system to be solved.
+        u_approx : dtype_u
+            Approximation of solution :math:`u` which is needed to solve
+            the implicit system.
+        factor : float
+            Abbrev. for the node-to-node stepsize.
         u0 : dtype_u
             Initial guess for solver.
         t : float
@@ -57,7 +62,7 @@ def solve_system(self, impl_sys, u0, t):
         me = self.dtype_u(self.init)
 
         def implSysFlatten(unknowns, **kwargs):
-            sys = impl_sys(unknowns.reshape(me.shape).view(type(u0)), **kwargs)
+            sys = impl_sys(unknowns.reshape(me.shape).view(type(u0)), self, factor, u_approx, t, **kwargs)
             return sys.flatten()
 
         opt = root(
@@ -69,3 +74,21 @@ def implSysFlatten(unknowns, **kwargs):
         me[:] = opt.x.reshape(me.shape)
         self.work_counters['newton'].niter += opt.nfev
         return me
+
+    def du_exact(self, t):
+        r"""
+        Routine for the derivative of the exact solution at time :math:`t \leq 1`.
+        For this problem, the exact solution is piecewise.
+
+        Parameters
+        ----------
+        t : float
+            Time of the exact solution.
+
+        Returns
+        -------
+        me : dtype_u
+            Derivative of exact solution.
+        """
+
+        raise NotImplementedError('ERROR: problem has to implement du_exact(self, t)!')

pySDC/projects/DAE/problems/DiscontinuousTestDAE.py

@@ -128,6 +128,33 @@ def u_exact(self, t, **kwargs):
             me.alg[0] = np.sinh(self.t_switch_exact)
         return me
 
+    def du_exact(self, t, **kwargs):
+        r"""
+        Routine for the derivative of the exact solution at time :math:`t \leq 1`.
+        For this problem, the exact solution is piecewise.
+
+        Parameters
+        ----------
+        t : float
+            Time of the exact solution.
+
+        Returns
+        -------
+        me : dtype_u
+            Derivative of exact solution.
+        """
+
+        assert t >= 1, 'ERROR: u_exact only available for t>=1'
+
+        me = self.dtype_u(self.init)
+        if t <= self.t_switch_exact:
+            me.diff[0] = np.sinh(t)
+            me.alg[0] = np.cosh(t)
+        else:
+            me.diff[0] = np.sinh(self.t_switch_exact)
+            me.alg[0] = np.cosh(self.t_switch_exact)
+        return me
+
     def get_switching_info(self, u, t):
         r"""
         Provides information about the state function of the problem. A change in sign of the state function

pySDC/projects/DAE/problems/simple_DAE.py

@@ -210,6 +210,26 @@ def u_exact(self, t):
         me.alg[0] = -np.exp(t) / (2 - t)
         return me
 
+    def du_exact(self, t):
+        """
+        Routine for the derivative of the exact solution.
+
+        Parameters
+        ----------
+        t : float
+            The time of the reference solution.
+
+        Returns
+        -------
+        me : dtype_u
+            The reference solution as mesh object containing three components.
+        """
+
+        me = self.dtype_u(self.init)
+        me.diff[:2] = (np.exp(t), np.exp(t))
+        me.alg[0] = (np.exp(t) * (t - 3)) / ((2 - t) ** 2)
+        return me
+
 
 class problematic_f(ptype_dae):
     r"""
@@ -293,3 +313,22 @@ def u_exact(self, t):
         me = self.dtype_u(self.init)
         me[:] = (np.sin(t), 0)
         return me
+
+    def du_exact(self, t):
+        """
+        Routine for the derivative of the exact solution.
+
+        Parameters
+        ----------
+        t : float
+            The time of the reference solution.
+
+        Returns
+        -------
+        me : dtype_u
+            The reference solution as mesh object containing two components.
+        """
+
+        me = self.dtype_u(self.init)
+        me[:] = (np.cos(t), 0)
+        return me

pySDC/projects/DAE/sweepers/RungeKuttaDAE.py

@@ -0,0 +1,158 @@
+from pySDC.projects.DAE.sweepers.fully_implicit_DAE import fully_implicit_DAE
+from pySDC.implementations.sweeper_classes.Runge_Kutta import (
+    RungeKutta,
+    BackwardEuler,
+    CrankNicholson,
+    EDIRK4,
+    DIRK43_2,
+)
+
+
+class RungeKuttaDAE(RungeKutta):
+    r"""
+    Custom sweeper class to implement Runge-Kutta (RK) methods for general differential-algebraic equations (DAEs)
+    of the form
+
+    .. math::
+        0 = F(u, u', t).
+    
+    RK methods for general DAEs have the form
+
+    .. math::
+        0 = F(u_0 + \Delta t \sum_{j=1}^M a_{i,j} U_j, U_m),
+
+    .. math::
+        u_M = u_0 + \Delta t \sum_{j=1}^M b_j U_j.
+
+    In pySDC, RK methods are implemented in the way that the coefficient matrix :math:`A` in the Butcher
+    tableau is a lower triangular matrix so that the stages are updated node-by-node. This class therefore only supports
+    RK methods with lower triangular matrices in the tableau.
+
+    Parameters
+    ----------
+    params : dict
+        Parameters passed to the sweeper.
+
+    Attributes
+    ----------
+    du_init : dtype_f
+        Stores the initial condition for each step.
+
+    Note
+    ----
+    When using a RK sweeper to simulate a problem make sure the DAE problem class has a ``du_exact`` method since RK methods need an initial
+    condition for :math:`u'(t_0)` as well.
+
+    In order to implement a new RK method for DAEs a new tableau can be added in ``pySDC.implementations.sweeper_classes.Runge_Kutta.py``.
+    For example, a new method called ``newRungeKuttaMethod`` with nodes :math:`c=(c_1, c_2, c_3)`, weights :math:`b=(b_1, b_2, b_3)` and
+    coefficient matrix
+
+    ..math::
+        \begin{eqnarray}
+            A = \begin{pmatrix}
+                a_{11} & 0 & 0 \\
+                a_{21} & a_{22} & 0 \\
+                a_{31} & a_{32} & & 0 \\
+            \end{pmatrix}
+        \end{eqnarray}
+
+    can be implemented as follows:
+
+    >>> class newRungeKuttaMethod(RungeKutta):
+    >>>     nodes = np.array([c1, c2, c3])
+    >>>     weights = np.array([b1, b2, b3])
+    >>>     matrix = np.zeros((3, 3))
+    >>>     matrix[0, 0] = a11
+    >>>     matrix[1, :2] = [a21, a22]
+    >>>     matrix[2, :] = [a31, a32, a33]
+    >>>     ButcherTableauClass = ButcherTableau
+
+    The new class ``newRungeKuttaMethodDAE`` can then be used by defining the DAE class inheriting from both, this base class and class containing
+    the Butcher tableau:
+
+    >>> class newRungeKuttaMethodDAE(RungeKuttaDAE, newRungeKuttaMethod):
+    >>>     pass
+
+    More details can be found [here](https://github.com/Parallel-in-Time/pySDC/blob/master/pySDC/implementations/sweeper_classes/Runge_Kutta.py).
+    """
+
+    def __init__(self, params):
+        super().__init__(params)
+        self.du_init = None
+        self.fully_initialized = False
+
+    def predict(self):
+        """
+        Predictor to fill values with zeros at nodes before first sweep.
+        """
+
+        # get current level and problem
+        lvl = self.level
+        prob = lvl.prob
+
+        if not self.fully_initialized:
+            self.du_init = prob.du_exact(lvl.time)
+            self.fully_initialized = True
+
+        lvl.f[0] = prob.dtype_f(self.du_init)
+        for m in range(1, self.coll.num_nodes + 1):
+            lvl.u[m] = prob.dtype_u(init=prob.init, val=0.0)
+            lvl.f[m] = prob.dtype_f(init=prob.init, val=0.0)
+
+        lvl.status.unlocked = True
+        lvl.status.updated = True
+
+    def update_nodes(self):
+        r"""
+        Updates the values of solution ``u`` and their gradient stored in ``f``.
+        """
+
+        # get current level and problem description
+        lvl = self.level
+        prob = lvl.prob
+
+        # only if the level has been touched before
+        assert lvl.status.unlocked
+        assert lvl.status.sweep <= 1, "RK schemes are direct solvers. Please perform only 1 iteration!"
+
+        M = self.coll.num_nodes
+        for m in range(M):
+            u_approx = prob.dtype_u(lvl.u[0])
+            for j in range(1, m + 1):
+                u_approx += lvl.dt * self.QI[m + 1, j] * lvl.f[j][:]
+
+            finit = lvl.f[m].flatten()
+            lvl.f[m + 1][:] = prob.solve_system(
+                fully_implicit_DAE.F,
+                u_approx,
+                lvl.dt * self.QI[m + 1, m + 1],
+                finit,
+                lvl.time + lvl.dt * self.coll.nodes[m + 1],
+            )
+
+        # Update numerical solution - update value only at last node
+        lvl.u[-1][:] = lvl.u[0]
+        for j in range(1, M + 1):
+            lvl.u[-1][:] += lvl.dt * self.coll.Qmat[-1, j] * lvl.f[j][:]
+
+        self.du_init = prob.dtype_f(lvl.f[-1])
+
+        lvl.status.updated = True
+
+        return None
+
+
+class BackwardEulerDAE(RungeKuttaDAE, BackwardEuler):
+    pass
+
+
+class TrapezoidalRuleDAE(RungeKuttaDAE, CrankNicholson):
+    pass
+
+
+class EDIRK4DAE(RungeKuttaDAE, EDIRK4):
+    pass
+
+
+class DIRK43_2DAE(RungeKuttaDAE, DIRK43_2):
+    pass