Serial implementation for ParaDiag with collocation methods (#530)

* Added serial implementation for ParaDiag for collocation methods

* Removed some old comments

* Added test that ParaDiag converges within the bounds for Dahlquist
problem

* Fix test on older Python versions

* Added smaller values of alpha to the tests

* Implemented separate Jacobian inversion to fix bug pointed out by
@JHopeCollins

* Update pySDC/tutorial/step_9/B_paradiag_for_nonlinear_problems.py

Co-authored-by: Josh Hope-Collins <[email protected]>

* Keep spatially extended residual in the levels after computation

* Refactored increment formulation

---------

Co-authored-by: Josh Hope-Collins <[email protected]>

Author: brownbaerchen

docs/source/index.rst

@@ -25,6 +25,7 @@ Tutorial
    tutorial/step_6.rst
    tutorial/step_7.rst
    tutorial/step_8.rst
+   tutorial/step_9.rst
 
 Playgrounds
 -----------

docs/source/tutorial/doc_step_9_A.rst

@@ -0,0 +1,7 @@
+Full code: `pySDC/tutorial/step_9/A_paradiag_for_linear_problems.py <https://github.com/Parallel-in-Time/pySDC/blob/master/pySDC/tutorial/step_9/A_paradiag_for_linear_problems.py>`_
+
+.. literalinclude:: ../../../pySDC/tutorial/step_9/A_paradiag_for_linear_problems.py
+
+Results:
+
+.. literalinclude:: ../../../data/step_9_A_out.txt

docs/source/tutorial/doc_step_9_B.rst

@@ -0,0 +1,7 @@
+Full code: `pySDC/tutorial/step_9/B_paradiag_for_nonlinear_problems.py <https://github.com/Parallel-in-Time/pySDC/blob/master/pySDC/tutorial/step_9/B_paradiag_for_nonlinear_problems.py>`_
+
+.. literalinclude:: ../../../pySDC/tutorial/step_9/B_paradiag_for_nonlinear_problems.py
+
+Results:
+
+.. literalinclude:: ../../../data/step_9_B_out.txt

docs/source/tutorial/doc_step_9_C.rst

@@ -0,0 +1,7 @@
+Full code: `pySDC/tutorial/step_9/C_paradiag_in_pySDC.py <https://github.com/Parallel-in-Time/pySDC/blob/master/pySDC/tutorial/step_9/C_paradiag_in_pySDC.py>`_
+
+.. literalinclude:: ../../../pySDC/tutorial/step_9/C_paradiag_in_pySDC.py
+
+Results:
+
+.. literalinclude:: ../../../data/step_9_C_out.txt

docs/source/tutorial/step_9.rst

@@ -0,0 +1 @@
+.. include:: /../../pySDC/tutorial/step_9/README.rst

pySDC/core/controller.py

@@ -41,6 +41,7 @@ def __init__(self, controller_params, description, useMPI=None):
             controller_params (dict): parameter set for the controller and the steps
         """
         self.useMPI = useMPI
+        self.description = description
 
         # check if we have a hook on this list. If not, use default class.
         self.__hooks = []
@@ -341,3 +342,65 @@ def return_stats(self):
         for hook in self.hooks:
             stats = {**stats, **hook.return_stats()}
         return stats
+
+
+class ParaDiagController(Controller):
+
+    def __init__(self, controller_params, description, n_steps, useMPI=None):
+        """
+        Initialization routine for ParaDiag controllers
+
+        Args:
+           num_procs: number of parallel time steps (still serial, though), can be 1
+           controller_params: parameter set for the controller and the steps
+           description: all the parameters to set up the rest (levels, problems, transfer, ...)
+           n_steps (int): Number of parallel steps
+           alpha (float): alpha parameter for ParaDiag
+        """
+        from pySDC.implementations.sweeper_classes.ParaDiagSweepers import QDiagonalization
+
+        if QDiagonalization in description['sweeper_class'].__mro__:
+            description['sweeper_params']['ignore_ic'] = True
+            description['sweeper_params']['update_f_evals'] = False
+        else:
+            logging.getLogger('controller').warning(
+                f'Warning: Your sweeper class {description["sweeper_class"]} is not derived from {QDiagonalization}. You probably want to use another sweeper class.'
+            )
+
+        if controller_params.get('all_to_done', False):
+            raise NotImplementedError('ParaDiag only implemented with option `all_to_done=True`')
+        if 'alpha' not in controller_params.keys():
+            from pySDC.core.errors import ParameterError
+
+            raise ParameterError('Please supply alpha as a parameter to the ParaDiag controller!')
+        controller_params['average_jacobian'] = controller_params.get('average_jacobian', True)
+
+        controller_params['all_to_done'] = True
+        super().__init__(controller_params=controller_params, description=description, useMPI=useMPI)
+
+        self.n_steps = n_steps
+
+    def FFT_in_time(self, quantity):
+        """
+        Compute weighted forward FFT in time. The weighting is determined by the alpha parameter in ParaDiag
+
+        Note: The implementation via matrix-vector multiplication may be inefficient and less stable compared to an FFT
+              with transposes!
+        """
+        if not hasattr(self, '__FFT_matrix'):
+            from pySDC.helpers.ParaDiagHelper import get_weighted_FFT_matrix
+
+            self.__FFT_matrix = get_weighted_FFT_matrix(self.n_steps, self.params.alpha)
+
+        self.apply_matrix(self.__FFT_matrix, quantity)
+
+    def iFFT_in_time(self, quantity):
+        """
+        Compute weighted backward FFT in time. The weighting is determined by the alpha parameter in ParaDiag
+        """
+        if not hasattr(self, '__iFFT_matrix'):
+            from pySDC.helpers.ParaDiagHelper import get_weighted_iFFT_matrix
+
+            self.__iFFT_matrix = get_weighted_iFFT_matrix(self.n_steps, self.params.alpha)
+
+        self.apply_matrix(self.__iFFT_matrix, quantity)

pySDC/core/level.py

@@ -82,6 +82,9 @@ def __init__(self, problem_class, problem_params, sweeper_class, sweeper_params,
         self.uend = None
         self.u = [None] * (self.sweep.coll.num_nodes + 1)
         self.uold = [None] * (self.sweep.coll.num_nodes + 1)
+        self.u_avg = [None] * self.sweep.coll.num_nodes
+        self.residual = [None] * self.sweep.coll.num_nodes
+        self.increment = [None] * self.sweep.coll.num_nodes
         self.f = [None] * (self.sweep.coll.num_nodes + 1)
         self.fold = [None] * (self.sweep.coll.num_nodes + 1)
 

pySDC/core/problem.py

@@ -176,3 +176,35 @@ def plot(self, u, t=None, fig=None):
         None
         """
         raise NotImplementedError
+
+    def solve_system(self, rhs, dt, u0, t):
+        """
+        Perform an Euler step.
+
+        Args:
+            rhs: Right hand side for the Euler step
+            dt (float): Step size for the Euler step
+            u0: Initial guess
+            t (float): Current time
+
+        Returns:
+            solution to the Euler step
+        """
+        raise NotImplementedError
+
+    def solve_jacobian(self, rhs, dt, u=None, u0=None, t=0, **kwargs):
+        """
+        Solve the Jacobian for an Euler step, linearized around u.
+        This defaults to an Euler step to accommodate linear problems.
+
+        Args:
+            rhs: Right hand side for the Euler step
+            dt (float): Step size for the Euler step
+            u: Solution to linearize around
+            u0: Initial guess
+            t (float): Current time
+
+        Returns:
+            Solution
+        """
+        return self.solve_system(rhs, dt, u0=u, t=t, **kwargs)

pySDC/core/sweeper.py

@@ -192,14 +192,14 @@ def compute_residual(self, stage=''):
 
         # build QF(u)
         res_norm = []
-        res = self.integrate()
+        L.residual = self.integrate()
         for m in range(self.coll.num_nodes):
-            res[m] += L.u[0] - L.u[m + 1]
+            L.residual[m] += L.u[0] - L.u[m + 1]
             # add tau if associated
             if L.tau[m] is not None:
-                res[m] += L.tau[m]
+                L.residual[m] += L.tau[m]
             # use abs function from data type here
-            res_norm.append(abs(res[m]))
+            res_norm.append(abs(L.residual[m]))
 
         # find maximal residual over the nodes
         if L.params.residual_type == 'full_abs':

pySDC/helpers/ParaDiagHelper.py

@@ -0,0 +1,131 @@
+import numpy as np
+import scipy.sparse as sp
+
+
+def get_FFT_matrix(N):
+    """
+    Get matrix for computing FFT of size N. Normalization is like "ortho" in numpy.
+    Compute inverse FFT by multiplying by the complex conjugate (numpy.conjugate) of this matrix
+
+    Args:
+        N (int): Size of the data to be transformed
+
+    Returns:
+        numpy.ndarray: Dense square matrix to compute forward transform
+    """
+    idx_1d = np.arange(N, dtype=complex)
+    i1, i2 = np.meshgrid(idx_1d, idx_1d)
+
+    return np.exp(-2 * np.pi * 1j * i1 * i2 / N) / np.sqrt(N)
+
+
+def get_E_matrix(N, alpha=0):
+    """
+    Get NxN matrix with -1 on the lower subdiagonal, -alpha in the top right and 0 elsewhere
+
+    Args:
+        N (int): Size of the matrix
+        alpha (float): Negative of value in the top right
+
+    Returns:
+        sparse E matrix
+    """
+    E = sp.diags(
+        [
+            -1.0,
+        ]
+        * (N - 1),
+        offsets=-1,
+    ).tolil()
+    E[0, -1] = -alpha
+    return E
+
+
+def get_J_matrix(N, alpha):
+    """
+    Get matrix for weights in the weighted inverse FFT
+
+    Args:
+        N (int): Size of the matrix
+        alpha (float): alpha parameter in ParaDiag
+
+    Returns:
+        sparse J matrix
+    """
+    gamma = alpha ** (-np.arange(N) / N)
+    return sp.diags(gamma)
+
+
+def get_J_inv_matrix(N, alpha):
+    """
+    Get matrix for weights in the weighted FFT
+
+    Args:
+        N (int): Size of the matrix
+        alpha (float): alpha parameter in ParaDiag
+
+    Returns:
+        sparse J_inv matrix
+    """
+    gamma = alpha ** (-np.arange(N) / N)
+    return sp.diags(1 / gamma)
+
+
+def get_weighted_FFT_matrix(N, alpha):
+    """
+    Get matrix for the weighted FFT
+
+    Args:
+        N (int): Size of the matrix
+        alpha (float): alpha parameter in ParaDiag
+
+    Returns:
+        Dense weighted FFT matrix
+    """
+    return get_FFT_matrix(N) @ get_J_inv_matrix(N, alpha)
+
+
+def get_weighted_iFFT_matrix(N, alpha):
+    """
+    Get matrix for the weighted inverse FFT
+
+    Args:
+        N (int): Size of the matrix
+        alpha (float): alpha parameter in ParaDiag
+
+    Returns:
+        Dense weighted FFT matrix
+    """
+    return get_J_matrix(N, alpha) @ np.conjugate(get_FFT_matrix(N))
+
+
+def get_H_matrix(N, sweeper_params):
+    """
+    Get sparse matrix for computing the collocation update. Requires not to do a collocation update!
+
+    Args:
+        N (int): Number of collocation nodes
+        sweeper_params (dict): Parameters for the sweeper
+
+    Returns:
+        Sparse matrix for collocation update
+    """
+    assert sweeper_params['quad_type'] == 'RADAU-RIGHT'
+    H = sp.eye(N).tolil() * 0
+    H[:, -1] = 1
+    return H
+
+
+def get_G_inv_matrix(l, L, alpha, sweeper_params):
+    M = sweeper_params['num_nodes']
+    I_M = sp.eye(M)
+    E_alpha = get_E_matrix(L, alpha)
+    H = get_H_matrix(M, sweeper_params)
+
+    gamma = alpha ** (-np.arange(L) / L)
+    diags = np.fft.fft(1 / gamma * E_alpha[:, 0].toarray().flatten(), norm='backward')
+    G = (diags[l] * H + I_M).tocsc()
+    if M > 1:
+        return sp.linalg.inv(G).toarray()
+    else:
+        return 1 / G.toarray()