utilities.parallel_utils#

Common general utilities for MPI-parallel simulations.

Functions#

`gather_vlen_array`(v, mpi_comm[, mpi_root])	Gather 1D arrays of possibly different lengths onto root process Other processes return None
`gather_vlen_arrays`(arrays, mpi_comm[, mpi_root])	Gather several 1D NumPy arrays of the same local length onto root process.
`parallel_run_from_cli`(parser, main)	Parse CLI arguments and run main, in parallel.

Details#

Common general utilities for MPI-parallel simulations.

gather_vlen_array(v, mpi_comm, mpi_root=0)[source]#

Gather 1D arrays of possibly different lengths onto root process Other processes return None

Parameters:

varray_like: Local 1D array. All ranks must use the same dtype.
mpi_commmpi4py.MPI.Comm: MPI communicator.
mpi_rootint, default=0: Rank of root process in MPI communicator.

Returns:

numpy.ndarray or None: Gathered 1D array on root. None on the other processes.

gather_vlen_arrays(arrays, mpi_comm, mpi_root=0)[source]#

Gather several 1D NumPy arrays of the same local length onto root process. Return None on the other processes. Length may vary between processes.

Parameters:

arraysiterable of 1D arrays: All local arrays must have the same length and dtype All ranks must pass the same number of arrays
mpi_commmpi4py.MPI.Comm: MPI communicator.
mpi_rootint, default=0: Rank of root process in MPI communicator.

Returns:

tuple[int, …] or None: On root, tuple of gathered 1D arrays. None on the other processes.

parallel_run_from_cli(parser, main)[source]#

Parse CLI arguments and run main, in parallel.

This function makes sure that the standard output remains clean when the size of the MPI communicator is > 1, in the case of argparse’s help or error messages. In addition, it prints the MPI communicator size and the CLI arguments. When not run (serially) from IPython, it also calls plt.show() to keep any Matplotlib windows open after main is run.

Parameters:

parsercallable: Function which parses the CLI arguments using the argparse library. It has no input arguments, and it returns an argparse Namespace. This function is only run by the root process, which has MPI rank 0. The output Namespace is then broadcasted to all the other processes.
maincallable: Main function to be run in parallel. Its arguments are all keyword- only, and most of them are extracted from the parser output using vars(). An additional important argument is mpi_comm, which is an instance of mpi4py.MPI.Comm. This function returns a dict with all its local variables, e.g. by using return locals(), which is useful for interactive usage. If matplotlib figures are generated, these are made visible with fig.show(), but without calling plt.show().

Returns:

dict: Dictionary with all the local variables. Useful for interactive usage.