Monte Carlo#
Trajectories that stop early
A batch of projectiles flies until each one lands or leaves the field. They stop at different steps, so the live set thins out. eagle keeps a map of the live samples, recomputes it every few steps, and runs the whole loop as one replayed graph.
Time: ~5 min · Runs on: CPU (a GPU is picked up automatically and runs a larger batch) · You need: eagle, hawk, NumPy and matplotlib; read first: Finished samples: compaction, reorder, and when not to
A Monte Carlo ensemble draws many launch conditions and follows each trajectory until it ends:
a landing, an escape, a collision. Here every projectile has its own speed and angle, feels
linear drag, and stops when it lands or crosses the far edge of the field. You will compose the
loop from ActiveSet, compaction_body and repeat_while, and then look at what group= changes. The projectile model and its
checks against the exact solution are on hawk’s Monte Carlo page; this page is the orchestration around the kernel.
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.
import sys, pathlib
sys.path.insert(0, str(pathlib.Path.cwd().parent.parent / "_shared"))
from nb_helpers import gpu_available, plot_style, to_numpy, xp_for
import numpy as np
DEVICE = "gpu" if gpu_available() else "cpu"
xp = xp_for(DEVICE)
n = 8192 if DEVICE == "gpu" else 512
print("running on", DEVICE, "with", n, "projectiles")
running on cpu with 512 projectiles
The projectile’s words#
Say the vocabulary first. A projectile has a position and a velocity, written by every step; the drag, gravity, time step and field length are the same for all; and a terminated flag marks the samples that have landed or left. Declare them once as a KernelKind, with the guard that reads the map of live samples. The step below is written in these names.
from hawk import Mutable, Param, Scalar, Terminated
from hawk.ext import Guard, KernelKind
class Projectiles(KernelKind, slug="projectile_ensemble"):
drag: Param
g: Param
dt: Param
xmax: Param
terminated: Terminated
x: Mutable[Scalar]
y: Mutable[Scalar]
vx: Mutable[Scalar]
vy: Mutable[Scalar]
guard = Guard(active_set=True) # a launch covers only the live samples
Projectiles
__main__.Projectiles
The kernel: one step, and the condition that ends it#
Units are SI: metres, metres per second, seconds. A sample’s state is a position and a velocity.
The step is one semi-implicit (symplectic) Euler update with linear drag k = 0.1 1/s and
dt = 0.02 s, so k * dt = 0.002 and the step is stable. It is first order in dt; the check
below confirms that. terminated is the one statement that says when a sample is done: it is set
once the projectile is below the ground or past xmax, so a landing is seen one step late (the
overshoot is at most dt * |vy|, about 0.6 m here). From then on the sample’s state no longer
changes.
The kernel is built with Guard(active_set=True), as in Finished samples: compaction, reorder, and when not to: it
reads a map of live samples, so a launch can cover only those.
import hawk
from hawk import Mutable, Param, Scalar, Terminated
from hawk.math import lor
def fly_step(drag, g, dt, xmax, terminated, x, y, vx, vy):
vx = vx - dt * drag * vx
vy = vy - dt * (g + drag * vy)
x = x + dt * vx
y = y + dt * vy
terminated = lor(y < 0.0, x > xmax)
fly = hawk.kernel(fly_step, kind=Projectiles.kind, steps=1) # `steps=` goes with the kernel, not with the vocabulary
import eagle
plan = eagle.deploy(fly)
G, K, XMAX = 9.81, 0.1, 40.0
def fresh(speed, angle):
m = len(speed)
state = dict(x=xp.zeros(m), y=xp.full(m, 1e-3),
vx=xp.asarray(speed * np.cos(angle)), vy=xp.asarray(speed * np.sin(angle)))
done = xp.zeros(m, dtype=xp.bool_)
return state, done, eagle.ActiveSet(done)
def bind(state, done, active, dt, xmax=XMAX):
return plan.bind(drag=K, g=G, dt=dt, xmax=xmax, terminated=done, **state,
finished_count=xp.zeros(1, xp.uint32), **active.planes())
Draw the ensemble#
Speeds and angles are drawn at random, so landing times spread out: a flat shot lands within a few steps, a lob takes many.
rng = np.random.default_rng(3)
speed = rng.uniform(8.0, 30.0, n)
angle = rng.uniform(0.15, 1.35, n)
DT = 0.02
The loop, composed#
compaction_body runs the step a few times (every), then recomputes the live-sample map on the
device, only when the finished counter moved. repeat_while turns that into one bounded loop that
stops when every sample is done. On a GPU the loop is one CUDA graph, built once and replayed, so
nothing is read back to the host between steps. On the CPU it runs as a plain loop.
from eagle import GraphPipeline, SkipGuard, compaction_body, repeat_while
def run_loop(speed, angle, every=4, max_steps=2000):
state, done, active = fresh(speed, angle)
counter = xp.zeros(1, xp.uint32)
run = plan.bind(drag=K, g=G, dt=DT, xmax=XMAX, terminated=done, **state,
finished_count=counter, **active.planes())
body = compaction_body(run.launch, active, every=every, finished=counter)
loop = repeat_while(body, SkipGuard(counter, 0, xp.asarray([len(speed)], xp.uint32), 0),
-(-max_steps // every))
if DEVICE == "gpu":
GraphPipeline().add(loop).build().launch()
xp.cuda.Device().synchronize()
else:
loop()
return state, done, loop
state, done, loop = run_loop(speed, angle)
landed = int(to_numpy(done).sum())
print(f"{loop.iterations()} loop passes, {landed} of {n} samples ended")
65 loop passes, 512 of 512 samples ended
A host loop, for the picture#
The graph replays the loop without showing you the steps in between. To draw the trajectories we step the same plan by hand and keep every position: after each step the map is recomputed on the device and read back. That read-back is the cost the graph above avoids, so use it for looking, not for running. The finished state is the same.
SHOWN = 48 # trajectories kept for the picture
def traced(speed, angle):
state, done, active = fresh(speed, angle)
run = bind(state, done, active, DT)
active.compact()
path = [tuple(to_numpy(state[k])[:SHOWN].copy() for k in "xy") + (to_numpy(done)[:SHOWN].copy(),)]
live, covered = [len(speed)], [len(speed)]
ended = np.full(len(speed), -1)
for step in range(1, 2000):
run.launch()
active.compact()
d = to_numpy(done)
ended[(d) & (ended < 0)] = step
live.append(len(speed) - int(d.sum()))
covered.append(int(to_numpy(active.count)[0]))
path.append(tuple(to_numpy(state[k])[:SHOWN].copy() for k in "xy") + (d[:SHOWN].copy(),))
if live[-1] == 0:
break
return state, path, live, covered, ended
state_h, path, live, covered, ended = traced(speed, angle)
same = all(np.array_equal(to_numpy(state[k]), to_numpy(state_h[k])) for k in "xy")
steps = len(live) - 1
print(f"every projectile ended within {steps} steps; same final state as the graph loop: {same}")
every projectile ended within 260 steps; same final state as the graph loop: True
What a launch covers, and what group= is for#
The kernel reads one lane at a time, so group=1 is the right set for it: a launch covers exactly
the live samples. group=32 is for kernels that declare a group, such as a sum over the fastest
axis of a lane space, a commit that every lane of a warp makes together, or neighbour reads. It
keeps a whole group of 32 lanes while any member is live, and pays back in contiguous,
coalesced access, not in fewer lanes. eagle refuses a group=32 set for this flat kernel.
What the group costs in lanes depends on lane order. Counts below for group=32 are what a set
would cover, worked out from the step each sample ended at; the group=1 line is the real map.
def lanes_per_step(ended, group):
# lanes a set of this group would cover at each step: live groups times the group size
pad = -len(ended) % group
last = np.concatenate([ended, np.full(pad, -1)]).reshape(-1, group)
last = np.where(last < 0, len(live), last).max(axis=1) # a group lives until its last member ends
return np.array([group * int((last > s).sum()) for s in range(len(live))])
lanes_32 = lanes_per_step(ended, 32)
print("lanes over the whole run: group=1", sum(covered), "| group=32 (random order, would cover)",
int(lanes_32.sum()), "| no compaction", n * steps)
lanes over the whole run: group=1 55235 | group=32 (random order, would cover) 117408 | no compaction 133120
Draw order is arbitrary, so sort it. The no-drag flight time, min(2 vy0 / g, xmax / vx0), is
a cheap guess at how long a sample lives. Samples that end together then share groups.
guess = np.minimum(2 * speed * np.sin(angle) / G, XMAX / (speed * np.cos(angle)))
order = np.argsort(guess)
_, _, _, _, ended_sorted = traced(speed[order], angle[order])
lanes_32_sorted = lanes_per_step(ended_sorted, 32)
print("group=32 after sorting by the flight-time guess (would cover):", int(lanes_32_sorted.sum()))
group=32 after sorting by the flight-time guess (would cover): 65216
import matplotlib.pyplot as plt
from themed_figure import show
plot_style()
fig, ax = plt.subplots(figsize=(5.5, 3.2))
ax.plot(live, label="still flying")
ax.plot(covered, "--", label="covered, group=1")
ax.plot(lanes_32, label="would cover, group=32")
ax.plot(lanes_32_sorted, label="would cover, group=32, sorted")
ax.set_xlabel("step")
ax.set_ylabel("lanes")
ax.set_title(f"{n:,} projectiles")
ax.legend()
fig.tight_layout()
show()
Without compaction every step covers all n lanes until the last projectile lands. Whether
gathering the live samples pays on your device depends on the batch size and the cost of the
step: Finished samples: compaction, reorder, and when not to has the card for each device.
The animation#
Trajectories on the left, the lanes a launch covers on the right. The same frames are written to
a GIF for the page (and to a short MP4 when imageio-ffmpeg is installed).
from anim import frame_of, write_anim
from themed_figure import DARK, LIGHT
pick = np.unique(np.linspace(0, len(path) - 1, min(len(path), 72)).astype(int))
def render(theme, k):
pal = LIGHT if theme == "light" else DARK
fig, (a, b) = plt.subplots(1, 2, figsize=(6.4, 2.9), dpi=60, facecolor=pal["bg"])
for ax in (a, b):
ax.set_facecolor(pal["bg"])
for s in ax.spines.values():
s.set_color(pal["rule"])
ax.tick_params(colors=pal["sub"], labelsize=7)
i = pick[k]
for j in range(SHOWN):
xs = [p[0][j] for p in path[: i + 1]]
ys = [p[1][j] for p in path[: i + 1]]
a.plot(xs, ys, lw=0.6, color=pal["us"], alpha=0.5)
a.plot(xs[-1], ys[-1], "o", ms=2.5, color=pal["off"] if path[i][2][j] else pal["on"])
a.set_xlim(0, 42)
a.set_ylim(0, 30)
a.set_xlabel("x (m)", color=pal["sub"], fontsize=8)
a.set_ylabel("height (m)", color=pal["sub"], fontsize=8)
b.plot(covered[: i + 1], color=pal["us"], lw=1.2, label="group=1")
b.plot(lanes_32[: i + 1], "--", color=pal["on"], lw=1.2, label="group=32, would cover")
b.set_xlim(0, len(path))
b.set_ylim(0, n * 1.05)
b.set_xlabel("step", color=pal["sub"], fontsize=8)
b.set_ylabel("lanes covered", color=pal["sub"], fontsize=8)
b.legend(fontsize=7, labelcolor=pal["ink"], facecolor=pal["bg"], edgecolor=pal["rule"], loc="upper right")
fig.tight_layout()
return frame_of(fig, pal["bg"])
for theme in ("light", "dark"):
frames = [render(theme, k) for k in range(len(pick))]
for name, size in write_anim(frames, f"monte_carlo_early_termination_{theme}").items():
print(name, f"{size / 1e6:.2f} MB")
monte_carlo_early_termination_light.gif 0.75 MB
monte_carlo_early_termination_dark.gif 0.71 MB
Try this#
Narrow the draws to speeds in (18, 22) and angles in (0.6, 0.8) and re-run: the projectiles
end within about 50 steps of each other, and the gap between the group=1 and group=32 totals
shrinks.
Next#
Finished samples: compaction, reorder, and when not to: when gathering the live samples pays, device by device.
A multi-kernel step, declaratively: the same loop with several kernels sharing one set.
Scan and compaction:
ActiveSetandgroup=in the reference.