Coverage for pySDC/tutorial/step_9/A_paradiag_for_linear_problems.py: 100%
96 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1# ---
2# jupyter:
3# jupytext:
4# formats: py:percent
5# kernelspec:
6# display_name: Python 3
7# name: python3
8# ---
10# %% [markdown]
11# # Part A: ParaDiag for linear problems
12#
13# We start with the linear case, where the composite collocation problem really can be written as a matrix and the
14# whole method is a few lines of linear algebra. It is recommended to view this code side by side with
15# [Gaya's paper on ParaDiag with collocation methods](https://arxiv.org/abs/2103.12571), as the code follows the
16# equations there closely without repeating their explanation.
17#
18# The test problem is the Dahlquist equation, on $N = 2$ unknowns, over $L = 4$ time steps with $M = 3$ collocation
19# nodes each.
21# %%
22import numpy as np
23import scipy.sparse as sp
25from pySDC.implementations.problem_classes.TestEquation_0D import testequation0d as problem_class
26from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit
27from pySDC.implementations.sweeper_classes.ParaDiagSweepers import QDiagonalization
29# setup parameters
30L = 4 # Number of parallel time steps
31M = 3 # Number of collocation nodes
32N = 2 # Number of spatial degrees of freedom
33alpha = 1e-4 # Circular perturbation parameter
34restol = 1e-10 # Residual tolerance for the composite collocation problem
35dt = 0.1 # step size
37print(f'Running ParaDiag test script with {L} time steps, {M} collocation nodes and {N} spatial degrees of freedom')
39# setup pySDC infrastructure for Dahlquist problem and quadrature
40prob = problem_class(lambdas=-1.0 * np.ones(shape=(N)), u0=1.0)
41sweeper_params = params = {'num_nodes': M, 'quad_type': 'RADAU-RIGHT'}
42sweep = generic_implicit(sweeper_params, None)
44# Setup a global NumPy array and insert initial conditions in the first step
45u = np.zeros((L, M, N), dtype=complex)
46u[0, :, :] = prob.u_exact(t=0)
48# %% [markdown]
49# ## The composite collocation problem
50#
51# We set up the matrices of the composite collocation problem, and note their sizes in the comments. The matrix $E$
52# propagates the solution of each step to be the initial condition of the next.
54# %%
55# Start with identity matrices (I) of various sizes
56I_L = sp.eye(L) # LxL
57I_MN = sp.eye((M) * N) # MNxMN
58I_N = sp.eye(N) # NxN
59I_M = sp.eye(M) # MxM
61# E matrix propagates the solution of the steps to be the initial condition for the next step
62E = sp.diags(
63 [
64 -1.0,
65 ]
66 * (L - 1),
67 offsets=-1,
68) # LxL
70# %% [markdown]
71# The $H$ matrix computes the solution at the end of an individual step from the solutions at the collocation nodes.
72# For the RADAU-RIGHT rule we use here, the right node coincides with the end of the interval, so this is simple. We
73# start with the $M \times M$ matrix $H_M$ on the node level, and extend it to the spatial dimension with a Kronecker
74# product.
76# %%
77H_M = sp.eye(M).tolil() * 0 # MxM
78H_M[:, -1] = 1
79H = sp.kron(H_M, I_N) # MNxMN
81# %% [markdown]
82# Now the collocation problem. Note that the Kronecker product of $Q$ and $A$ is only possible when there is an $A$,
83# i.e. when the problem is linear; nonlinear problems follow in [Part B](B_paradiag_for_nonlinear_problems).
85# %%
86Q = sweep.coll.Qmat[1:, 1:] # MxM
87C_coll = I_MN - dt * sp.kron(Q, prob.A) # MNxMN
89# Set up the composite collocation / all-at-once problem
90C = (sp.kron(I_L, C_coll) + sp.kron(E, H)).tocsc() # LMNxLMN
92# %% [markdown]
93# With the full composite collocation problem as one large matrix, we can solve it directly for a reference solution.
94# Of course, this is prohibitively expensive for any actual application, and we would never do this in practice.
96# %%
97sol_direct = sp.linalg.spsolve(C, u.flatten()).reshape(u.shape)
99# %% [markdown]
100# The normal time-stepping approach solves the composite collocation problem by forward substitution, one step after
101# the other. This only solves $MN \times MN$ systems rather than $LMN \times LMN$ ones, but that is still really
102# expensive in practice, which is why there is SDC, for example.
104# %%
105sol_stepping = u.copy()
106for l in range(L):
107 # solve the current step (sol_stepping[l] currently contains the initial conditions at step l)
108 sol_stepping[l, :] = sp.linalg.spsolve(C_coll, sol_stepping[l].flatten()).reshape(sol_stepping[l].shape)
110 # place the solution to the current step as the initial conditions to the next step
111 if l < L - 1:
112 sol_stepping[l + 1, ...] = sol_stepping[l, -1, :]
114assert np.allclose(sol_stepping, sol_direct)
116# %% [markdown]
117# ## ParaDiag
118#
119# So far, so serial and boring. We now parallelize this with ParaDiag. We solve the composite collocation problem
120# with preconditioned Picard iterations,
121#
122# $$
123# C_\alpha \delta = u_0 - C u^k, \qquad u^{k+1} = u^k + \delta ,
124# $$
125#
126# where the right-hand side is the residual of the composite collocation problem. The trick behind ParaDiag is to
127# choose the preconditioner $C_\alpha$ as a time-periodic approximation to $C$ that can be diagonalized, and
128# therefore inverted in parallel. What changes in $C_\alpha$ compared to $C$ is the matrix $E$, which now also feeds
129# the solution of the last step back into the first one, scaled by $\alpha$:
131# %%
132E_alpha = sp.diags(
133 [
134 -1.0,
135 ]
136 * (L - 1),
137 offsets=-1,
138).tolil() # LxL
139E_alpha[0, -1] = -alpha # make the problem time-periodic
141# %% [markdown]
142# To diagonalize $C_\alpha$ on the step level, we need to diagonalize $I_L$ and $E_\alpha$ simultaneously. Both are
143# $\alpha$-circulant matrices, which can be diagonalized simultaneously by a weighted Fourier transform. We set up
144# the weighting matrices of the Fourier transforms, and compute the diagonal entries of the diagonal version
145# $D_\alpha$ of $E_\alpha$. We do not set up the preconditioner itself, because we will not use the expanded version.
147# %%
148gamma = alpha ** (-np.arange(L) / L)
149J = sp.diags(gamma) # LxL
150J_inv = sp.diags(1 / gamma) # LxL
152# compute diagonal entries via Fourier transform
153D_alpha_diag_vals = np.fft.fft(1 / gamma * E_alpha[:, 0].toarray().flatten(), norm='backward')
155# %% [markdown]
156# Two convenience functions, for matrix-vector products on the steps and for the residual of the composite
157# collocation problem:
160# %%
161def mat_vec(mat, vec):
162 """
163 Matrix vector product
165 Args:
166 mat (np.ndarray or scipy.sparse) : Matrix
167 vec (np.ndarray) : vector
169 Returns:
170 np.ndarray: mat @ vec
171 """
172 res = np.zeros_like(vec).astype(complex)
173 for l in range(vec.shape[0]):
174 for k in range(vec.shape[0]):
175 res[l] += mat[l, k] * vec[k]
176 return res
179def residual(_u, u0):
180 """
181 Compute the residual of the composite collocation problem
183 Args:
184 _u (np.ndarray): Current iterate
185 u0 (np.ndarray): Initial conditions
187 Returns:
188 np.ndarray: LMN size array with the residual
189 """
190 res = _u * 0j
191 for l in range(L):
192 # build step local residual
194 # communicate initial conditions for each step
195 if l == 0:
196 res[l, ...] = u0[l, ...]
197 else:
198 res[l, ...] = _u[l - 1, -1, ...]
200 # evaluate and subtract integral over right hand side functions
201 f_evals = np.array([prob.eval_f(_u[l, m], 0) for m in range(M)])
202 Qf = mat_vec(Q, f_evals)
203 for m in range(M):
204 res[l, m, ...] -= (_u[l] - dt * Qf)[m]
206 return res
209# %% [markdown]
210# ### Parallel across the steps
211#
212# First, we parallelize across the $L$ steps, but solve the collocation problems directly and in serial: a weighted
213# FFT in time, $L$ independent solves, the inverse transform.
215# %%
216sol_ParaDiag_L = u.copy()
217u0 = u.copy()
218niter_ParaDiag_L = 0
220res = residual(sol_ParaDiag_L, u0)
221while np.linalg.norm(res) > restol:
222 # compute weighted FFT in time to go to diagonal base of C_alpha
223 x = np.fft.fft(
224 mat_vec(J_inv.tolil(), res),
225 axis=0,
226 norm='ortho',
227 )
229 # solve the collocation problems in parallel on the steps
230 y = np.empty_like(x)
231 for l in range(L):
232 # construct local matrix of "collocation problem"
233 local_matrix = (D_alpha_diag_vals[l] * H + C_coll).tocsc()
235 # solve local "collocation problem" directly
236 y[l, ...] = sp.linalg.spsolve(local_matrix, x[l, ...].flatten()).reshape(x[l, ...].shape)
238 # compute inverse weighted FFT in time to go back from diagonal base of C_alpha
239 sol_ParaDiag_L += mat_vec(J.tolil(), np.fft.ifft(y, axis=0, norm='ortho'))
241 # update residual
242 res = residual(sol_ParaDiag_L, u0)
243 niter_ParaDiag_L += 1
244print(
245 f'Needed {niter_ParaDiag_L} iterations in parallel across the steps ParaDiag. Stopped at residual {np.linalg.norm(res):.2e}'
246)
247assert np.allclose(sol_ParaDiag_L, sol_direct)
249# %% [markdown]
250# ### Parallel across the steps and the nodes
251#
252# The work is distributed across $L$ tasks now, but each of them still solves a perturbed collocation problem
253# directly, which is very expensive. So we additionally diagonalize the quadrature matrix $Q$, to distribute the work
254# on $LM$ tasks that solve $N \times N$ systems each. We rearrange the contribution of $E_\alpha$ to arrive at a
255# problem $(I - \Delta t\, Q G^{-1} A) u = u_0$. After diagonalizing $Q G^{-1}$, we can simply use the implicit Euler
256# solves that pySDC problems implement, keeping in mind that they need complex "step sizes".
257#
258# We set up $G$ and $G^{-1}$, and pySDC sweepers that diagonalize $Q G^{-1}$. We do not use the sweepers themselves
259# here, only the diagonalization they compute, to make clearer what is going on.
261# %%
262G = [(D_alpha_diag_vals[l] * H_M + I_M).tocsc() for l in range(L)] # MxM
263G_inv = [sp.linalg.inv(_G).toarray() for _G in G] # MxM
264sweepers = [QDiagonalization(params={**sweeper_params, 'G_inv': _G_inv}, level=None) for _G_inv in G_inv]
267sol_ParaDiag = u.copy().astype(complex)
268res = residual(sol_ParaDiag, u0)
269niter = 0
270while np.max(np.abs(residual(sol_ParaDiag, u0))) > restol:
272 # weighted FFT in time
273 x = np.fft.fft(
274 mat_vec(J_inv.tolil(), res),
275 axis=0,
276 norm='ortho',
277 )
279 # perform local solves of "collocation problems" on the steps in parallel
280 y = np.empty_like(x)
281 for l in range(L):
283 # diagonalize QG^-1 matrix
284 w, S, S_inv = sweepers[l].w, sweepers[l].S, sweepers[l].S_inv
286 # perform local solves on the collocation nodes in parallel
287 x1 = S_inv @ x[l]
288 x2 = np.empty_like(x1)
289 for m in range(M):
290 x2[m, :] = prob.solve_system(rhs=x1[m], factor=w[m] * dt, u0=x1[m], t=0)
291 z = S @ x2
292 y[l, ...] = G_inv[l] @ z
294 # inverse weighted FFT in time
295 sol_ParaDiag += mat_vec(J.tolil(), np.fft.ifft(y, axis=0, norm='ortho'))
297 res = residual(sol_ParaDiag, u0)
298 niter += 1
299print(
300 f'Needed {niter} iterations in parallel and local paradiag with increment formulation, stopped at residual {np.linalg.norm(res):.2e}'
301)
303# %% [markdown]
304# :::{admonition} Important things to note
305# - The diagonalization happens across the time steps, not across the collocation nodes; the second variant adds a
306# diagonalization across the nodes on top.
307# - The $\alpha$-circulant approximation is what makes the diagonalization possible in the first place.
308# :::
309#
310# The checks the tests run: both variants arrive at the direct solution, in the same number of iterations.
312# %%
313assert np.allclose(sol_ParaDiag, sol_direct)
314assert np.allclose(niter, niter_ParaDiag_L)