Step-9: ParaDiag¶
ParaDiag is a parallel-in-time method of a rather different flavour than PFASST. Instead of iterating on a hierarchy of levels and passing information forward step by step, it diagonalizes the “top layer” of Kronecker products that makes up the composite collocation problem. After the diagonalization, the collocation problems on the individual steps decouple and can be solved concurrently, which is where the parallelism comes from.
The price is an approximation: the time-stepping matrix is replaced by an \(\alpha\)-circulant one, which is diagonalizable by a weighted Fourier transform. The outer iteration then corrects for that perturbation, and \(\alpha\) trades approximation quality against the conditioning of the diagonalization.
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.
Important things to note:
The diagonalization happens across the time-steps, not across the collocation nodes.
The \(\alpha\)-circulant approximation is what makes the diagonalization possible in the first place.
Full code: pySDC/tutorial/step_9/A_paradiag_for_linear_problems.py
"""
This script introduces ParaDiag for linear problems.
It is recommended to view this code side by side with `Gaya's paper on ParaDiag with collocation methods
<https://arxiv.org/abs/2103.12571>`_ as the code is close to the equations presented there but offers no explanations
about them.
"""
import numpy as np
import scipy.sparse as sp
import sys
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 output
out_file = open('data/step_9_A_out.txt', 'w')
def my_print(*args, **kwargs):
for output in [sys.stdout, out_file]:
print(*args, **kwargs, file=output)
# 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
my_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)
# setup matrices for composite collocation problem. We note the sizes of the matrices in comments after generating them.
# 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 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 building the MxM matrix H_M on the node level and then extend 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
"""
Set up collocation problem.
Note that the Kronecker product from Q and A is only possible when there is an A, i.e. when the problem is linear.
We will discuss non-linear problems in later steps in this tutorial
"""
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
"""
Now that we have the full composite collocation problem as one large matrix, we can just solve it directly to get a reference solution.
Of course, this is prohibitively expensive for any actual application and we would never want to do this in practice.
"""
sol_direct = sp.linalg.spsolve(C, u.flatten()).reshape(u.shape)
"""
The normal time-stepping approach is to solve the composite collocation problem with forward substitution
"""
sol_stepping = u.copy()
for l in range(L):
"""
Solve the current step (sol_stepping[l] currently contains the initial conditions at step l)
Here, we only solve MNxMN systems rather than LMNxLMN systems. This is still really expensive in practice, which is why there is SDC, for example.
"""
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)
"""
So far, so serial and boring. We will now parallelize this using ParaDiag.
We will solve the composite collocation problem using preconditioned Picard iterations:
C_alpha delta = u_0 - Cu^k = < residual of the composite collocation problem >
u^{k+1} = u^k + delta
The trick behind ParaDiag is to choose the preconditioner C_alpha to be a time-periodic approximation to C that can be diagonalized and therefore inverted in parallel.
What we change in C_alpha compared to C is the E matrix that propagates the solutions between steps, which we amend to feed the solution to the last step back into the first step.
"""
E_alpha = sp.diags(
[
-1.0,
]
* (L - 1),
offsets=-1,
).tolil() # LxL
E_alpha[0, -1] = -alpha # make the problem time-periodic
"""
In order to diagonalize C_alpha, on the step level, we need to diagonalize I_L and E_alpha simultaneously.
I_L and E_alpha are alpha-circular matrices which can be simultaneously diagonalized by a weighted Fourier transform.
We start by setting the weighting matrices for the Fourier transforms and then compute the diagonal entries of the diagonal version D_alpha of E_alpha.
We refrain from actually setting up the preconditioner because we will not use the expanded version here.
"""
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')
"""
We need some convenience functions for computing matrix vector multiplication and the composite collocation problem residual here
"""
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)[-1]
res[l, m, ...] -= (_u[l] - dt * Qf)[m]
# res[l, m, ...] -= np.mean((_u[l] - dt * Qf), axis=0)
return res
"""
We will start with ParaDiag where we parallelize across the L steps but solve the collocation problems directly in serial.
"""
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
my_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)
"""
While we have distributed the work across L tasks, we are still solving perturbed collocation problems directly on a single task here.
This is very expensive, and we will now additionally diagonalize the quadrature matrix Q in order to distribute the work on LM tasks, where we solve NxN systems each.
We rearrange the contribution of E_alpha to arrive at a problem (I - dtQG^{-1}A)u = u0.
After diagonalizing QG^{-1}, we can simply utilize the Euler solves that are implemented in pySDC, but need to keep in mind that complex valued "step sizes" are required.
We start by setting up the G and G^{-1} matrices. Then we will setup pySDC sweepers that solve QG^{-1} with diagonalization.
Here, we will not use the sweepers, but just the diagonalization computed there in order to make more clear 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
my_print(
f'Needed {niter} iterations in parallel and local paradiag with increment formulation, stopped at residual {np.linalg.norm(res):.2e}'
)
assert np.allclose(sol_ParaDiag, sol_direct)
assert np.allclose(niter, niter_ParaDiag_L)
Results:
Running ParaDiag test script with 4 time steps, 3 collocation nodes and 2 spatial degrees of freedom
Needed 3 iterations in parallel across the steps ParaDiag. Stopped at residual 7.38e-13
Needed 3 iterations in parallel and local paradiag with increment formulation, stopped at residual 7.38e-13
Part B: ParaDiag for nonlinear problems¶
For nonlinear problems the composite collocation problem cannot be written as a matrix, so the diagonalization needs a linear operator to work with. This part shows the two ways out: IMEX splitting, where only the linear implicit part enters the preconditioner, and averaging the Jacobian across the steps.
Important things to note:
Averaging the Jacobian requires communicating the average solution, which is why
average_jacobianis off by default for linear problems.We do a single Newton iteration per ParaDiag iteration, so the number of Newton iterations per node equals the number of ParaDiag iterations.
Full code: pySDC/tutorial/step_9/B_paradiag_for_nonlinear_problems.py
"""
This script introduces ParaDiag for nonlinear problems with the van der Pol oscillator as an example.
ParaDiag works by diagonalizing the "top layer" of Kronecker products that make up the circularized composite
collocation problem.
However, in nonlinear problems, the problem cannot be written as a matrix and therefore we cannot write the composite
collocation problem as a matrix.
There are two approaches for dealing with this. We can do IMEX splitting, where we treat only the linear part implicitly.
The ParaDiag preconditioner is then only made up of the linear implicit part and we can again write this as a matrix and
do the diagonalization just like for linear problems. The non-linear part then comes in via the residual on the right
hand side.
The second approach is to average Jacobians. The non-linear problems are solved with a Newton scheme, where the Jacobian
matrix is computed based on the current solution and then inverted in each Newton iteration. In order to write the
ParaDiag preconditioner as a matrix with Kronecker products and then only diagonalize the outermost part, we need to
have the same Jacobian on all steps.
The ParaDiag iteration then proceeds as follows:
- (1) Compute residual of composite collocation problem
- (2) Average the solution across the steps and nodes as preparation for computing the average Jacobian
- (3) Weighted FFT in time to diagonalize E_alpha
- (4) Solve for the increment by inverting the averaged Jacobian from (2) on the subproblems on the different steps
and nodes.
- (5) Weighted iFFT in time
- (6) Increment solution
As IMEX ParaDiag is a trivial extension of ParaDiag for linear problems, we focus on the second approach here.
"""
import numpy as np
import scipy.sparse as sp
import sys
from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit as sweeper_class
from pySDC.implementations.problem_classes.Van_der_Pol_implicit import vanderpol
# setup output
out_file = open('data/step_9_B_out.txt', 'w')
def my_print(*args, **kwargs):
for output in [sys.stdout, out_file]:
print(*args, **kwargs, file=output)
# setup parameters
L = 4
M = 3
alpha = 1e-4
restol = 1e-8
dt = 0.1
# setup infrastructure
prob = vanderpol(newton_maxiter=1, mu=1e0, crash_at_maxiter=False)
N = prob.init[0]
# make problem work on complex data
prob.init = tuple([*prob.init[:2]] + [np.dtype('complex128')])
# setup global solution array
u = np.zeros((L, M, N), dtype=complex)
# setup collocation problem
sweep = sweeper_class({'num_nodes': M, 'quad_type': 'RADAU-RIGHT'}, None)
# initial conditions
u[0, :, :] = prob.u_exact(t=0)
my_print(
f'Running ParaDiag test script for van der Pol with mu={prob.mu} and {L} time steps and {M} collocation nodes.'
)
"""
Setup matrices that make up the composite collocation problem. We do not set up the full composite collocation problem
here, however. See https://arxiv.org/abs/2103.12571 for the meaning of the matrices.
"""
I_M = sp.eye(M)
H_M = sp.eye(M).tolil() * 0
H_M[:, -1] = 1
Q = sweep.coll.Qmat[1:, 1:]
E_alpha = sp.diags(
[
-1.0,
]
* (L - 1),
offsets=-1,
).tolil()
E_alpha[0, -1] = -alpha
gamma = alpha ** (-np.arange(L) / L)
D_alpha_diag_vals = np.fft.fft(1 / gamma * E_alpha[:, 0].toarray().flatten(), norm='backward')
J = sp.diags(gamma)
J_inv = sp.diags(1 / gamma)
G = [(D_alpha_diag_vals[l] * H_M + I_M).tocsc() for l in range(L)] # MxM
# prepare diagonalization of QG^{-1}
w = []
S = []
S_inv = []
for l in range(L):
# diagonalize QG^-1 matrix
if M > 1:
_w, _S = np.linalg.eig(Q @ sp.linalg.inv(G[l]).toarray())
else:
_w, _S = np.linalg.eig(Q / (G[l].toarray()))
_S_inv = np.linalg.inv(_S)
w.append(_w)
S.append(_S)
S_inv.append(_S_inv)
"""
Setup functions for computing matrix-vector productions on the steps and for computing 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)
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)
res[l, ...] -= _u[l] - dt * Qf
return res
# do ParaDiag
sol_paradiag = u.copy() * 0j
u0 = u.copy()
niter = 0
res = residual(sol_paradiag, u0)
while np.max(np.abs(res)) > restol:
# compute all-at-once residual
res = residual(sol_paradiag, u0)
# compute solution averaged across the L steps and M nodes. This is the difference to ParaDiag for linear problems.
u_avg = prob.u_init
u_avg[:] = np.mean(sol_paradiag, axis=(0, 1))
# weighted FFT in time
x = np.fft.fft(mat_vec(J_inv.toarray(), res), axis=0)
# perform local solves of "collocation problems" on the steps in parallel
y = np.empty_like(x)
for l in range(L):
# perform local solves on the collocation nodes in parallel
x1 = S_inv[l] @ x[l]
x2 = np.empty_like(x1)
for m in range(M):
x2[m, :] = prob.solve_jacobian(x1[m], w[l][m] * dt, u=u_avg, t=l * dt)
z = S[l] @ x2
y[l, ...] = sp.linalg.spsolve(G[l], z)
# inverse FFT in time and increment
sol_paradiag += mat_vec(J.toarray(), np.fft.ifft(y, axis=0))
res = residual(sol_paradiag, u0)
niter += 1
assert niter < 99, 'ParaDiag did not converge for nonlinear problem!'
my_print(f'Needed {niter} ParaDiag iterations, stopped at residual {np.max(np.abs(res)):.2e}')
Results:
Running ParaDiag test script for van der Pol with mu=1.0 and 4 time steps and 3 collocation nodes.
Needed 5 ParaDiag iterations, stopped at residual 7.54e-10
Part C: ParaDiag in pySDC¶
Here we leave the hand-written linear algebra behind and set ParaDiag up through pySDC’s controllers, comparing it to single-level PFASST in Jacobi mode and to serial time stepping. Both schemes are used without any optimization, so please refrain from computing parallel efficiency from these numbers.
Important things to note:
ParaDiag needs its own sweeper (
QDiagonalization) and its own controller.The solution becomes complex, because the diagonalization is.
ParaDiag converges in very few iterations for the hyperbolic advection example, where PFASST struggles, and the picture reverses for the van der Pol oscillator.
Full code: pySDC/tutorial/step_9/C_paradiag_in_pySDC.py
"""
This script shows how to setup ParaDiag in pySDC for two examples and compares performance to single-level PFASST in
Jacobi mode and serial time stepping.
In PFASST, we use a diagonal preconditioner, which allows for the same amount of parallelism as ParaDiag.
We show iteration counts per step here, but both schemes have further concurrency across the nodes.
We have a linear advection example, discretized with finite differences, where ParaDiag converges in very few iterations.
PFASST, on the hand, needs a lot more iterations for this hyperbolic problem.
Note that we did not optimize either setup. With different choice of alpha in ParaDiag, or inexactness and coarsening
in PFASST, both schemes could be improved significantly.
Second is the nonlinear van der Pol oscillator. We choose the mu parameter such that the problem is not overly stiff.
Here, ParaDiag needs many iterations compared to PFASST, but remember that we only perform one Newton iteration per
ParaDiag iteration. So per node, the number of Newton iterations is equal to the number of ParaDiag iterations.
In PFASST, on the other hand, we solve the systems to some accuracy and allow more iterations. Here, ParaDiag needs
fewer Newton iterations per step in total, leaving it with greater speedup. Again, inexactness could improve PFASST.
This script is not meant to show that one parallelization scheme is better than the other. It does, however, demonstrate
that both schemes, without optimization, need fewer iterations per task than serial time stepping. Kindly refrain from
computing parallel efficiency for these examples, however. ;)
"""
import numpy as np
import sys
from pySDC.helpers.stats_helper import get_sorted
# prepare output
out_file = open('data/step_9_C_out.txt', 'w')
def my_print(*args, **kwargs):
for output in [sys.stdout, out_file]:
print(*args, **kwargs, file=output)
def get_description(problem='advection', mode='ParaDiag'):
level_params = {}
level_params['dt'] = 0.1
level_params['restol'] = 1e-6
sweeper_params = {}
sweeper_params['quad_type'] = 'RADAU-RIGHT'
sweeper_params['num_nodes'] = 3
sweeper_params['initial_guess'] = 'copy'
if mode == 'ParaDiag':
from pySDC.implementations.sweeper_classes.ParaDiagSweepers import QDiagonalization as sweeper_class
# we only want to use the averaged Jacobian and do only one Newton iteration per ParaDiag iteration!
else:
from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit as sweeper_class
# need diagonal preconditioner for same concurrency as ParaDiag
sweeper_params['QI'] = 'MIN-SR-S'
if problem == 'advection':
from pySDC.implementations.problem_classes.AdvectionEquation_ND_FD import advectionNd as problem_class
problem_params = {'nvars': 64, 'order': 8, 'c': 1, 'solver_type': 'GMRES', 'lintol': 1e-8}
elif problem == 'vdp':
from pySDC.implementations.problem_classes.Van_der_Pol_implicit import vanderpol as problem_class
# need to not raise an error when Newton has not converged because we do only one iteration
problem_params = {'newton_maxiter': 99, 'crash_at_maxiter': False, 'mu': 1, 'newton_tol': 1e-9}
step_params = {}
step_params['maxiter'] = 99
description = {}
description['problem_class'] = problem_class
description['problem_params'] = problem_params
description['sweeper_class'] = sweeper_class
description['sweeper_params'] = sweeper_params
description['level_params'] = level_params
description['step_params'] = step_params
return description
def get_controller_params(problem='advection', mode='ParaDiag'):
from pySDC.implementations.hooks.log_errors import LogGlobalErrorPostRun
from pySDC.implementations.hooks.log_work import LogWork, LogSDCIterations
controller_params = {}
controller_params['logger_level'] = 30
controller_params['hook_class'] = [LogGlobalErrorPostRun, LogWork, LogSDCIterations]
if mode == 'ParaDiag':
controller_params['alpha'] = 1e-4
# For nonlinear problems, we need to communicate the average solution, which allows to compute the average
# Jacobian locally. For linear problems, we do not want the extra communication.
if problem == 'advection':
controller_params['average_jacobian'] = False
elif problem == 'vdp':
controller_params['average_jacobian'] = True
else:
# We do Block-Jacobi multi-step SDC here. It's a bit silly but it's better for comparing "speedup"
controller_params['mssdc_jac'] = True
return controller_params
def run_problem(
n_steps=4,
problem='advection',
mode='ParaDiag',
):
if mode == 'ParaDiag':
from pySDC.implementations.controller_classes.controller_ParaDiag_nonMPI import (
controller_ParaDiag_nonMPI as controller_class,
)
else:
from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI as controller_class
if mode == 'serial':
num_procs = 1
else:
num_procs = n_steps
description = get_description(problem, mode)
controller_params = get_controller_params(problem, mode)
controller = controller_class(num_procs=num_procs, description=description, controller_params=controller_params)
for S in controller.MS:
S.levels[0].prob.init = tuple([*S.levels[0].prob.init[:2]] + [np.dtype('complex128')])
P = controller.MS[0].levels[0].prob
t0 = 0.0
uinit = P.u_exact(t0)
uend, stats = controller.run(u0=uinit, t0=t0, Tend=n_steps * controller.MS[0].levels[0].dt)
return uend, stats
def compare_ParaDiag_and_PFASST(n_steps, problem):
my_print(f'Running {problem} with {n_steps} steps')
uend_PD, stats_PD = run_problem(n_steps, problem, mode='ParaDiag')
uend_PF, stats_PF = run_problem(n_steps, problem, mode='PFASST')
uend_S, stats_S = run_problem(n_steps, problem, mode='serial')
assert np.allclose(uend_PD, uend_PF)
assert np.allclose(uend_S, uend_PD)
assert (
abs(uend_PD - uend_PF) > 0
) # two different iterative methods should not give identical results for non-zero tolerance
k_PD = get_sorted(stats_PD, type='k')
k_PF = get_sorted(stats_PF, type='k')
my_print(
f'Needed {max(me[1] for me in k_PD)} ParaDiag iterations and {max(me[1] for me in k_PF)} single-level PFASST iterations'
)
if problem == 'advection':
k_GMRES_PD = get_sorted(stats_PD, type='work_GMRES')
k_GMRES_PF = get_sorted(stats_PF, type='work_GMRES')
k_GMRES_S = get_sorted(stats_S, type='work_GMRES')
my_print(
f'Maximum GMRES iterations on each step: {max(me[1] for me in k_GMRES_PD)} in ParaDiag, {max(me[1] for me in k_GMRES_PF)} in single-level PFASST and {sum(me[1] for me in k_GMRES_S)} total GMRES iterations in serial'
)
elif problem == 'vdp':
k_Jac_PD = get_sorted(stats_PD, type='work_jacobian_solves')
k_Jac_PF = get_sorted(stats_PF, type='work_jacobian_solves')
k_Jac_S = get_sorted(stats_S, type='work_jacobian_solves')
my_print(
f'Maximum Jacabian solves on each step: {max(me[1] for me in k_Jac_PD)} in ParaDiag, {max(me[1] for me in k_Jac_PF)} in single-level PFASST and {sum(me[1] for me in k_Jac_S)} total Jacobian solves in serial'
)
my_print()
if __name__ == '__main__':
out_file = open('data/step_9_C_out.txt', 'w')
params = {
'n_steps': 16,
}
# compare_ParaDiag_and_PFASST(**params, problem='advection')
compare_ParaDiag_and_PFASST(**params, problem='vdp')
Results:
Running advection with 16 steps
Needed 3 ParaDiag iterations and 35 single-level PFASST iterations
Maximum GMRES iterations on each step: 94 in ParaDiag, 861 in single-level PFASST and 1951 total GMRES iterations in serial
Running vdp with 16 steps
Needed 10 ParaDiag iterations and 24 single-level PFASST iterations
Maximum Jacabian solves on each step: 30 in ParaDiag, 143 in single-level PFASST and 233 total Jacobian solves in serial
Part D: MPI-parallel ParaDiag¶
Parts A to C all ran ParaDiag with the “virtually parallel” controller, which keeps every step in a single process.
That is what you want while developing, but it does not actually run in parallel.
This part uses controller_ParaDiag_MPI instead, with one time-step per rank and the communicator spanning the block that is diagonalized.
Nothing about the method changes: the description and the controller parameters are the same, only the controller class differs.
We always integrate the same total number of time-steps and only vary how many of them run in parallel. With four steps in total and a block size of one, two or four, the controller has to window through four, two or one block respectively. Windowing works the same way in both controllers, so we run each block size with the MPI controller and the virtually parallel one and compare.
Important things to note:
All steps of a block iterate together. In PFASST an early step can converge and drop out, which is what makes it pipelined; ParaDiag cannot do that, because the transform in time needs every step. A step that stopped early would leave the others waiting.
Consequently the block is always full. If the end time does not divide into whole blocks, ParaDiag solves past it rather than truncating, and says so.
Neither the iteration counts nor the error depend on how the time domain is split into blocks, which is what we check here. Spreading a block over more processes must not change the method, and windowing through more blocks must not either.
alphamay also be a list or a callable of the iteration index, if you want to start with a well-conditioned value and tighten it later.
Full code: pySDC/tutorial/step_9/D_paradiag_MPI.py
"""
This script shows how to run ParaDiag with actual MPI parallelism across the time-steps.
Part C ran ParaDiag with the "virtually parallel" controller, which holds all steps in one process and
is what you want for developing and debugging. Here we use the MPI controller instead: one time-step
per rank, with the communicator spanning the block that ParaDiag diagonalizes across.
The point of this part is that nothing about the *method* changes. The description and the controller
parameters are the same ones Part C would use; only the controller class differs. So you can develop a
setup serially and then run it in parallel without touching it.
We always integrate the same total number of time-steps and only vary how many of them are done in
parallel. With four steps in total and a block size of one, two or four, that means four, two or one
block respectively, so the controller has to window through the time domain block by block. Windowing
works the same way in both controllers, which is what we check: for every block size we run the MPI
controller and the virtually parallel one and compare.
Two properties of ParaDiag are worth keeping in mind when going parallel, because they are different
from PFASST:
- All steps of a block have to iterate together. In PFASST an early step can converge and drop out of
the iteration, which is what makes it pipelined. ParaDiag cannot do that: the transform in time
needs every step, so a step that stopped early would leave the others waiting forever.
- The block is therefore always full. If the end time does not divide into whole blocks, ParaDiag
solves past it rather than truncating, and says so.
"""
import os
import subprocess
# we always do this many time-steps in total, no matter how many of them run in parallel
num_steps_total = 4
def get_description():
"""
Set up the same advection problem as in Part C.
Returns:
dict: the description for the ParaDiag controller
"""
from pySDC.implementations.problem_classes.AdvectionEquation_ND_FD import advectionNd
from pySDC.implementations.sweeper_classes.ParaDiagSweepers import QDiagonalization
level_params = {}
level_params['dt'] = 0.1
level_params['restol'] = 1e-6
sweeper_params = {}
sweeper_params['quad_type'] = 'RADAU-RIGHT'
sweeper_params['num_nodes'] = 3
sweeper_params['initial_guess'] = 'copy'
# Part C uses GMRES here to count linear solver work. We only care about the parallelism, and the
# complex shifted systems ParaDiag produces are hard for GMRES, so we solve them directly instead.
problem_params = {'nvars': 64, 'order': 8, 'c': 1, 'solver_type': 'direct'}
step_params = {}
step_params['maxiter'] = 99
description = {}
description['problem_class'] = advectionNd
description['problem_params'] = problem_params
description['sweeper_class'] = QDiagonalization
description['sweeper_params'] = sweeper_params
description['level_params'] = level_params
description['step_params'] = step_params
return description
def get_controller_params():
"""
Set up controller parameters for ParaDiag.
`alpha` may also be a list or a callable of the iteration index if you want to start with a
well-conditioned value and tighten it later. We keep it fixed here.
Returns:
dict: the controller parameters
"""
controller_params = {}
controller_params['logger_level'] = 30
controller_params['alpha'] = 1e-4
# the advection problem is linear, so we do not need the extra communication for average Jacobians
controller_params['average_jacobian'] = False
return controller_params
def format_result(mode, block_size, niter, error):
"""
One line of output, in the same shape for both controllers so they can be compared.
Args:
mode (str): 'MPI' or 'virtual'
block_size (int): number of time-steps done in parallel
niter (list): number of iterations of each step
error (float): error against the exact solution at the end
Returns:
str: the formatted line
"""
num_blocks = num_steps_total // block_size
return (
f'{mode:>7s}: block size {block_size}, {num_blocks} block(s) of {block_size} step(s), '
f'iterations {niter}, error {error:.4e}'
)
def run_virtual(block_size):
"""
Run the same setup with the virtually parallel controller.
Args:
block_size (int): number of time-steps in one block
Returns:
tuple: the end value, the iteration counts and the error
"""
import numpy as np
from pySDC.helpers.stats_helper import get_sorted
from pySDC.implementations.controller_classes.controller_ParaDiag_nonMPI import controller_ParaDiag_nonMPI
controller_params = get_controller_params()
controller_params['mssdc_jac'] = False
controller = controller_ParaDiag_nonMPI(
controller_params=controller_params, description=get_description(), num_procs=block_size
)
# ParaDiag diagonalizes in time, so the solution becomes complex
for S in controller.MS:
S.levels[0].prob.init = tuple([*S.levels[0].prob.init[:2]] + [np.dtype('complex128')])
P = controller.MS[0].levels[0].prob
dt = controller.MS[0].levels[0].params.dt
Tend = num_steps_total * dt
uend, stats = controller.run(u0=P.u_exact(0.0), t0=0.0, Tend=Tend)
niter = [int(me[1]) for me in get_sorted(stats, type='niter', sortby='time')]
return uend, niter, abs(uend - P.u_exact(Tend))
def main(cwd):
"""
A simple test program to test the MPI-parallel ParaDiag controller
Args:
cwd (str): current working directory
"""
# try to import MPI here, will fail if things go wrong (and not in the subprocess part)
try:
import mpi4py
del mpi4py
except ImportError as e:
raise ImportError('ParaDiag with MPI needs mpi4py') from e
import numpy as np
# Set python path once
my_env = os.environ.copy()
my_env['PYTHONPATH'] = '../../..:.'
my_env['COVERAGE_PROCESS_START'] = 'pyproject.toml'
# one time-step per rank, so the number of ranks is the block size
block_sizes = [1, 2, 4]
# set up new/empty file for output
fname = 'step_9_D_out.txt'
f = open(cwd + '/../../../data/' + fname, 'w')
f.close()
# run the MPI controller with different block sizes, always doing num_steps_total steps in total
for block_size in block_sizes:
print('Running ParaDiag with block size %2i...' % block_size)
cmd = ('mpirun -np ' + str(block_size) + ' python playground_ParaDiag_MPI.py ../../../../data/' + fname).split()
p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=my_env, cwd=cwd)
p.wait()
assert p.returncode == 0, 'ERROR: did not get return code 0, got %s with %2i processes' % (
p.returncode,
block_size,
)
# now do the same with the virtually parallel controller and append the results
f = open(cwd + '/../../../data/' + fname, 'a')
virtual = {}
for block_size in block_sizes:
uend, niter, error = run_virtual(block_size)
virtual[block_size] = uend
out = format_result('virtual', block_size, niter, error)
f.write(out + '\n')
print(out)
f.close()
# windowing must not change the answer: every block size integrates to the same end time
reference = virtual[block_sizes[0]]
for block_size in block_sizes[1:]:
assert np.allclose(
virtual[block_size], reference, atol=1e-9
), 'ERROR: virtual ParaDiag gives different results for different block sizes'
if __name__ == "__main__":
main('.')
The script running on each rank: pySDC/tutorial/step_9/playground_ParaDiag_MPI.py
import sys
from pathlib import Path
from mpi4py import MPI
import numpy as np
from pySDC.helpers.stats_helper import get_sorted
from pySDC.implementations.controller_classes.controller_ParaDiag_MPI import controller_ParaDiag_MPI
from pySDC.tutorial.step_9.D_paradiag_MPI import (
get_description,
get_controller_params,
format_result,
num_steps_total,
)
if __name__ == "__main__":
"""
A simple test program to do MPI-parallel ParaDiag runs
One time-step per rank, so the communicator spans the block that ParaDiag diagonalizes across. We
always integrate `num_steps_total` steps, so with fewer ranks the controller simply windows
through more blocks.
"""
# set MPI communicator
comm = MPI.COMM_WORLD
# one step per rank, so the number of ranks is the block size
block_size = comm.size
# instantiate the controller
controller = controller_ParaDiag_MPI(
controller_params=get_controller_params(), description=get_description(), comm=comm
)
# ParaDiag diagonalizes in time, so the solution becomes complex
P = controller.S.levels[0].prob
P.init = tuple([*P.init[:2]] + [np.dtype('complex128')])
# get initial values
t0 = 0.0
dt = controller.S.levels[0].params.dt
Tend = num_steps_total * dt
uinit = P.u_exact(t0)
# call main function to get things done...
uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend)
# gathering the iteration counts is collective, so every rank has to take part
niter = [int(me[1]) for me in get_sorted(stats, type='niter', sortby='time', comm=comm)]
# only the last rank has the end point of the block, so only it writes the output
if comm.rank == comm.size - 1:
fname = sys.argv[1] if len(sys.argv) == 2 else 'step_9_D_out.txt'
Path("data").mkdir(parents=True, exist_ok=True)
f = open('data/' + fname, 'a')
out = format_result('MPI', block_size, niter, abs(uend - P.u_exact(Tend)))
f.write(out + '\n')
print(out)
f.close()
Results:
MPI: block size 1, 4 block(s) of 1 step(s), iterations [3, 3, 3, 3], error 3.3560e-05
MPI: block size 2, 2 block(s) of 2 step(s), iterations [3, 3, 3, 3], error 3.3560e-05
MPI: block size 4, 1 block(s) of 4 step(s), iterations [3, 3, 3, 3], error 3.3560e-05
virtual: block size 1, 4 block(s) of 1 step(s), iterations [3, 3, 3, 3], error 3.3560e-05
virtual: block size 2, 2 block(s) of 2 step(s), iterations [3, 3, 3, 3], error 3.3560e-05
virtual: block size 4, 1 block(s) of 4 step(s), iterations [3, 3, 3, 3], error 3.3560e-05
Part E: Adaptive alpha¶
Every part so far picked \(\alpha\) by hand and kept it fixed, which means committing to one compromise for the whole run. A small \(\alpha\) approximates the original problem better and converges in fewer iterations, but conditions the diagonalization worse, so round-off and inexact inner solves get amplified. The right balance shifts as the residual falls, so a fixed value is wrong at one end of the run or the other.
The AdaptiveAlpha convergence controller updates \(\alpha\) after every iteration instead, following Čaklović et al.:
\[\gamma = L (3 \epsilon + \tau), \quad
\alpha_{k} = \sqrt{\frac{\gamma r_k}{e_k}}, \quad
e_{k+1} = 2 \sqrt{\gamma e_k r_k},\]
with \(L\) the number of steps in the block, \(\epsilon\) machine precision, \(\tau\) the inner solver tolerance, \(r_k\) the residual and \(e_k\) a running bound on the error.
Important things to note:
\(\gamma\) is an accuracy floor. There is no point pushing \(\alpha\) below the level at which round-off and the inner solver dominate anyway, which is why
inner_tolenters: a looser inner solve should get a larger \(\alpha\).The interesting result is not that the adaptive strategy wins on iteration count. It ties with the best fixed value we tried, but it gets there without being told, and it keeps \(\alpha\) orders of magnitude larger while doing so, which is exactly the margin that protects you once the inner solves are inexact.
The residual is taken over the whole block, so every rank computes the same \(\alpha\) and the controllers stay in step.
\(\alpha\) is a property of the method, not of the parallelization. The adaptive controller therefore has to produce the same iteration counts whether ParaDiag runs virtually or across MPI ranks, which is what we check here.
Full code: pySDC/tutorial/step_9/E_adaptive_alpha.py
"""
This script shows how to let ParaDiag choose its alpha by itself.
ParaDiag replaces the time-stepping matrix by an alpha-circulant approximation. Alpha trades two error
sources against each other: a small value approximates the original problem better and converges in
fewer iterations, but conditions the diagonalization worse, so round-off and inexact inner solves get
amplified. A single fixed alpha therefore has to be a compromise for the whole run, even though the
right balance shifts as the residual falls.
The `AdaptiveAlpha` convergence controller updates alpha after every iteration instead, following
`Caklovic et al. <https://doi.org/10.2140/camcos.2023.18.55>`_:
gamma = L * (3 * eps + tau)
alpha_k = sqrt(gamma * r_k / e_k)
e_{k+1} = 2 * sqrt(gamma * e_k * r_k)
with L the block size, eps machine precision, tau the inner solver tolerance, r_k the residual and
e_k a running bound on the error. Gamma is an accuracy floor: there is no point pushing alpha below
the level at which round-off and the inner solver dominate anyway.
We compare a few fixed alphas against the adaptive one on the same advection problem used in part D.
The interesting result is not that adaptive wins on iteration count -- it ties with the best fixed
value -- but that it gets there without being told, and while keeping alpha orders of magnitude
larger, which is exactly the margin that protects you once the inner solves are inexact.
Since alpha is a property of the method and not of the parallelization, the adaptive controller has to
give the same answer whether ParaDiag runs virtually or across MPI ranks. We check that too.
"""
import os
import subprocess
from pySDC.tutorial.step_9.D_paradiag_MPI import get_description, num_steps_total
# the fixed values we compare against, plus the adaptive strategy
alpha_settings = [1e-2, 1e-4, 1e-8, 'adaptive']
def get_controller_params(alpha):
"""
Controller parameters for one alpha setting.
Args:
alpha: a number, or the string 'adaptive'
Returns:
tuple: the controller parameters and the extra description entries
"""
from pySDC.implementations.convergence_controller_classes.adaptive_alpha import AdaptiveAlpha
controller_params = {}
controller_params['logger_level'] = 30
controller_params['average_jacobian'] = False
extra_description = {}
if alpha == 'adaptive':
# the adaptive controller overwrites this from the first iteration onwards, but ParaDiag needs
# some alpha to build its first transform with
controller_params['alpha'] = 1e-4
extra_description['convergence_controllers'] = {AdaptiveAlpha: {}}
else:
controller_params['alpha'] = alpha
return controller_params, extra_description
def format_result(mode, alpha, niter, error, final_alpha):
"""
One line of output, in the same shape for both controllers so they can be compared.
Args:
mode (str): 'MPI' or 'virtual'
alpha: the alpha setting used
niter (int): number of iterations needed
error (float): error against the exact solution
final_alpha (float): the alpha in use when the run finished
Returns:
str: the formatted line
"""
return (
f'{mode:>7s}: alpha {str(alpha):>9s} -> {niter:2d} iterations, '
f'error {error:.4e}, final alpha {final_alpha:.3e}'
)
def run(alpha, block_size, comm=None):
"""
Run the advection problem from part D with one alpha setting.
Args:
alpha: a number, or the string 'adaptive'
block_size (int): number of time-steps in one block
comm: MPI communicator, or None for the virtually parallel controller
Returns:
tuple: the end value, the iteration count, the error and the final alpha
"""
import numpy as np
from pySDC.helpers.stats_helper import get_sorted
controller_params, extra_description = get_controller_params(alpha)
description = {**get_description(), **extra_description}
if comm is None:
from pySDC.implementations.controller_classes.controller_ParaDiag_nonMPI import controller_ParaDiag_nonMPI
controller_params['mssdc_jac'] = False
controller = controller_ParaDiag_nonMPI(
controller_params=controller_params, description=description, num_procs=block_size
)
steps = controller.MS
else:
from pySDC.implementations.controller_classes.controller_ParaDiag_MPI import controller_ParaDiag_MPI
controller = controller_ParaDiag_MPI(controller_params=controller_params, description=description, comm=comm)
steps = [controller.S]
# ParaDiag diagonalizes in time, so the solution becomes complex
for S in steps:
S.levels[0].prob.init = tuple([*S.levels[0].prob.init[:2]] + [np.dtype('complex128')])
P = steps[0].levels[0].prob
dt = steps[0].levels[0].params.dt
Tend = num_steps_total * dt
uend, stats = controller.run(u0=P.u_exact(0.0), t0=0.0, Tend=Tend)
niter = max(int(me[1]) for me in get_sorted(stats, type='niter', sortby='time', comm=comm))
return uend, niter, abs(uend - P.u_exact(Tend)), controller.params.alpha
def main(cwd):
"""
Compare fixed and adaptive alpha, with both controllers.
Args:
cwd (str): current working directory
"""
try:
import mpi4py
del mpi4py
except ImportError as e:
raise ImportError('ParaDiag with MPI needs mpi4py') from e
import numpy as np
my_env = os.environ.copy()
my_env['PYTHONPATH'] = '../../..:.'
my_env['COVERAGE_PROCESS_START'] = 'pyproject.toml'
block_size = num_steps_total
fname = 'step_9_E_out.txt'
f = open(cwd + '/../../../data/' + fname, 'w')
f.close()
# the MPI controller, one rank per time-step, all alpha settings in one run
print('Running ParaDiag with %2i ranks...' % block_size)
cmd = ('mpirun -np ' + str(block_size) + ' python playground_adaptive_alpha.py ../../../../data/' + fname).split()
p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=my_env, cwd=cwd)
p.wait()
assert p.returncode == 0, 'ERROR: did not get return code 0, got %s' % p.returncode
# and the same with the virtually parallel controller
f = open(cwd + '/../../../data/' + fname, 'a')
results = {}
for alpha in alpha_settings:
uend, niter, error, final_alpha = run(alpha, block_size)
results[alpha] = (uend, niter)
out = format_result('virtual', alpha, niter, error, final_alpha)
f.write(out + '\n')
print(out)
f.close()
# the adaptive strategy should need no more iterations than the best fixed alpha we tried
best_fixed = min(results[a][1] for a in alpha_settings if a != 'adaptive')
assert (
results['adaptive'][1] <= best_fixed
), 'ERROR: adaptive alpha needed %s iterations, the best fixed alpha only %s' % (results['adaptive'][1], best_fixed)
# alpha changes the iteration, not the problem, so all settings solve the same thing
reference = results[alpha_settings[0]][0]
for alpha in alpha_settings[1:]:
assert np.allclose(results[alpha][0], reference, atol=1e-5), (
'ERROR: alpha %s gives a different solution' % alpha
)
if __name__ == "__main__":
main('.')
The script running on each rank: pySDC/tutorial/step_9/playground_adaptive_alpha.py
import sys
from pathlib import Path
from mpi4py import MPI
from pySDC.tutorial.step_9.E_adaptive_alpha import alpha_settings, format_result, run
if __name__ == "__main__":
"""
Compare fixed and adaptive alpha with the MPI-parallel ParaDiag controller.
One time-step per rank, so the communicator spans the block ParaDiag diagonalizes across. Alpha is
a property of the method rather than of the parallelization, so these numbers have to match the
ones the virtually parallel controller produces.
"""
comm = MPI.COMM_WORLD
lines = []
for alpha in alpha_settings:
uend, niter, error, final_alpha = run(alpha, comm.size, comm=comm)
lines.append(format_result('MPI', alpha, niter, error, final_alpha))
# only the last rank has the end point of the block, so only it writes the output
if comm.rank == comm.size - 1:
fname = sys.argv[1] if len(sys.argv) == 2 else 'step_9_E_out.txt'
Path("data").mkdir(parents=True, exist_ok=True)
with open('data/' + fname, 'a') as f:
for line in lines:
f.write(line + '\n')
print(line)
Results:
MPI: alpha 0.01 -> 5 iterations, error 3.3560e-05, final alpha 1.000e-02
MPI: alpha 0.0001 -> 3 iterations, error 3.3560e-05, final alpha 1.000e-04
MPI: alpha 1e-08 -> 2 iterations, error 3.3560e-05, final alpha 1.000e-08
MPI: alpha adaptive -> 2 iterations, error 3.3560e-05, final alpha 2.976e-06
virtual: alpha 0.01 -> 5 iterations, error 3.3560e-05, final alpha 1.000e-02
virtual: alpha 0.0001 -> 3 iterations, error 3.3560e-05, final alpha 1.000e-04
virtual: alpha 1e-08 -> 2 iterations, error 3.3560e-05, final alpha 1.000e-08
virtual: alpha adaptive -> 2 iterations, error 3.3560e-05, final alpha 2.976e-06