Refactored increment formulation

Author: brownbaerchen

pySDC/core/controller.py

@@ -6,7 +6,6 @@
 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
 
@@ -378,12 +377,10 @@ def __init__(self, controller_params, description, n_steps, useMPI=None):
 
         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):
+    def FFT_in_time(self, quantity):
         """
         Compute weighted forward FFT in time. The weighting is determined by the alpha parameter in ParaDiag
 
@@ -395,9 +392,9 @@ def FFT_in_time(self):
 
             self.__FFT_matrix = get_weighted_FFT_matrix(self.n_steps, self.params.alpha)
 
-        self.apply_matrix(self.__FFT_matrix)
+        self.apply_matrix(self.__FFT_matrix, quantity)
 
-    def iFFT_in_time(self):
+    def iFFT_in_time(self, quantity):
         """
         Compute weighted backward FFT in time. The weighting is determined by the alpha parameter in ParaDiag
         """
@@ -406,4 +403,4 @@ def iFFT_in_time(self):
 
             self.__iFFT_matrix = get_weighted_iFFT_matrix(self.n_steps, self.params.alpha)
 
-        self.apply_matrix(self.__iFFT_matrix)
+        self.apply_matrix(self.__iFFT_matrix, quantity)

pySDC/core/level.py

@@ -84,6 +84,7 @@ def __init__(self, problem_class, problem_params, sweeper_class, sweeper_params,
         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/implementations/controller_classes/controller_ParaDiag_nonMPI.py

@@ -1,11 +1,9 @@
 import itertools
-import copy as cp
 import numpy as np
-import dill
 
 from pySDC.core.controller import ParaDiagController
 from pySDC.core import step as stepclass
-from pySDC.core.errors import ControllerError, CommunicationError
+from pySDC.core.errors import ControllerError
 from pySDC.implementations.convergence_controller_classes.basic_restarting import BasicRestarting
 from pySDC.helpers.ParaDiagHelper import get_G_inv_matrix
 
@@ -92,7 +90,7 @@ def ParaDiag(self, local_MS_active):
 
         return all(S.status.done for S in local_MS_active)
 
-    def apply_matrix(self, mat):
+    def apply_matrix(self, mat, quantity):
         """
         Apply a matrix on the step level. Needs to be square. Puts the result back into the controller.
 
@@ -112,29 +110,33 @@ def apply_matrix(self, mat):
             None,
         ] * L
 
+        if quantity == 'residual':
+            me = [S.levels[0].residual for S in self.MS]
+        elif quantity == 'increment':
+            me = [S.levels[0].increment for S in self.MS]
+        else:
+            raise NotImplementedError
+
         # compute matrix-vector product
         for i in range(mat.shape[0]):
-            res[i] = [prob.u_init for _ in range(M + 1)]
+            res[i] = [prob.u_init for _ in range(M)]
             for j in range(mat.shape[1]):
-                for m in range(M + 1):
-                    res[i][m] += mat[i, j] * self.MS[j].levels[0].u[m]
+                for m in range(M):
+                    res[i][m] += mat[i, j] * me[j][m]
 
         # put the result in the "output"
         for i in range(mat.shape[0]):
-            for m in range(M + 1):
-                self.MS[i].levels[0].u[m] = res[i][m]
+            for m in range(M):
+                me[i][m] = res[i][m]
 
-    def swap_solution_for_all_at_once_residual(self, local_MS_running):
+    def compute_all_at_once_residual(self, local_MS_running):
         """
-        Replace the solution values in the steps with the all-at-once residual.
-
         This requires to communicate the solutions at the end of the steps to be the initial conditions for the next
         steps. Afterwards, the residual can be computed locally on the steps.
 
         Args:
             local_MS_running (list): list of currently running steps
         """
-        prob = self.MS[0].levels[0].prob
 
         for S in local_MS_running:
             # communicate initial conditions
@@ -143,9 +145,7 @@ def swap_solution_for_all_at_once_residual(self, local_MS_running):
             for hook in self.hooks:
                 hook.pre_comm(step=S, level_number=0)
 
-            if S.status.first:
-                S.levels[0].u[0] = prob.dtype_u(self.ParaDiag_block_u0)
-            else:
+            if not S.status.first:
                 S.levels[0].u[0] = S.prev.levels[0].uend
 
             for hook in self.hooks:
@@ -154,25 +154,16 @@ def swap_solution_for_all_at_once_residual(self, local_MS_running):
             # compute residuals locally
             S.levels[0].sweep.compute_residual()
 
-            # put residual in the solution variables
-            for m in range(S.levels[0].sweep.coll.num_nodes):
-                S.levels[0].u[m + 1] = S.levels[0].residual[m]
-
-    def swap_increment_for_solution(self, local_MS_running):
+    def update_solution(self, local_MS_running):
         """
-        After inversion of the preconditioner, the values stored in the steps are the increment. This function adds the
-        solution after the previous iteration to arrive at the solution after the current iteration.
-        Note that we also need to put in the initial conditions back in the first step because they will be perturbed by
-        the circular preconditioner.
+        Since we solve for the increment, we need to update the solution between iterations by adding the increment.
 
         Args:
             local_MS_running (list): list of currently running steps
         """
         for S in local_MS_running:
-            for m in range(S.levels[0].sweep.coll.num_nodes + 1):
-                S.levels[0].u[m] = S.levels[0].uold[m] + S.levels[0].u[m]
-            if S.status.first:
-                S.levels[0].u[0] = self.ParaDiag_block_u0
+            for m in range(S.levels[0].sweep.coll.num_nodes):
+                S.levels[0].u[m + 1] += S.levels[0].increment[m]
 
     def prepare_Jacobians(self, local_MS_running):
         # get solutions for constructing average Jacobians
@@ -215,22 +206,22 @@ def it_ParaDiag(self, local_MS_running):
         # communicate average residual for setting up Jacobians for non-linear problems
         self.prepare_Jacobians(local_MS_running)
 
-        # replace the values stored in the steps with the residuals in order to compute the increment
-        self.swap_solution_for_all_at_once_residual(local_MS_running)
+        # compute the all-at-once residual to use as right hand side
+        self.compute_all_at_once_residual(local_MS_running)
 
-        # weighted FFT in time
-        self.FFT_in_time()
+        # weighted FFT of the residual in time
+        self.FFT_in_time(quantity='residual')
 
         # perform local solves of "collocation problems" on the steps (can be done in parallel)
         for S in local_MS_running:
             assert len(S.levels) == 1, 'Multi-level SDC not implemented in ParaDiag'
             S.levels[0].sweep.update_nodes()
 
-        # inverse FFT in time
-        self.iFFT_in_time()
+        # inverse FFT of the increment in time
+        self.iFFT_in_time(quantity='increment')
 
-        # replace the values stored in the steps with the previous solution plus the increment
-        self.swap_increment_for_solution(local_MS_running)
+        # get the next iterate by adding increment to previous iterate
+        self.update_solution(local_MS_running)
 
         for S in local_MS_running:
             for hook in self.hooks:
@@ -438,7 +429,6 @@ def restart_block(self, active_slots, time, u0):
             u0: initial value to distribute across the steps
 
         """
-        self.ParaDiag_block_u0 = u0  # need this for computing residual
 
         for j in range(len(active_slots)):
             # get slot number

pySDC/implementations/sweeper_classes/ParaDiagSweepers.py

@@ -102,7 +102,7 @@ def update_nodes(self):
 
         # perform local solves on the collocation nodes, can be parallelized!
         if self.params.ignore_ic:
-            x1 = self.mat_vec(self.S_inv, [self.level.u[m + 1] for m in range(M)])
+            x1 = self.mat_vec(self.S_inv, [self.level.residual[m] for m in range(M)])
         else:
             x1 = self.mat_vec(self.S_inv, [self.level.u[0] for _ in range(M)])
 
@@ -120,7 +120,10 @@ def update_nodes(self):
 
         # update solution and evaluate right hand side
         for m in range(M):
-            L.u[m + 1] = y[m]
+            if self.params.ignore_ic:
+                L.increment[m] = y[m]
+            else:
+                L.u[m + 1] = y[m]
             if self.params.update_f_evals:
                 L.f[m + 1] = P.eval_f(L.u[m + 1], L.time + L.dt * self.coll.nodes[m])
 

pySDC/tests/test_controllers/test_controller_ParaDiag_nonMPI.py

@@ -261,9 +261,44 @@ def test_ParaDiag_convergence_rate(L, M, N, alpha):
     ), f'Convergence rate {convergence_rate:.2e} exceeds upper bound of {convergence_bound:.2e}!'
 
 
+@pytest.mark.base
+@pytest.mark.parametrize('L', [4, 12])
+@pytest.mark.parametrize('M', [3, 4])
+@pytest.mark.parametrize('N', [1, 2])
+def test_fft(L, M, N):
+    import numpy as np
+    from pySDC.helpers.ParaDiagHelper import get_FFT_matrix
+
+    dt = 1e-2
+    controller, prob = get_composite_collocation_problem(L, M, N, alpha=1e-1, dt=dt, problem='Dahlquist')
+    # generate random data
+    data = np.random.random((L, M, N))
+    data = np.ones((L, M, N))
+
+    for l in range(L):
+        for m in range(M):
+            controller.MS[l].levels[0].residual[m] = prob.u_init
+            controller.MS[l].levels[0].residual[m][:] = data[l, m]
+
+    fft_matrix = get_FFT_matrix(L)
+    controller.apply_matrix(fft_matrix, 'residual')
+    data_fft = np.fft.fft(data, axis=0, norm='ortho')
+
+    for l in range(L):
+        for m in range(M):
+            assert np.allclose(controller.MS[l].levels[0].residual[m], data_fft[l, m])
+
+    controller.apply_matrix(np.conjugate(fft_matrix), 'residual')
+    for l in range(L):
+        for m in range(M):
+            assert np.allclose(controller.MS[l].levels[0].residual[m], data[l, m])
+
+
 if __name__ == '__main__':
-    test_ParaDiag_convergence_rate(4, 3, 1, 1e-4)
+    test_fft(3, 2, 2)
+    # test_ParaDiag_convergence_rate(4, 3, 1, 1e-4)
     # test_ParaDiag_vs_PFASST(4, 3, 2, 'Dahlquist')
     # test_ParaDiag_convergence(4, 3, 1, 1e-4, 'vdp')
     # test_IMEX_ParaDiag_convergence(4, 3, 64, 1e-4)
     # test_ParaDiag_order(3, 3, 1, 1e-4)
+    print('done')

pySDC/tests/test_sweepers/test_ParaDiag_sweepers.py

@@ -75,6 +75,8 @@ def test_direct_solve(M, N, ignore_ic):
         level.u[m] = prob.u_exact(0)
         level.f[m] = prob.eval_f(level.u[m], 0)
 
+    level.sweep.compute_residual()
+
     if ignore_ic:
         level.u[0][:] = None
 
@@ -92,6 +94,8 @@ def test_direct_solve(M, N, ignore_ic):
     u = sp.linalg.spsolve(C_coll, u0.flatten()).reshape(u0.shape)
 
     for m in range(M):
+        if ignore_ic:
+            level.u[m + 1] = level.u[m + 1] + level.increment[m]
         assert np.allclose(u[m], level.u[m + 1])
 
     if not ignore_ic:
@@ -100,4 +104,4 @@ def test_direct_solve(M, N, ignore_ic):
 
 
 if __name__ == '__main__':
-    test_direct_solve(2, 1, False)
+    test_direct_solve(2, 1, True)

pySDC/tutorial/step_9/C_paradiag_in_pySDC.py

@@ -176,5 +176,5 @@ def compare_ParaDiag_and_PFASST(n_steps, problem):
     params = {
         'n_steps': 16,
     }
-    compare_ParaDiag_and_PFASST(**params, problem='advection')
+    # compare_ParaDiag_and_PFASST(**params, problem='advection')
     compare_ParaDiag_and_PFASST(**params, problem='vdp')