Part B: mpi4py-fft for parallel Fourier transforms#
The most prominent parallel solver is, probably, the FFT. While many implementations or wrappers for Python exist, we decided to use mpi4py-fft, which provided the easiest installation, a simple API and good parallel scaling. As an example, we solve the nonlinear Schrödinger equation in 2D, with the IMEX sweeper treating the nonlinear parts explicitly. The code works both in real and in spectral space, the latter usually being faster. This example tests SDC, MLSDC and PFASST.
One run#
The problem class nonlinearschroedinger_imex gets the MPI communicator, over which mpi4py-fft distributes the
grid, and fft_to_fft transfers between the levels. Each run appends its error, statistics of the iterations and
the time to solution to data/step_7_B_out.txt, from rank 0.
import numpy as np
from pathlib import Path
from mpi4py import MPI
from pySDC.helpers.stats_helper import get_sorted
from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order
from pySDC.implementations.problem_classes.NonlinearSchroedinger_MPIFFT import nonlinearschroedinger_imex
from pySDC.implementations.transfer_classes.TransferMesh_MPIFFT import fft_to_fft
def run_simulation(spectral=None, ml=None, num_procs=None):
"""
A test program to do SDC, MLSDC and PFASST runs for the 2D NLS equation
Args:
spectral (bool): run in real or spectral space
ml (bool): single or multiple levels
num_procs (int): number of parallel processors
"""
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
# initialize level parameters
level_params = dict()
level_params['restol'] = 1e-08
level_params['dt'] = 1e-01 / 2
level_params['nsweeps'] = [1]
# initialize sweeper parameters
sweeper_params = dict()
sweeper_params['quad_type'] = 'RADAU-RIGHT'
sweeper_params['num_nodes'] = [3]
sweeper_params['QI'] = ['LU'] # For the IMEX sweeper, the LU-trick can be activated for the implicit part
sweeper_params['initial_guess'] = 'zero'
# initialize problem parameters
problem_params = dict()
if ml:
problem_params['nvars'] = [(128, 128), (32, 32)]
else:
problem_params['nvars'] = [(128, 128)]
problem_params['spectral'] = spectral
problem_params['comm'] = comm
# initialize step parameters
step_params = dict()
step_params['maxiter'] = 50
# initialize controller parameters
controller_params = dict()
controller_params['logger_level'] = 30 if rank == 0 else 99
# controller_params['predict_type'] = 'fine_only'
# fill description dictionary for easy step instantiation
description = dict()
description['problem_params'] = problem_params # pass problem parameters
description['problem_class'] = nonlinearschroedinger_imex
description['sweeper_class'] = imex_1st_order
description['sweeper_params'] = sweeper_params # pass sweeper parameters
description['level_params'] = level_params # pass level parameters
description['step_params'] = step_params # pass step parameters
description['space_transfer_class'] = fft_to_fft
# set time parameters
t0 = 0.0
Tend = 1.0
f = None
if rank == 0:
Path("data").mkdir(parents=True, exist_ok=True)
f = open('data/step_7_B_out.txt', 'a')
out = f'Running with ml={ml} and num_procs={num_procs}...'
f.write(out + '\n')
print(out)
# instantiate controller
controller = controller_nonMPI(num_procs=num_procs, controller_params=controller_params, description=description)
# get initial values on finest level
P = controller.MS[0].levels[0].prob
uinit = P.u_exact(t0)
# call main function to get things done...
uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend)
uex = P.u_exact(Tend)
err = abs(uex - uend)
if rank == 0:
# filter statistics by type (number of iterations)
iter_counts = get_sorted(stats, type='niter', sortby='time')
niters = np.array([item[1] for item in iter_counts])
out = (
f' Min/Mean/Max number of iterations: '
f'{np.min(niters):4.2f} / {np.mean(niters):4.2f} / {np.max(niters):4.2f}'
)
f.write(out + '\n')
print(out)
out = ' Range of values for number of iterations: %2i ' % np.ptp(niters)
f.write(out + '\n')
print(out)
out = ' Position of max/min number of iterations: %2i -- %2i' % (
int(np.argmax(niters)),
int(np.argmin(niters)),
)
f.write(out + '\n')
print(out)
out = ' Std and var for number of iterations: %4.2f -- %4.2f' % (float(np.std(niters)), float(np.var(niters)))
f.write(out + '\n')
print(out)
out = f'Error: {err:6.4e}'
f.write(out + '\n')
print(out)
timing = get_sorted(stats, type='timing_run', sortby='time')
out = f'Time to solution: {timing[0][1]:6.4f} sec.'
f.write(out + '\n')
print(out)
assert err <= 1.133e-05, 'Error is too high, got %s' % err
if ml:
if num_procs > 1:
maxmean = 12.5
else:
maxmean = 6.6
else:
maxmean = 12.7
assert np.mean(niters) <= maxmean, 'Mean number of iterations is too high, got %s' % np.mean(niters)
f.write('\n')
print()
f.close()
All runs#
SDC, MLSDC and PFASST with 10 steps in parallel (emulated in one process by controller_nonMPI), each in real
and in spectral space. The code runs in serial with
python B_pySDC_with_mpi4pyfft.py, and in parallel in space with, e.g., mpirun -np 2 python B_pySDC_with_mpi4pyfft.py.
def main():
"""
Little helper routine to run the whole thing
Note: This can also be run with "mpirun -np 2 python B_pySDC_with_mpi4pyfft.py"
"""
run_simulation(spectral=False, ml=False, num_procs=1)
run_simulation(spectral=True, ml=False, num_procs=1)
run_simulation(spectral=False, ml=True, num_procs=1)
run_simulation(spectral=True, ml=True, num_procs=1)
run_simulation(spectral=False, ml=True, num_procs=10)
run_simulation(spectral=True, ml=True, num_procs=10)
if __name__ == "__main__":
main()
Results#
mpi4py-fft does not run in the browser, nor in the environment this website is built in. These are the results of our CI, which runs this part in an environment with mpi4py-fft, in the run that built this page:
Running with ml=False and num_procs=1...
Min/Mean/Max number of iterations: 9.00 / 12.70 / 16.00
Range of values for number of iterations: 7
Position of max/min number of iterations: 0 -- 19
Std and var for number of iterations: 2.10 -- 4.41
Error: 1.1321e-05
Time to solution: 1.4272 sec.
Running with ml=False and num_procs=1...
Min/Mean/Max number of iterations: 8.00 / 11.40 / 15.00
Range of values for number of iterations: 7
Position of max/min number of iterations: 0 -- 19
Std and var for number of iterations: 2.03 -- 4.14
Error: 4.1749e-06
Time to solution: 1.1746 sec.
Running with ml=True and num_procs=1...
Min/Mean/Max number of iterations: 5.00 / 6.60 / 8.00
Range of values for number of iterations: 3
Position of max/min number of iterations: 0 -- 16
Std and var for number of iterations: 1.07 -- 1.14
Error: 1.1316e-05
Time to solution: 1.4305 sec.
Running with ml=True and num_procs=1...
Min/Mean/Max number of iterations: 4.00 / 5.95 / 8.00
Range of values for number of iterations: 4
Position of max/min number of iterations: 0 -- 19
Std and var for number of iterations: 1.02 -- 1.05
Error: 4.1744e-06
Time to solution: 1.3196 sec.
Running with ml=True and num_procs=10...
Min/Mean/Max number of iterations: 7.00 / 12.45 / 18.00
Range of values for number of iterations: 11
Position of max/min number of iterations: 9 -- 10
Std and var for number of iterations: 3.11 -- 9.65
Error: 1.1306e-05
Time to solution: 3.0203 sec.
Running with ml=True and num_procs=10...
Min/Mean/Max number of iterations: 6.00 / 11.50 / 17.00
Range of values for number of iterations: 11
Position of max/min number of iterations: 9 -- 10
Std and var for number of iterations: 3.04 -- 9.25
Error: 4.1688e-06
Time to solution: 2.7796 sec.
Important things to note
The nonlinear Schrödinger example is not expected to work well with PFASST. In fact, SDC and MLSDC converge for larger time steps, but PFASST does not.