Added serial implementation for ParaDiag for collocation methods

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

@@ -6,6 +6,7 @@
 from pySDC.core.base_transfer import BaseTransfer
 from pySDC.helpers.pysdc_helper import FrozenClass
 from pySDC.implementations.convergence_controller_classes.check_convergence import CheckConvergence
+from pySDC.implementations.convergence_controller_classes.store_uold import StoreUOld
 from pySDC.implementations.hooks.default_hook import DefaultHooks
 from pySDC.implementations.hooks.log_timings import CPUTimings
 
@@ -41,6 +42,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 +343,68 @@ 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
+        """
+        # TODO: where should I put alpha? When I want to adapt it, maybe it shouldn't be in the controller?
+        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.base_convergence_controllers += [StoreUOld]
+
+        self.ParaDiag_block_u0 = None
+        self.n_steps = n_steps
+
+    def FFT_in_time(self):
+        """
+        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)
+
+    def iFFT_in_time(self):
+        """
+        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)

pySDC/core/level.py

@@ -82,6 +82,7 @@ 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.f = [None] * (self.sweep.coll.num_nodes + 1)
         self.fold = [None] * (self.sweep.coll.num_nodes + 1)
 

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()