Metadata-Version: 2.4
Name: mpiarray
Version: 0.2.1
Summary: MPI domain decomposition and distributed NumPy/CuPy arrays with halo exchange.
Author: Max
License: MIT License
        
        Copyright (c) 2026 Max Lindqvist
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Project-URL: Source, https://github.com/max-models/mpiarray
Keywords: array,cupy,domain decomposition,halo exchange,mpi,numpy
Classifier: Development Status :: 3 - Alpha
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE.txt
Requires-Dist: cunumpy>=0.6.3
Requires-Dist: maybempi>=0.1.2
Requires-Dist: numpy
Provides-Extra: dev
Requires-Dist: jupyterlab; extra == "dev"
Requires-Dist: pre-commit; extra == "dev"
Requires-Dist: pyproject-fmt; extra == "dev"
Requires-Dist: pyright; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Requires-Dist: ty; extra == "dev"
Requires-Dist: mpiarray[docs,mpi,test]; extra == "dev"
Provides-Extra: docs
Requires-Dist: astropy; extra == "docs"
Requires-Dist: h5py; extra == "docs"
Requires-Dist: griffe; extra == "docs"
Requires-Dist: ipykernel; extra == "docs"
Requires-Dist: matplotlib; extra == "docs"
Requires-Dist: nbclient>=0.10; extra == "docs"
Requires-Dist: nbformat>=5.10; extra == "docs"
Requires-Dist: plasma-plots>=0.2.1; extra == "docs"
Requires-Dist: xarray; extra == "docs"
Provides-Extra: hdf5
Requires-Dist: h5py; extra == "hdf5"
Provides-Extra: mpi
Requires-Dist: mpi4py; extra == "mpi"
Provides-Extra: units
Requires-Dist: astropy; extra == "units"
Provides-Extra: xarray
Requires-Dist: xarray; extra == "xarray"
Provides-Extra: test
Requires-Dist: astropy; extra == "test"
Requires-Dist: h5py; extra == "test"
Requires-Dist: matplotlib; extra == "test"
Requires-Dist: nbformat>=5.10; extra == "test"
Requires-Dist: pytest; extra == "test"
Requires-Dist: pytest-cov; extra == "test"
Requires-Dist: xarray; extra == "test"
Dynamic: license-file

# mpiarray


<!-- README.md is generated from README.qmd: edit the .qmd and run `make readme`. -->

[![Tests](https://github.com/max-models/mpiarray/actions/workflows/test_pytest.yml/badge.svg)](https://github.com/max-models/mpiarray/actions/workflows/test_pytest.yml)
[![Static
analysis](https://github.com/max-models/mpiarray/actions/workflows/static_analysis.yml/badge.svg)](https://github.com/max-models/mpiarray/actions/workflows/static_analysis.yml)
[![Docs](https://github.com/max-models/mpiarray/actions/workflows/docs.yml/badge.svg)](https://max-models.github.io/mpiarray/)
[![codecov](https://codecov.io/gh/max-models/mpiarray/branch/main/graph/badge.svg)](https://codecov.io/gh/max-models/mpiarray)
[![PyPI](https://img.shields.io/pypi/v/mpiarray.png)](https://pypi.org/project/mpiarray/)
[![Python](https://img.shields.io/pypi/pyversions/mpiarray.png)](https://pypi.org/project/mpiarray/)

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

``` python
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
```

``` bash
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](https://github.com/max-models/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

``` bash
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](https://mpi4py.readthedocs.io/), 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](https://docs.astral.sh/uv/):

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

or with pip, in a Python 3.10+ environment:

``` bash
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](https://docs.astral.sh/ruff/), type
checking [pyright](https://microsoft.github.io/pyright/) and
[ty](https://docs.astral.sh/ty/), run by
[pre-commit](https://pre-commit.com/) and in CI:

``` bash
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:

``` bash
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](https://www.conventionalcommits.org/); see
[CONTRIBUTING.md](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](https://ewels.github.io/starlight-pydocs/). It needs
Node 22 or newer.

``` bash
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](https://quarto.org/):

``` bash
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](https://max-models.github.io/mpiarray/development/publishing/).
