Moving markers between ranks

Download notebookView on GitHubRun when the site was built.

A ParticleStore holds the markers of one MPI rank. When markers move, the ones that leave a rank’s part of the domain are sent to their new owner with migrate. This tutorial loads markers into a box split between four ranks, pushes them around a periodic box and checks that no marker or id is lost.

The notebook kernel is a single process, so the MPI cells write a script and launch it with mpiexec -n 4; without mpiexec those cells say so and the rest still runs.

The store of one rank

Storage is component-major, as in a particle-in-cell code: positions is (ndim, capacity), velocities (nvel, capacity), and weights, alive and ids have one entry per column. The first length columns are in use; a dead marker keeps its column until compact.

store = ParticleStore(ndim=2, nvel=2, capacity=4)
rng = np.random.default_rng(0)
ids = store.add(rng.uniform(size=(2, 6)), rng.normal(size=(2, 6)))
print("ids", ids, "| length", store.length, "| capacity", store.capacity)

store.deactivate(xp.asarray([False, True, False, False, True, False]))
print("alive", store.num_alive, "of", store.length)
store.compact()
print("after compact: length", store.length, "ids", store.ids[: store.length])

ids [0 1 2 3 4 5] | length 6 | capacity 8
alive 4 of 6
after compact: length 4 ids [0 2 3 5]

Every marker has an int64 id that stays with it, also when it moves to another rank. Ids are handed out by an IdAllocator that gives rank r of n the ids r, r + n, r + 2n, ..., so they are unique without communication.

Splitting the box

A Decomposition tells which rank owns a position. BoxDecomposition is the reference one: equal blocks, ranks in row-major order. The script below runs on every rank. Each rank draws the same 2000 markers and keeps its own (keep_owned), then every step moves all markers by their velocity and calls migrate. A marker that crosses into a neighbouring block is sent there; the others never leave memory.

%%writefile moving_markers.py
import cunumpy as xp
import maybempi
import numpy as np

from mpiparticles import (
    BoxDecomposition,
    IdAllocator,
    ParticleStore,
    gather_columns,
    global_count,
    keep_owned,
    migrate,
)

comm = maybempi.get_mpi().COMM_WORLD  # the serial stand-in without mpiexec
rank, size = comm.Get_rank(), comm.Get_size()
proc_sizes = {1: (1, 1), 2: (2, 1), 4: (2, 2)}[size]
decomposition = BoxDecomposition(
    comm, [0.0, 0.0], [1.0, 1.0], proc_sizes, periodic=True
)

# every rank loads the same markers, then keeps the ones in its block
rng = np.random.default_rng(42)
count = 2000
store = ParticleStore(2, 2, id_allocator=IdAllocator.from_comm(comm))
store.add(
    rng.uniform(size=(2, count)),
    rng.normal(scale=0.05, size=(2, count)),
    ids=np.arange(count),
)
keep_owned(store, decomposition)
per_rank = comm.allgather(store.num_alive)

dt = 1.0
sent_total = 0
for step in range(20):
    end = store.length
    store.positions[:, :end] += dt * store.velocities[:, :end]
    sent, received = migrate(store, decomposition)  # collective
    sent_total += comm.allreduce(sent)
    store.compact_if_needed()

# nothing was lost or duplicated, and every marker sits where it belongs
ids = gather_columns(store.active_ids(), comm, root=None)
assert global_count(store, comm) == count
assert sorted(ids.tolist()) == list(range(count))
assert bool(xp.all(decomposition.owned_mask(store.active_positions())))

final = comm.allgather(store.num_alive)
if rank == 0:
    print(f"{size} rank(s), blocks {proc_sizes}")
    print("markers per rank at the start:", per_rank)
    print("markers per rank at the end:  ", final)
    print("markers sent between ranks in 20 steps:", sent_total)
    print("all", count, "ids present and unique")

Writing moving_markers.py

Run it on four ranks. The remaining ranks take part in the same collective calls, so the output is the same wherever it is printed.

run_mpi(4, sys.executable, "moving_markers.py")

mpiexec is not installed: skipping the MPI run

Only the leavers travel. migrate first tests every marker against the rank’s own block (two comparisons per axis), looks up the owner of the few that failed, and sends those as two arrays: one float matrix with a (position, velocity, weight) row per marker and the ids as a separate int64 array, so ids are never converted to floats and stay exact above $2^{53}$. The leavers are marked dead locally and the arrivals are appended with their ids.

The same script without MPI

Without a launcher, maybempi provides a one-rank stand-in for mpi4py, and migrate returns (0, 0) without any MPI call:

result = subprocess.run(
    [sys.executable, "moving_markers.py"], capture_output=True, text=True, cwd=WORKDIR
)
print(result.stdout + result.stderr[-1000:], end="")

1 rank(s), blocks (1, 1)
markers per rank at the start: [2000]
markers per rank at the end:   [2000]
markers sent between ranks in 20 steps: 0
all 2000 ids present and unique

Your own decomposition

migrate(store, decomposition) takes anything that follows the Decomposition protocol: a comm, rank and size, an owners(positions) that returns the owning rank of every column of a (ndim, n) position array, and a cheap owned_mask(positions) that says which columns this rank owns. A code with its own domain decomposition wraps it in a small class like that and keeps using the same store and migrate.