mpiarray

Distributed NumPy/CuPy arrays over MPI, with halo exchange, that you create like NumPy arrays:

import numpy as np

import mpiarray as mpa

a = mpa.array([1, 2, 3, 4])  # mpiexec -n 2: rank 0 holds [1, 2], rank 1 holds [3, 4]
b = mpa.zeros((64, 48), split=(0, 1), halo=1, periodic=(True, False))
c = 2 * a + np.sin(a)  # elementwise: no communication

total = c.sum()  # collective: the same value on every rank
full = c.gather()  # collective: the whole array on every rank
print(c)  # this rank's block; printing never communicates

mpiexec -n 2 python example.py   # or just: python example.py

• Creation like NumPy: array, zeros, ones, full, empty, the *_like versions, arange, linspace and fromfunction, which compute only the local block. Arrays are split along their first axis by default; split= chooses other axes, split=None gives every rank the whole array.
• Halo cells per axis: update_halos() before a stencil, accumulate_halos() after depositing particles near block edges; both can overlap with computation (wait=False).
• Operators, ufuncs and reductions as in NumPy (sum, max, mean, var, argmax, cumsum, norm, vdot, …), and NumPy’s functions (np.mean(a, axis=0), np.allclose, np.where, …) through __array_function__. Global reductions return the same host scalar on every rank; reductions along axes and slices (a[2:10, :]) return distributed arrays, without gathering anything.
• Redistribution and load balancing: redistribute moves an array to any other layout with one Alltoallv; explicit or weighted cut points (Layout(bounds=...), Layout.weighted, rebalance) and layout.aligned for cell and node arrays.
• Particles: layout.owners finds the rank of many cells at once, mpa.migrate sends rows of NumPy or CuPy arrays to their ranks.
• NumPy or CuPy: arrays come from cunumpy, so the same code runs on the GPU. Device buffers are copied through host memory unless the program declares its MPI CUDA-aware (xp.mpi.set_mpi_cuda_aware(True) or xp.mpi.mpi_is_cuda_aware(comm)).
• With or without MPI: without an MPI launcher the script runs as one rank on maybempi’s stand-in for mpi4py.MPI, without importing mpi4py.
• Data in and out: from_local builds an array from the pieces the ranks hold; save/load write and read ordinary .npy files in parallel with MPI-IO, and save_hdf5/load_hdf5 several arrays and attributes in one HDF5 file (mpiarray[hdf5]).
• Halo boundary conditions at walls: constant, "edge", "symmetric", "reflect".
• Debugging: with MPIARRAY_DEBUG=1, a collective called on only some ranks raises an error instead of hanging.

Each array has a layout (a mpa.Layout) with its process grid, neighbours and owned index ranges; most code only reads it.

Documentation: https://max-models.github.io/mpiarray/

Install

pip install "mpiarray[mpi]"   # with mpi4py, for runs under mpiexec
pip install mpiarray          # serial only, no MPI library needed
pip install "mpiarray[hdf5]"  # with h5py, for save_hdf5/load_hdf5

The mpi extra installs mpi4py, which needs an MPI library, e.g. brew install open-mpi or sudo apt-get install libopenmpi-dev openmpi-bin. Without it, mpiarray runs on maybempi’s serial stand-in for mpi4py.MPI, as one rank holding the whole array. Starting such an installation with mpiexec gives a warning, and every process then computes the whole problem on its own.

For development, with uv:

make install    # uv sync --extra dev, plus the pre-commit hooks

or with pip, in a Python 3.10+ environment:

pip install -e ".[dev]"

The test, docs and dev extras install the test runner, the documentation tooling and the linters; dev includes mpi.

Development

Formatting and linting use ruff, type checking pyright and ty, run by pre-commit and in CI:

make lint     # ruff check, ruff format --check, pyright, ty
make test     # pytest with coverage

The tests run serially and under MPI; some only run on 2 or 6 ranks:

mpiexec -n 2 .venv/bin/python -m pytest
mpiexec -n 6 .venv/bin/python -m pytest
make coverage   # serial and 2, 3, 4, 6 ranks, combined; fails below 100% line and branch coverage

The CuPy code paths run on the CPU with cunumpy’s fake CuPy (CUNUMPY_FAKE_CUPY=1 CUNUMPY_BACKEND=cupy); .gitlab-ci.yml runs the tests on NVIDIA GPUs.

Commit messages follow Conventional Commits; see CONTRIBUTING.md.

Build docs

The documentation in docs/ is an Astro + Starlight site: hand-written pages, the notebooks in guides/ executed and published as pages, and the API reference generated from the docstrings with starlight-pydocs. It needs Node 22 or newer.

make docs-install     # npm packages and the Python docs extra
make docs-notebooks   # execute guides/*.ipynb and convert them to pages
make docs-dev         # live preview at http://localhost:4321/mpiarray/
make docs-build       # the static site in docs/dist

Build the README

README.md is rendered from README.qmd with Quarto:

make readme

Releases

Before merging a release to main, update the version in pyproject.toml, src/mpiarray/__init__.py and CITATION.cff (including its release date), and add the release notes to CHANGELOG.md. The push to main creates a GitHub release with a vX.Y.Z tag and publishes the package to PyPI with trusted publishing (OIDC). The one-time PyPI and GitHub configuration is described in the publishing guide.