Added solve_system to all DAE problem classes (#371)

* Added solve_system method to all DAE classes

* Removed comment

* Moved solve_system to generic ProblemDAE class

* Used possibility to update documentation :D

* Moved work_counters for rhs also to ProblemDAE

* Inheritance from generic_implicit

pySDC/projects/DAE/misc/ProblemDAE.py

@@ -1,25 +1,67 @@
 import numpy as np
+from scipy.optimize import root
 
-from pySDC.core.Problem import ptype
+from pySDC.core.Problem import ptype, WorkCounter
 from pySDC.implementations.datatype_classes.mesh import mesh
 
 
 class ptype_dae(ptype):
-    """
-    Interface class for DAE problems. Ensures that all parameters are passed that are needed by DAE sweepers
+    r"""
+    This class implements a generic DAE class and illustrates the interface class for DAE problems.
+    It ensures that all parameters are passed that are needed by DAE sweepers.
+
+    Parameters
+    ----------
+    nvars : int
+        Number of unknowns of the problem class.
+    newton_tol : float
+        Tolerance for the nonlinear solver.
+
+    Attributes
+    ----------
+    work_counters : WorkCounter
+        Counts the work, here the number of function calls during the nonlinear solve is logged and stored
+        in work_counters['newton']. The number of each function class of the right-hand side is then stored
+        in work_counters['rhs']
     """
 
     dtype_u = mesh
     dtype_f = mesh
 
     def __init__(self, nvars, newton_tol):
-        """
-        Initialization routine
-
-        Args:
-            problem_params (dict): custom parameters for the example
-            dtype_u: mesh data type (will be passed parent class)
-            dtype_f: mesh data type (will be passed parent class)
-        """
+        """Initialization routine"""
         super().__init__((nvars, None, np.dtype('float64')))
         self._makeAttributeAndRegister('nvars', 'newton_tol', localVars=locals(), readOnly=True)
+
+        self.work_counters['newton'] = WorkCounter()
+        self.work_counters['rhs'] = WorkCounter()
+
+    def solve_system(self, impl_sys, u0, t):
+        r"""
+        Solver for nonlinear implicit system (defined in sweeper).
+
+        Parameters
+        ----------
+        impl_sys : callable
+            Implicit system to be solved.
+        u0 : dtype_u
+            Initial guess for solver.
+        t : float
+            Current time :math:`t`.
+
+        Returns
+        -------
+        me : dtype_u
+            Numerical solution.
+        """
+
+        me = self.dtype_u(self.init)
+        opt = root(
+            impl_sys,
+            u0,
+            method='hybr',
+            tol=self.newton_tol,
+        )
+        me[:] = opt.x
+        self.work_counters['newton'].niter += opt.nfev
+        return me

pySDC/projects/DAE/problems/simple_DAE.py

@@ -73,6 +73,7 @@ def eval_f(self, u, du, t):
         # weight somehow
         f = self.dtype_f(self.init)
         f[:] = (du[0] - u[2], du[1] - u[3], du[2] + u[4] * u[0], du[3] + u[4] * u[1] + g, u[0] ** 2 + u[1] ** 2 - 1)
+        self.work_counters['rhs']()
         return f
 
     def u_exact(self, t):
@@ -164,6 +165,7 @@ def eval_f(self, u, du, t):
             -du[1] + (1 - a) / (t - 2) * u[0] - u[1] + (a - 1) * u[2] + 2 * np.exp(t),
             (t + 2) * u[0] + (t**2 - 4) * u[1] - (t**2 + t - 2) * np.exp(t),
         )
+        self.work_counters['rhs']()
         return f
 
     def u_exact(self, t):
@@ -207,7 +209,7 @@ class problematic_f(ptype_dae):
 
     Attributes
     ----------
-    eta: float
+    eta : float
         Specific parameter of the problem.
 
     References
@@ -244,6 +246,7 @@ def eval_f(self, u, du, t):
             u[0] + self.eta * t * u[1] - np.sin(t),
             du[0] + self.eta * t * du[1] + (1 + self.eta) * u[1] - np.cos(t),
         )
+        self.work_counters['rhs']()
         return f
 
     def u_exact(self, t):

pySDC/projects/DAE/problems/synchronous_machine.py

@@ -257,6 +257,7 @@ def eval_f(self, u, du, t):
             -psi_Q1 + self.L_mq * i_q + self.L_Q1 * i_Q1 + self.L_mq * i_Q2,
             -psi_Q2 + self.L_mq * i_q + self.L_mq * i_Q1 + self.L_Q2 * i_Q2,
         )
+        self.work_counters['rhs']()
         return f
 
     def u_exact(self, t):

pySDC/projects/DAE/problems/transistor_amplifier.py

@@ -104,6 +104,7 @@ def eval_f(self, u, du, t):
             (u_b - u[3]) / r_k + c_3 * (du[4] - du[3]) - alpha * _transistor(u[1] - u[2]),
             -u[4] / r_k + c_3 * (du[3] - du[4]),
         )
+        self.work_counters['rhs']()
         return f
 
     def u_exact(self, t):
@@ -240,6 +241,7 @@ def eval_f(self, u, du, t):
             (u_b - u[6]) / r_k - c_5 * (du[6] - du[7]) - alpha * _transistor(u[4] - u[5]),
             -u[7] / r_k + c_5 * (du[6] - du[7]),
         )
+        self.work_counters['rhs']()
         return f
 
     def u_exact(self, t):

pySDC/projects/DAE/sweepers/fully_implicit_DAE.py

@@ -2,19 +2,47 @@
 from scipy import optimize
 
 from pySDC.core.Errors import ParameterError
-from pySDC.core.Sweeper import sweeper
+from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit
 
 
-class fully_implicit_DAE(sweeper):
-    """
-    Custom sweeper class, implements Sweeper.py
+class fully_implicit_DAE(generic_implicit):
+    r"""
+    Custom sweeper class to implement the fully-implicit SDC for solving DAEs. It solves fully-implicit DAE problems
+    of the form
+
+    .. math::
+        F(t, u, u') = 0.
+
+    It solves a collocation problem of the form
+
+    .. math::
+        F(\tau, \vec{U}_0 + \Delta t (\mathbf{Q} \otimes \mathbf{I}_n) \vec{U}, \vec{U}) = 0,
+
+    where
+
+    - :math:`\tau=(\tau_1,..,\tau_M) in \mathbb{R}^M` the vector of collocation nodes,
+    - :math:`\vec{U}_0 = (u_0,..,u_0) \in \mathbb{R}^{Mn}` the vector of initial condition spread to each node,
+    - spectral integration matrix :math:`\mathbf{Q} \in \mathbb{R}^{M \times M}`,
+    - :math:`\vec{U}=(U_1,..,U_M) \in \mathbb{R}^{Mn}` the vector of unknown derivatives
+      :math:`U_m \approx U(\tau_m) = u'(\tau_m) \in \mathbb{R}^n`,
+    - and identity matrix :math:`\mathbf{I}_n \in \mathbb{R}^{n \times n}`.
+
+    The construction of this sweeper is based on the concepts outlined in [1]_.
 
-    Sweeper to solve first order differential equations in fully implicit form
-    Primarily implemented to be used with differential algebraic equations
-    Based on the concepts outlined in "Arbitrary order Krylov deferred correction methods for differential algebraic equations" by Huang et al.
+    Parameters
+    ----------
+    params : dict
+        Parameters passed to the sweeper.
 
-    Attributes:
-        QI: implicit Euler integration matrix
+    Attributes
+    ----------
+    QI : np.2darray
+        Implicit Euler integration matrix.
+
+    References
+    ----------
+    .. [1] J. Huang, J. Jun, M. L. Minion. Arbitrary order Krylov deferred correction methods for differential algebraic equation.
+       J. Comput. Phys. Vol. 221 No. 2 (2007).
     """
 
     def __init__(self, params):
@@ -37,39 +65,10 @@ def __init__(self, params):
 
         self.QI = self.get_Qdelta_implicit(coll=self.coll, qd_type=self.params.QI)
 
-    # TODO: hijacking this function to return solution from its gradient i.e. fundamental theorem of calculus.
-    # This works well since (ab)using level.f to store the gradient. Might need to change this for release?
-    def integrate(self):
-        """
-        Returns the solution by integrating its gradient (fundamental theorem of calculus)
-        Note that level.f stores the gradient values in the fully implicit case, rather than the evaluation of the rhs as in the ODE case
-
-        Returns:
-            list of dtype_u: containing the integral as values
-        """
-
-        # get current level and problem description
-        L = self.level
-        P = L.prob
-        M = self.coll.num_nodes
-
-        me = []
-
-        # integrate gradient over all collocation nodes
-        for m in range(1, M + 1):
-            # new instance of dtype_u, initialize values with 0
-            me.append(P.dtype_u(P.init, val=0.0))
-            for j in range(1, M + 1):
-                me[-1] += L.dt * self.coll.Qmat[m, j] * L.f[j]
-
-        return me
-
     def update_nodes(self):
-        """
-        Update the u- and f-values at the collocation nodes -> corresponds to a single iteration of the preconditioned Richardson iteration in "ordinary" SDC
-
-        Returns:
-            None
+        r"""
+        Updates values of ``u`` and ``f`` at collocation nodes. This correspond to a single iteration of the
+        preconditioned Richardson iteration in **"ordinary"** SDC.
         """
 
         # get current level and problem description
@@ -104,32 +103,45 @@ def update_nodes(self):
                 u_approx += L.dt * self.QI[m, j] * L.f[j]
 
             # params contains U = u'
-            def impl_fn(params):
-                # make params into a mesh object
+            def implSystem(params):
+                """
+                Build implicit system to solve in order to find the unknowns.
+
+                Parameters
+                ----------
+                params : dtype_u
+                    Unknowns of the system.
+
+                Returns
+                -------
+                sys :
+                    System to be solved as implicit function.
+                """
+
                 params_mesh = P.dtype_f(P.init)
                 params_mesh[:] = params
+
                 # build parameters to pass to implicit function
                 local_u_approx = u_approx
+
                 # note that derivatives of algebraic variables are taken into account here too
                 # these do not directly affect the output of eval_f but rather indirectly via QI
                 local_u_approx += L.dt * self.QI[m, m] * params_mesh
-                return P.eval_f(local_u_approx, params_mesh, L.time + L.dt * self.coll.nodes[m - 1])
+
+                sys = P.eval_f(local_u_approx, params_mesh, L.time + L.dt * self.coll.nodes[m - 1])
+                return sys
 
             # get U_k+1
             # note: not using solve_system here because this solve step is the same for any problem
             # See link for how different methods use the default tol parameter
             # https://github.com/scipy/scipy/blob/8a6f1a0621542f059a532953661cd43b8167fce0/scipy/optimize/_root.py#L220
             # options['xtol'] = P.params.newton_tol
             # options['eps'] = 1e-16
-            opt = optimize.root(
-                impl_fn,
-                L.f[m],
-                method='hybr',
-                tol=P.newton_tol
-                # callback= lambda x, f: print("solution:", x, " residual: ", f)
-            )
+
+            u_new = P.solve_system(implSystem, L.f[m], L.time + L.dt * self.coll.nodes[m - 1])
+
             # update gradient (recall L.f is being used to store the gradient)
-            L.f[m][:] = opt.x
+            L.f[m][:] = u_new
 
         # Update solution approximation
         integral = self.integrate()
@@ -141,12 +153,16 @@ def impl_fn(params):
         return None
 
     def predict(self):
-        """
-        Predictor to fill values at nodes before first sweep
+        r"""
+        Predictor to fill values at nodes before first sweep. It can decide whether the
+
+            - initial condition is spread to each node ('initial_guess' = 'spread'),
+            - zero values are spread to each node ('initial_guess' = 'zero'),
+            - or random values are spread to each collocation node ('initial_guess' = 'random').
 
-        Default prediction for the sweepers, only copies the values to all collocation nodes
-        This function overrides the base implementation by always initialising level.f to zero
-        This is necessary since level.f stores the solution derivative in the fully implicit case, which is not initially known
+        Default prediction for the sweepers, only copies the values to all collocation nodes. This function
+        overrides the base implementation by always initialising ``level.f`` to zero. This is necessary since
+        ``level.f`` stores the solution derivative in the fully implicit case, which is not initially known.
         """
         # get current level and problem description
         L = self.level
@@ -158,7 +174,6 @@ def predict(self):
             if self.params.initial_guess == 'spread':
                 L.u[m] = P.dtype_u(L.u[0])
                 L.f[m] = P.dtype_f(init=P.init, val=0.0)
-            # start with zero everywhere
             elif self.params.initial_guess == 'zero':
                 L.u[m] = P.dtype_u(init=P.init, val=0.0)
                 L.f[m] = P.dtype_f(init=P.init, val=0.0)
@@ -174,15 +189,18 @@ def predict(self):
         L.status.updated = True
 
     def compute_residual(self, stage=None):
-        """
-        Overrides the base implementation
-        Uses the absolute value of the implicit function ||F(u', u, t)|| as the residual
+        r"""
+        Uses the absolute value of the DAE system
 
-        Args:
-            stage (str): The current stage of the step the level belongs to
+        .. math::
+            ||F(t, u, u')||
+
+        for computing the residual.
 
-        Returns:
-            None
+        Parameters
+        ----------
+        stage : str, optional
+            The current stage of the step the level belongs to.
         """
 
         # get current level and problem description
@@ -224,29 +242,3 @@ def compute_residual(self, stage=None):
         L.status.updated = False
 
         return None
-
-    def compute_end_point(self):
-        """
-        Compute u at the right point of the interval
-
-        The value uend computed here is a full evaluation of the Picard formulation unless do_full_update==False
-
-        Returns:
-            None
-        """
-
-        # get current level and problem description
-        L = self.level
-        P = L.prob
-
-        # check if Mth node is equal to right point and do_coll_update is false, perform a simple copy
-        if self.coll.right_is_node and not self.params.do_coll_update:
-            # a copy is sufficient
-            L.uend = P.dtype_u(L.u[-1])
-        else:
-            # start with u0 and add integral over the full interval (using coll.weights)
-            L.uend = P.dtype_u(L.u[0])
-            for m in range(self.coll.num_nodes):
-                L.uend += L.dt * self.coll.weights[m] * L.f[m + 1]
-
-        return None