Building, reducing and saving arrays

Download notebookView on GitHubRun when the site was built.

Simulation data usually starts on the ranks that computed it. This tutorial builds distributed arrays from those pieces with from_local, reduces them along axes, reads parts of them, and writes them to .npy files that NumPy reads, every rank handling only its own block. Run it serially or with mpiexec -n 4 python working_with_data.py.

import os
import tempfile

import cunumpy as xp
import matplotlib.pyplot as plt
import numpy as np

import mpiarray as mpa

comm = mpa.default_comm()  # MPI.COMM_WORLD, or a serial stand-in without mpiexec
rank, size = comm.Get_rank(), comm.Get_size()

From the pieces the ranks hold

Say every rank recorded a different number of samples, three values each. from_local stacks the pieces in rank order along the first axis and, since their lengths differ from the near-even split, redistributes them once. Nothing global is built.

rng = np.random.default_rng(rank)
samples = rng.normal(loc=rank, size=(10 + 5 * rank, 3))  # 10, 15, 20, ... rows
records = mpa.from_local(samples)
if rank == 0:
    print("global shape:", records.shape)
print(
    f"rank {rank}: had {len(samples)} rows, now holds rows {records.layout.index_bounds[0]}"
)

global shape: (10, 3)
rank 0: had 10 rows, now holds rows (0, 10)

Reductions along an axis

Reducing along an axis does not gather the array: every rank reduces its own block, and only these partial results travel, to the ranks that own them in the result, which is again a distributed array. Here a 2D field and its profiles along each axis:

n = 96


def wave(i, j):
    return xp.sin(2 * xp.pi * i / n) * xp.cos(4 * xp.pi * j / n) + i / n


field = mpa.fromfunction(wave, (n, n), split=(0, 1))
column_mean = field.mean(axis=1)  # a distributed array: one value per row
row_max = field.max(axis=0)
peak = field.max()  # over everything: a number on every rank
column_mean, row_max = column_mean.to_numpy(), row_max.to_numpy()  # small: to plot
if rank == 0:
    print(
        "profiles:",
        column_mean.shape,
        row_max.shape,
        "| global maximum:",
        round(float(peak), 4),
    )
    fig, ax = plt.subplots(figsize=(5, 3))
    ax.plot(column_mean, label="mean over j")
    ax.plot(row_max, label="max over i")
    ax.set(xlabel="index", title="profiles of the field")
    ax.legend()

profiles: (96,) (96,) | global maximum: 1.7623

Reading parts

Indexing with global indices is collective. A single element comes back as a number on every rank; a slice is a new distributed array, made by sending only the selected cells to the ranks that own them in the result. to_numpy() turns it into a NumPy array:

corner = field[:3, :3]  # a new distributed array, made of only these cells
value = field[10, 20]  # one element, broadcast from its owner
small = corner.to_numpy()  # collective: outside the `if`
if rank == 0:
    print(small.round(3))
    print(
        "field[10, 20] =",
        round(float(value), 4),
        "| exact:",
        round(float(wave(10, 20)), 4),
    )

[[0.    0.    0.   ]
 [0.076 0.075 0.074]
 [0.151 0.15  0.147]]
field[10, 20] = -0.423 | exact: -0.423

Saving and loading

mpa.save writes an ordinary .npy file with MPI-IO: rank 0 writes the header, every rank its own block. numpy.load reads it, and mpa.load reads it back with any layout, again each rank reading only its block. The file must be on a file system all ranks see; rank 0 chooses the name here.

directory = comm.bcast(tempfile.mkdtemp() if rank == 0 else None, root=0)
path = os.path.join(directory, "field.npy")
mpa.save(path, field)

again = mpa.load(path, split=1, halo=1)  # a different layout than when saved
same = np.array_equal(again.to_numpy(), field.to_numpy())
if rank == 0:
    print("numpy reads it:", np.load(path).shape, "| loaded == saved:", same)
    print(again.layout)

numpy reads it: (96, 96) | loaded == saved: True
Layout(shape=(96, 96), split=(1,), halo=(1, 1), periodic=(False, False), process_grid=(1, 1), rank=0 of 1)

When something hangs

Every call above that communicates is collective: all ranks must make it, in the same order. A collective inside if rank == 0: makes the program hang. With MPIARRAY_DEBUG=1 in the environment, mpiarray checks every collective call: ranks in different calls raise an error naming both, and a rank left waiting raises after MPIARRAY_DEBUG_TIMEOUT seconds (default 30) instead of hanging.

MPIARRAY_DEBUG=1 mpiexec -n 4 python working_with_data.py