Part A: ParaDiag for linear problems#
We start with the linear case, where the composite collocation problem really can be written as a matrix and the whole method is a few lines of linear algebra. It is recommended to view this code side by side with Gaya’s paper on ParaDiag with collocation methods, as the code follows the equations there closely without repeating their explanation.
The test problem is the Dahlquist equation, on \(N = 2\) unknowns, over \(L = 4\) time steps with \(M = 3\) collocation nodes each.
import numpy as np
import scipy.sparse as sp
from pySDC.implementations.problem_classes.TestEquation_0D import testequation0d as problem_class
from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit
from pySDC.implementations.sweeper_classes.ParaDiagSweepers import QDiagonalization
# setup parameters
L = 4 # Number of parallel time steps
M = 3 # Number of collocation nodes
N = 2 # Number of spatial degrees of freedom
alpha = 1e-4 # Circular perturbation parameter
restol = 1e-10 # Residual tolerance for the composite collocation problem
dt = 0.1 # step size
print(f'Running ParaDiag test script with {L} time steps, {M} collocation nodes and {N} spatial degrees of freedom')
# setup pySDC infrastructure for Dahlquist problem and quadrature
prob = problem_class(lambdas=-1.0 * np.ones(shape=(N)), u0=1.0)
sweeper_params = params = {'num_nodes': M, 'quad_type': 'RADAU-RIGHT'}
sweep = generic_implicit(sweeper_params, None)
# Setup a global NumPy array and insert initial conditions in the first step
u = np.zeros((L, M, N), dtype=complex)
u[0, :, :] = prob.u_exact(t=0)
Running ParaDiag test script with 4 time steps, 3 collocation nodes and 2 spatial degrees of freedom
The composite collocation problem#
We set up the matrices of the composite collocation problem, and note their sizes in the comments. The matrix \(E\) propagates the solution of each step to be the initial condition of the next.
# Start with identity matrices (I) of various sizes
I_L = sp.eye(L) # LxL
I_MN = sp.eye((M) * N) # MNxMN
I_N = sp.eye(N) # NxN
I_M = sp.eye(M) # MxM
# E matrix propagates the solution of the steps to be the initial condition for the next step
E = sp.diags(
[
-1.0,
]
* (L - 1),
offsets=-1,
) # LxL
The \(H\) matrix computes the solution at the end of an individual step from the solutions at the collocation nodes. For the RADAU-RIGHT rule we use here, the right node coincides with the end of the interval, so this is simple. We start with the \(M \times M\) matrix \(H_M\) on the node level, and extend it to the spatial dimension with a Kronecker product.
H_M = sp.eye(M).tolil() * 0 # MxM
H_M[:, -1] = 1
H = sp.kron(H_M, I_N) # MNxMN
Now the collocation problem. Note that the Kronecker product of \(Q\) and \(A\) is only possible when there is an \(A\), i.e. when the problem is linear; nonlinear problems follow in Part B.
Q = sweep.coll.Qmat[1:, 1:] # MxM
C_coll = I_MN - dt * sp.kron(Q, prob.A) # MNxMN
# Set up the composite collocation / all-at-once problem
C = (sp.kron(I_L, C_coll) + sp.kron(E, H)).tocsc() # LMNxLMN
With the full composite collocation problem as one large matrix, we can solve it directly for a reference solution. Of course, this is prohibitively expensive for any actual application, and we would never do this in practice.
sol_direct = sp.linalg.spsolve(C, u.flatten()).reshape(u.shape)
The normal time-stepping approach solves the composite collocation problem by forward substitution, one step after the other. This only solves \(MN \times MN\) systems rather than \(LMN \times LMN\) ones, but that is still really expensive in practice, which is why there is SDC, for example.
sol_stepping = u.copy()
for l in range(L):
# solve the current step (sol_stepping[l] currently contains the initial conditions at step l)
sol_stepping[l, :] = sp.linalg.spsolve(C_coll, sol_stepping[l].flatten()).reshape(sol_stepping[l].shape)
# place the solution to the current step as the initial conditions to the next step
if l < L - 1:
sol_stepping[l + 1, ...] = sol_stepping[l, -1, :]
assert np.allclose(sol_stepping, sol_direct)
ParaDiag#
So far, so serial and boring. We now parallelize this with ParaDiag. We solve the composite collocation problem with preconditioned Picard iterations,
where the right-hand side is the residual of the composite collocation problem. The trick behind ParaDiag is to choose the preconditioner \(C_\alpha\) as a time-periodic approximation to \(C\) that can be diagonalized, and therefore inverted in parallel. What changes in \(C_\alpha\) compared to \(C\) is the matrix \(E\), which now also feeds the solution of the last step back into the first one, scaled by \(\alpha\):
E_alpha = sp.diags(
[
-1.0,
]
* (L - 1),
offsets=-1,
).tolil() # LxL
E_alpha[0, -1] = -alpha # make the problem time-periodic
To diagonalize \(C_\alpha\) on the step level, we need to diagonalize \(I_L\) and \(E_\alpha\) simultaneously. Both are \(\alpha\)-circulant matrices, which can be diagonalized simultaneously by a weighted Fourier transform. We set up the weighting matrices of the Fourier transforms, and compute the diagonal entries of the diagonal version \(D_\alpha\) of \(E_\alpha\). We do not set up the preconditioner itself, because we will not use the expanded version.
gamma = alpha ** (-np.arange(L) / L)
J = sp.diags(gamma) # LxL
J_inv = sp.diags(1 / gamma) # LxL
# compute diagonal entries via Fourier transform
D_alpha_diag_vals = np.fft.fft(1 / gamma * E_alpha[:, 0].toarray().flatten(), norm='backward')
Two convenience functions, for matrix-vector products on the steps and for the residual of the composite collocation problem:
def mat_vec(mat, vec):
"""
Matrix vector product
Args:
mat (np.ndarray or scipy.sparse) : Matrix
vec (np.ndarray) : vector
Returns:
np.ndarray: mat @ vec
"""
res = np.zeros_like(vec).astype(complex)
for l in range(vec.shape[0]):
for k in range(vec.shape[0]):
res[l] += mat[l, k] * vec[k]
return res
def residual(_u, u0):
"""
Compute the residual of the composite collocation problem
Args:
_u (np.ndarray): Current iterate
u0 (np.ndarray): Initial conditions
Returns:
np.ndarray: LMN size array with the residual
"""
res = _u * 0j
for l in range(L):
# build step local residual
# communicate initial conditions for each step
if l == 0:
res[l, ...] = u0[l, ...]
else:
res[l, ...] = _u[l - 1, -1, ...]
# evaluate and subtract integral over right hand side functions
f_evals = np.array([prob.eval_f(_u[l, m], 0) for m in range(M)])
Qf = mat_vec(Q, f_evals)
for m in range(M):
res[l, m, ...] -= (_u[l] - dt * Qf)[m]
return res
Parallel across the steps#
First, we parallelize across the \(L\) steps, but solve the collocation problems directly and in serial: a weighted FFT in time, \(L\) independent solves, the inverse transform.
sol_ParaDiag_L = u.copy()
u0 = u.copy()
niter_ParaDiag_L = 0
res = residual(sol_ParaDiag_L, u0)
while np.linalg.norm(res) > restol:
# compute weighted FFT in time to go to diagonal base of C_alpha
x = np.fft.fft(
mat_vec(J_inv.tolil(), res),
axis=0,
norm='ortho',
)
# solve the collocation problems in parallel on the steps
y = np.empty_like(x)
for l in range(L):
# construct local matrix of "collocation problem"
local_matrix = (D_alpha_diag_vals[l] * H + C_coll).tocsc()
# solve local "collocation problem" directly
y[l, ...] = sp.linalg.spsolve(local_matrix, x[l, ...].flatten()).reshape(x[l, ...].shape)
# compute inverse weighted FFT in time to go back from diagonal base of C_alpha
sol_ParaDiag_L += mat_vec(J.tolil(), np.fft.ifft(y, axis=0, norm='ortho'))
# update residual
res = residual(sol_ParaDiag_L, u0)
niter_ParaDiag_L += 1
print(
f'Needed {niter_ParaDiag_L} iterations in parallel across the steps ParaDiag. Stopped at residual {np.linalg.norm(res):.2e}'
)
assert np.allclose(sol_ParaDiag_L, sol_direct)
Needed 3 iterations in parallel across the steps ParaDiag. Stopped at residual 7.38e-13
Parallel across the steps and the nodes#
The work is distributed across \(L\) tasks now, but each of them still solves a perturbed collocation problem directly, which is very expensive. So we additionally diagonalize the quadrature matrix \(Q\), to distribute the work on \(LM\) tasks that solve \(N \times N\) systems each. We rearrange the contribution of \(E_\alpha\) to arrive at a problem \((I - \Delta t\, Q G^{-1} A) u = u_0\). After diagonalizing \(Q G^{-1}\), we can simply use the implicit Euler solves that pySDC problems implement, keeping in mind that they need complex “step sizes”.
We set up \(G\) and \(G^{-1}\), and pySDC sweepers that diagonalize \(Q G^{-1}\). We do not use the sweepers themselves here, only the diagonalization they compute, to make clearer what is going on.
G = [(D_alpha_diag_vals[l] * H_M + I_M).tocsc() for l in range(L)] # MxM
G_inv = [sp.linalg.inv(_G).toarray() for _G in G] # MxM
sweepers = [QDiagonalization(params={**sweeper_params, 'G_inv': _G_inv}, level=None) for _G_inv in G_inv]
sol_ParaDiag = u.copy().astype(complex)
res = residual(sol_ParaDiag, u0)
niter = 0
while np.max(np.abs(residual(sol_ParaDiag, u0))) > restol:
# weighted FFT in time
x = np.fft.fft(
mat_vec(J_inv.tolil(), res),
axis=0,
norm='ortho',
)
# perform local solves of "collocation problems" on the steps in parallel
y = np.empty_like(x)
for l in range(L):
# diagonalize QG^-1 matrix
w, S, S_inv = sweepers[l].w, sweepers[l].S, sweepers[l].S_inv
# perform local solves on the collocation nodes in parallel
x1 = S_inv @ x[l]
x2 = np.empty_like(x1)
for m in range(M):
x2[m, :] = prob.solve_system(rhs=x1[m], factor=w[m] * dt, u0=x1[m], t=0)
z = S @ x2
y[l, ...] = G_inv[l] @ z
# inverse weighted FFT in time
sol_ParaDiag += mat_vec(J.tolil(), np.fft.ifft(y, axis=0, norm='ortho'))
res = residual(sol_ParaDiag, u0)
niter += 1
print(
f'Needed {niter} iterations in parallel and local paradiag with increment formulation, stopped at residual {np.linalg.norm(res):.2e}'
)
Needed 3 iterations in parallel and local paradiag with increment formulation, stopped at residual 7.38e-13
Important things to note
The diagonalization happens across the time steps, not across the collocation nodes; the second variant adds a diagonalization across the nodes on top.
The \(\alpha\)-circulant approximation is what makes the diagonalization possible in the first place.
The checks the tests run: both variants arrive at the direct solution, in the same number of iterations.
assert np.allclose(sol_ParaDiag, sol_direct)
assert np.allclose(niter, niter_ParaDiag_L)