NuMPI

NuMPI is a collection of numerical tools for MPI-parallelized Python codes. NuMPI presently contains:

• An (incomplete) stub implementation of the mpi4py interface to the MPI libraries. This allows running serial versions of MPI parallel code without having mpi4py (and hence a full MPI stack) installed.
• Parallel file IO in numpy's .npy format using MPI I/O.
• MPI-parallel L-BFGS optimizers:
  • l_bfgs — unconstrained, with a strong-Wolfe line search.
  • l_bfgs_bounded — box-constrained (lo <= x <= hi) with optional index pinning, two-loop recursion and projected Armijo backtracking.
  • l_bfgs_projected — a single linear equality <a, x> = target plus optional box bounds.
• An MPI-parallel bound constrained conjugate gradients algorithm.

Build status

Installation

python3 -m pip install NuMPI

Development Installation

Clone the repository.

To use the code, install the current package as editable:

pip install -e .[test]

Testing

You have to do a development installation to be able to run the tests.

From the main installation directory, run the test suite with pytest:

python -m pytest

This runs serially and exercises the MPI stub implementation, so it works without mpi4py installed. To run a single test file or test:

python -m pytest test/Optimization/test_cg.py
python -m pytest test/Optimization/test_cg.py::test_directions -s

To run the suite under real MPI at a given number of ranks (this is what CI does, for 1, 2, 4, 8 and 10 ranks):

mpirun --oversubscribe -n 4 python -m pytest --verbose

MPI Conventions

All of NuMPI's parallel algorithms operate on distributed arrays: each MPI rank holds a slice of the global data, and scalar quantities (energies, norms, convergence tolerances, Lagrange multipliers) are globally reduced — the same value on every rank. Understanding the split between local and global is essential to using the optimizers correctly; this section spells it out.

Distributed vs. global

User-supplied callbacks

The solvers call back into user code in a few places; each has a specific contract.

• Objective fun(x) -> (energy, gradient) (when jac=True) or separate fun(x) -> energy and jac(x) -> gradient:

  • energy must be a globally reduced scalar. All ranks must return the same number. The standard way to do this is to compute a local quantity and reduce it with pnp.sum(...).item() (or equivalent), where pnp is the Reduction(comm) wrapper. Returning a local energy is the single most common MPI mistake: ranks will silently disagree in line-search acceptance tests and the optimisation will diverge or hang.
  • gradient is local — only the current rank's slice.

• callback(x) receives the current local iterate. If the caller needs the global state (for plotting or logging from rank 0), they must gather explicitly.

• hessp(x, d) (CG) returns a local Hessian-vector product.

Building distributed inputs

Use NuMPI.Tools.Reduction(comm) to obtain a pnp object whose sum, max, min, mean, dot methods perform MPI_Allreduce across the communicator. When mpi4py is not installed, NuMPI.MPIStub provides the same interface with a single "rank", so the same code runs serially too.

A typical setup with a communicator-provided subdomain looks like:

from NuMPI.Tools import Reduction
from NuMPI.Optimization import LinearConstraint, l_bfgs_projected

pnp = Reduction(comm)

# a_local: this rank's slice of the global weight vector, shape matching x
# target: global scalar, same on every rank
lc = LinearConstraint(a_local, target, pnp=pnp)

def fun(x):                      # x is the local slice
    # compute local integrand, then REDUCE for the scalar return
    local_energy = 0.5 * np.sum((x - y_local) ** 2)
    return pnp.sum(local_energy).item(), (x - y_local)   # gradient stays local

res = l_bfgs_projected(fun, x0_local, lc, jac=True,
                       bounds_lo=0.0, bounds_hi=1.0,
                       comm=comm, gtol=1e-5)

The returned res.x is the local slice of the solution; res.fun, res.multiplier, and res.max_grad are globally reduced scalars.

See NuMPI/Optimization/__init__.py for optimizer-specific notes and test/Optimization/MPIMinimizationProblems.py::MPI_Quadratic for a reference implementation of a distributed objective.

Development & Funding

Development of this project is funded by the European Research Council within Starting Grant 757343 and by the Deutsche Forschungsgemeinschaft within project EXC 2193.