Particle binning#
Keyed sums: the mass in each cell
A simulation tracks many particles; you want the mass in each cell of a mesh, and how many particles fell there. That is a keyed sum: every particle adds its weight at the key of its cell. Many particles share a key, and hawk combines them so there is no race.
Time: ~5 min · Runs on: CPU (same code on GPU) · You need: hawk, eagle, NumPy and matplotlib; read first: Across lanes
Colab or Kaggle: run this first cell (it installs the packages and downloads the helper files this notebook uses). Elsewhere it only checks the imports.
Particles#
Positions in [0, 1) drawn from a skewed distribution, each with a weight.
from nb_helpers import batch_size
import numpy as np
NB = 16 # cells
n = batch_size(gpu=20_000_000, cpu=20_000)
rng = np.random.default_rng(4)
x = rng.beta(2, 5, size=n)
w = rng.uniform(0.5, 1.5, size=n)
print(f"{n} particles, {NB} cells")
20000000 particles, 16 cells
The particle’s words#
Say the vocabulary first. A particle has a position x and a weight w; once its cell is known, the cell number is a key, and when the particles are kept sorted by cell it is a Sorted[Index]. Declare those once as a KernelKind. The kernels below are written in these names, and a kernel takes only the ones it needs.
from hawk import Index, Scalar, Sorted
from hawk.ext import DATA_ONLY, KernelKind
class Particle(KernelKind, slug="particle"):
x: Scalar # position in [0, 1)
w: Scalar # weight
cell: Sorted[Index] # the cell number, for particles kept sorted by cell
guard = DATA_ONLY
Particle
__main__.Particle
The kernel#
m.index(x * NB, NB) turns a position into a cell number, kept inside 0 .. NB-1. hawk.scatter(mass)[cell] += w
adds the weight there. A sink declared with a name and an extent, Mutable["bin":NB], is an output you do not have to
supply: it starts at zero.
policy="compensated" asks for the tightest sum: the error in a cell is at most k * eps * sum|w|, with k = 1 on the card. The default, "auto", is the fast
form. See Across lanes for the three policies.
import eagle
import hawk
import hawk.math as m
from hawk import Index, Mutable, Scalar, Sorted
@Particle
def deposit(x, w, mass: Mutable["bin":NB], count: Mutable["bin":NB]):
cell = m.index(x * NB, NB)
hawk.scatter(mass, policy="compensated")[cell] += w
hawk.scatter(count)[cell] += 1.0
out = eagle.deploy(deposit).run(x=x, w=w, returns="numpy")
cell = np.minimum((x * NB).astype(int), NB - 1)
assert np.array_equal(out["count"], np.bincount(cell, minlength=NB))
print("count matches NumPy; total count", int(out["count"].sum()))
count matches NumPy; total count 20000000
Check against the exact sum#
math.fsum gives each cell’s correctly rounded sum. The compensated mass must stay within eps * sum|w| of it
(the weights are positive, so sum|w| is the cell’s mass), and nothing may be lost: the cells add up to the total
weight and the counts to n.
import math
exact = np.array([math.fsum(w[cell == k]) for k in range(NB)])
bound = np.finfo(float).eps * np.bincount(cell, np.abs(w), NB)
assert np.all(np.abs(out["mass"] - exact) <= bound)
assert math.isclose(out["mass"].sum(), math.fsum(w), rel_tol=1e-12) and out["count"].sum() == n
used = bound > 0 # an empty cell has nothing to bound
print(f"worst cell error {np.max(np.abs(out['mass'] - exact)[used] / bound[used]):.2f} of the bound; mass and count conserved")
worst cell error 0.00 of the bound; mass and count conserved
When the keys are already sorted#
If the particles are sorted by cell, equal keys sit side by side. Particle-in-cell codes keep their particles sorted by cell, which is when this applies. Declare the key as Sorted[Index] and hawk adds each
run of equal keys once, instead of once per particle. The promise is yours: values that decrease are not detected unless
you set HAWK_CHECK_SORTED.
@Particle
def deposit_sorted(cell, w, mass: Mutable["bin":NB]):
hawk.scatter(mass)[cell] += w
order = np.argsort(cell, kind="stable")
mass_sorted = eagle.deploy(deposit_sorted).run(cell=cell[order].astype(np.int32), w=w[order], returns="numpy")
assert np.all(np.abs(mass_sorted - exact) <= bound) # the default sum meets the same bound
print("sorted keys give the same mass")
sorted keys give the same mass
The histogram#
The positions follow a Beta(2, 5) law, whose cumulative form is 1 - (1-x)^5 (1+5x); the mean weight is 1. The
outlines show the expected counts and mass density.
import matplotlib.pyplot as plt
edges = np.arange(NB + 1) / NB
cdf = 1 - (1 - edges) ** 5 * (1 + 5 * edges)
expect = n * np.diff(cdf) # expected particles per cell
fig, (a0, a1) = plt.subplots(1, 2, figsize=(7.2, 3.0), constrained_layout=True)
a0.stairs(out["count"], edges, fill=True, alpha=0.6, label="deposited")
a0.stairs(expect, edges, color="#898781", label="expected")
a0.set_xlabel("position x / L")
a0.set_ylabel("particles")
a0.legend()
a1.stairs(out["mass"] * NB, edges, fill=True, alpha=0.6)
a1.stairs(expect * NB, edges, color="#898781")
a1.set_xlabel("position x / L")
a1.set_ylabel("mass per unit length")
show(fig, "Particle count and mass per unit length against the expected Beta distribution")
What just happened#
hawk.scatter(sink)[key] += vis a keyed sum; many lanes may share a key.policy="compensated"picked the tightest sum for the mass; the count used the default.Sorted[Index]told hawk the keys never decrease, so it added each run of equal keys once.The mass stays within
eps * sum|w|of the exactly rounded sum in every cell, and the totals are conserved.
Across lanes has the other policies and the sums over axes of a grid.