Particles: migration and deposition

Download notebookView on GitHubRun when the site was built.

Particle-in-cell codes keep particles with the rank that owns the grid cell they are in, and deposit their charge on the grid. This tutorial moves particles in a periodic 2D box, sends each one to its new owner with Layout.owners and mpa.migrate, deposits the charge with cloud-in-cell weights into the blocks and their halo cells, and folds the halo contributions into the neighbours with accumulate_halos. The particle arrays are the backend’s (NumPy or CuPy, through cunumpy), so the same code runs on GPUs. Run it serially or with mpiexec -n 4 python particles.py.

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

import mpiarray as mpa

n = 64  # cells per axis of the unit box
dx = 1.0 / n
rho = mpa.zeros((n, n), split=(0, 1), halo=1, periodic=True)
layout = rho.layout
comm, rank = layout.comm, layout.rank
print(f"rank {rank} owns cells {layout.index_bounds}")

rank 0 owns cells ((0, 64), (0, 64))

Particles in each block

Every rank starts with particles inside its own block of cells. The positions are in box units, $[0, 1)^2$.

rng = np.random.default_rng(rank)
(i0, i1), (j0, j1) = layout.index_bounds
count = 2000
positions = xp.asarray(  # drawn on the host, then moved to the backend
    np.column_stack(
        [rng.uniform(i0 * dx, i1 * dx, count), rng.uniform(j0 * dx, j1 * dx, count)]
    )
)
velocities = xp.asarray(rng.normal(0.0, 0.05, (count, 2)))
total = comm.allreduce(len(positions))
if rank == 0:
    print("particles:", total)

particles: 2000

Moving and migrating

After a push, some particles have left the block. owners maps their cells to ranks for all particles at once; mpa.migrate then sends every row of the position and velocity arrays to its rank, as raw buffers in one Alltoallv per array.

for _ in range(20):
    positions = (positions + 0.02 * velocities) % 1.0
    destination = layout.owners(xp.floor(positions / dx).astype(int) % n)
    positions, velocities = mpa.migrate(destination, positions, velocities, comm=comm)

cells = xp.floor(positions / dx).astype(int)
everyone_home = comm.allreduce(
    bool((layout.owners(cells) == rank).all()), op=mpa.layout.MPI.LAND
)
total_now = comm.allreduce(len(positions))
if rank == 0:
    print("all particles with their owner:", everyone_home, "| particles:", total_now)

all particles with their owner: True | particles: 2000

Depositing the charge

Cloud-in-cell spreads each particle over the four cells around it, with weights that add up to one. A particle near the edge of a block also writes into the halo cells. bincount adds the weights per cell (on NumPy or CuPy, with repeated cells counted correctly). accumulate_halos then adds the halo contributions to the neighbouring blocks (wrapping around the periodic box) and clears the halos, so no charge is lost.

def cloud_in_cell(grid, positions, start, wrap=False):
    """Add one unit of charge per particle to ``grid``, a storage array."""
    shape = grid.shape
    p = positions / dx - 0.5  # in cell units, cell centres at integers
    corner = xp.floor(p).astype(int)
    weight = p - corner
    index = corner - xp.asarray(start)  # into the storage
    for di in (0, 1):
        for dj in (0, 1):
            w = (weight[:, 0] if di else 1 - weight[:, 0]) * (
                weight[:, 1] if dj else 1 - weight[:, 1]
            )
            i, j = index[:, 0] + di, index[:, 1] + dj
            if wrap:
                i, j = i % shape[0], j % shape[1]
            cell = i * shape[1] + j
            grid += xp.bincount(cell, weights=w, minlength=grid.size).reshape(shape)


def deposit(rho, positions):
    rho.local_with_halos[...] = 0
    cloud_in_cell(rho.local_with_halos, positions, (i0 - 1, j0 - 1))  # halo layer
    rho.accumulate_halos()


deposit(rho, positions)
deposited = rho.sum()
if rank == 0:
    print(f"deposited charge: {deposited:.6f} (particles: {total_now})")

deposited charge: 2000.000000 (particles: 2000)

Every particle lies within one cell of its own block, so with one halo layer all its four cells are in the storage. After accumulate_halos the deposited charge equals the number of particles.

The same with a replicated grid

Codes that split only the particles keep the whole grid on every rank (split=None): each rank deposits its own particles on its copy, and allreduce_replicated adds the copies up. The result is the same density.

full = mpa.zeros((n, n), split=None)
cloud_in_cell(full.local, positions, (0, 0), wrap=True)
full.allreduce_replicated()
density = rho.to_numpy(root=0)  # NumPy on rank 0, for plotting
if rank == 0:
    print("same density:", np.allclose(density, xp.to_numpy(full.local)))
    fig, ax = plt.subplots(figsize=(4, 3.5))
    image = ax.imshow(density.T, origin="lower", extent=(0, 1, 0, 1), cmap="viridis")
    fig.colorbar(image, ax=ax, label="particles per cell")
    ax.set(xlabel="x", ylabel="y", title="deposited charge")

same density: True