L05 · Galactic Collision

"""L05 Galactic Collision: direct-sum N-body in Numba, leapfrog vs Euler, and a Toomre-style merger.

Run from the repo root:   uv run python L05_galactic_collision/sim.py          (full run, ~5 min)
                          uv run python L05_galactic_collision/sim.py --quick  (smoke test, ~1 min)
"""
import os

os.environ.setdefault("NUMBA_NUM_THREADS", "4")  # shared machine: cap threads BEFORE numba loads
os.environ.setdefault("OMP_NUM_THREADS", "4")
import math
import sys
import time
import timeit
from multiprocessing import Pool
from pathlib import Path

import imageio_ffmpeg
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FFMpegWriter
from numba import njit, prange
from scipy.ndimage import gaussian_filter, gaussian_filter1d
from scipy.spatial import cKDTree
from scipy.spatial.transform import Rotation

QUICK = "--quick" in sys.argv
OUT = Path(__file__).parent / ("output/quick" if QUICK else "output")   # smoke tests never clobber the real media
plt.rcParams["animation.ffmpeg_path"] = imageio_ffmpeg.get_ffmpeg_exe()

# N-body units: G = 1, length = 3.5 kpc (disk scale length), mass = 3e10 Msun (one stellar disk).
L_KPC, M_SUN = 3.5, 3e10
V_KMS = math.sqrt(4.3009e-6 * M_SUN / L_KPC)  # G in kpc (km/s)^2 / Msun -> ~192 km/s
T_MYR = L_KPC * 3.0857e16 / V_KMS / 3.156e13  # kpc->km, s->Myr        -> ~17.8 Myr

# Collision set-up (sized so the full run fits in ~5 min on 4 threads)
N_DISK, N_HALO = (1500, 1000) if QUICK else (5000, 3000)  # per galaxy
EPS = 0.08  # Plummer softening (~280 pc)
DT = 0.05  # ~0.9 Myr. ponytail: one global dt, per-particle hierarchical steps (GADGET, SWIFT) for cusps
T_END = 40.0 if QUICK else 120.0  # ~2.1 Gyr
FRAME_EVERY = 4
RP, D0 = 5.0, 20.0  # pericentre of the point-mass parabola and starting separation (N-body lengths)


# ----------------------------------------------------------------------------- gravity kernels
def accel(px, py, pz, m, eps2, ax, ay, az):
    """Direct-sum acceleration on every particle (G = 1, Plummer softening), O(N^2).

    One source, five speeds: run as plain Python on lists, or compiled by the njit variants below.
    Structure-of-arrays (separate x, y, z) keeps the inner j-loop on contiguous memory -> SIMD.
    The i == j term is harmless: dx = 0 and eps2 > 0.
    """
    n = len(px)  # ponytail: O(N^2) pairs, Barnes-Hut tree / FMM / TreePM once N >> 1e5
    for i in prange(n):
        xi, yi, zi = px[i], py[i], pz[i]
        sx = sy = sz = 0.0
        for j in range(n):
            dx = px[j] - xi
            dy = py[j] - yi
            dz = pz[j] - zi
            r2 = dx * dx + dy * dy + dz * dz + eps2
            w = m[j] / (r2 * math.sqrt(r2))
            sx += dx * w
            sy += dy * w
            sz += dz * w
        ax[i] = sx
        ay[i] = sy
        az[i] = sz


accel_jit = njit(accel)  # compiled, strict IEEE maths, 1 thread
# + fastmath (sums may be reordered) and the numpy error model (no ZeroDivisionError branch in the loop):
# both are needed before LLVM will vectorise the j-loop with AVX (4 doubles / 8 floats per instruction)
accel_simd = njit(fastmath=True, error_model="numpy")(accel)
accel_par = njit(parallel=True, fastmath=True, error_model="numpy")(accel)  # + i-loop over NUMBA_NUM_THREADS


@njit(parallel=True, fastmath=True)
def accel_aos(x, m, eps2, a):
    """Same maths on an (N, 3) array-of-structs: strided loads, so the compiler struggles to vectorise."""
    for i in prange(x.shape[0]):
        sx = sy = sz = 0.0
        for j in range(x.shape[0]):
            dx = x[j, 0] - x[i, 0]
            dy = x[j, 1] - x[i, 1]
            dz = x[j, 2] - x[i, 2]
            r2 = dx * dx + dy * dy + dz * dz + eps2
            w = m[j] / (r2 * math.sqrt(r2))
            sx += dx * w
            sy += dy * w
            sz += dz * w
        a[i, 0] = sx
        a[i, 1] = sy
        a[i, 2] = sz


def accel_numpy(x, m, eps2):
    """NumPy broadcasting: builds (3, N, N) temporaries, so memory caps N at a few thousand."""
    d = x[:, None, :] - x[:, :, None]  # d[k, i, j] = x_j - x_i
    r2 = (d * d).sum(0) + eps2
    return (d * (m / (r2 * np.sqrt(r2)))).sum(2)


@njit(parallel=True, fastmath=True)
def potential(px, py, pz, m, eps2):
    """Total softened potential energy: sum over pairs of -m_i m_j / sqrt(r^2 + eps^2)."""
    n = len(px)
    w = 0.0
    for i in prange(n):
        s = 0.0
        for j in range(i + 1, n):
            dx = px[j] - px[i]
            dy = py[j] - py[i]
            dz = pz[j] - pz[i]
            s += m[j] / math.sqrt(dx * dx + dy * dy + dz * dz + eps2)
        w -= m[i] * s
    return w


def gravity(m, eps, kernel=accel_par, dtype=np.float64):
    """force(x) -> accelerations, both (3, N). dtype=float32 is the GPU-style mixed-precision mode:
    pair maths in float32, sums in float64, state (x, v) stays float64."""
    mk, e2 = m.astype(dtype), dtype(eps * eps)

    def force(x):
        xk = x.astype(dtype)
        a = np.empty_like(xk)
        kernel(xk[0], xk[1], xk[2], mk, e2, a[0], a[1], a[2])
        return a

    return force


def energy(x, v, m, eps):
    return 0.5 * np.sum(m * (v * v).sum(0)) + potential(x[0], x[1], x[2], m, eps * eps)


# ----------------------------------------------------------------------------- integrators
def leapfrog(x, v, force, dt, nsteps):
    """Kick-drift-kick leapfrog (symplectic, time-reversible, 2nd order). Updates x, v in place."""
    a = force(x)
    for n in range(1, nsteps + 1):
        v += 0.5 * dt * a
        x += dt * v
        a = force(x)
        v += 0.5 * dt * a
        yield n, x, v


def euler(x, v, force, dt, nsteps):
    """Explicit (forward) Euler: 1st order and NOT symplectic -> energy drifts secularly."""
    for n in range(1, nsteps + 1):
        a = force(x)
        x += dt * v
        v += dt * a
        yield n, x, v


# ----------------------------------------------------------------------------- verification
def verify():
    """Two-body Kepler orbit: period vs 2 pi sqrt(a^3 / G M), and leapfrog vs Euler energy error."""
    m, a_, e, eps = np.array([1.0, 0.5]), 1.0, 0.5, 1e-6  # eps << r: effectively Newtonian
    M = m.sum()
    T = 2 * np.pi * np.sqrt(a_**3 / M)
    r0, v0 = a_ * (1 + e), np.sqrt(M * (1 - e) / (a_ * (1 + e)))  # start at apocentre
    x0 = np.zeros((3, 2))
    u0 = np.zeros((3, 2))
    x0[0] = [-m[1] / M * r0, m[0] / M * r0]  # centre-of-mass frame
    u0[1] = [-m[1] / M * v0, m[0] / M * v0]
    orbits, per = 20, 500
    dt, n = T / per, orbits * per
    force = gravity(m, eps, accel_jit)
    E0 = energy(x0, u0, m, eps)
    res = {}
    for name, stepper in (("leapfrog", leapfrog), ("Euler", euler)):
        rel, dE = [], []
        for _, x, v in stepper(x0.copy(), u0.copy(), force, dt, n):
            rel.append(x[:2, 1] - x[:2, 0])
            dE.append(abs(energy(x, v, m, eps) / E0 - 1))
        res[name] = np.array(rel), np.array(dE)
    rel = res["leapfrog"][0]
    theta = np.unwrap(np.arctan2(rel[:, 1], rel[:, 0]))
    t = dt * np.arange(1, n + 1)
    k = int(theta[-1] // (2 * np.pi))
    T_meas = np.interp(2 * np.pi * k, theta, t) / k  # time of the k-th full revolution / k
    err_T = abs(T_meas / T - 1)
    lf, eu = res["leapfrog"][1], res["Euler"][1]
    print(f"[verify] Kepler period: analytic {T:.6f}, leapfrog {T_meas:.6f} ({k} orbits), error {err_T:.2e}")
    print(f"[verify] |dE/E| after {orbits} orbits: leapfrog max {lf.max():.2e} (last orbit {lf[-per:].max():.2e}),"
          f" Euler final {eu[-1]:.2e}")
    assert err_T < 0.01, "Kepler period off by more than 1%"
    assert lf[-per:].max() < 2 * lf[:per].max(), "leapfrog energy error is growing (should stay bounded)"
    assert eu[-1] > 100 * lf.max(), "Euler should drift far beyond leapfrog"

    fig, (a1, a2) = plt.subplots(1, 2, figsize=(12, 5), gridspec_kw={"width_ratios": [1, 1.4]})
    for (name, (r, _)), c in zip(res.items(), ("#4cc9f0", "#f72585")):
        a1.plot(r[:, 0], r[:, 1], lw=0.6, color=c, label=name)
    a1.plot(0, 0, "*", color="w", ms=10)
    a1.set(aspect="equal", xlabel="x (relative, N-body length)", ylabel="y (relative)",
           title=f"Relative orbit, e = {e}, {orbits} orbits, {per} steps/orbit")
    a1.legend(loc="upper left")
    a2.semilogy(t / T, lf, color="#4cc9f0", lw=0.8, label="leapfrog (symplectic): bounded")
    a2.semilogy(t / T, eu, color="#f72585", lw=1.2, label="explicit Euler: secular drift")
    a2.set(xlabel="time (orbital periods)", ylabel="|E(t)/E0 - 1|",
           title=f"Energy error. Period check: {T_meas:.5f} vs {T:.5f} ({err_T:.1e})")
    a2.legend(loc="lower right")
    a2.grid(alpha=0.2)
    fig.tight_layout()
    fig.savefig(OUT / "verify_kepler_energy.png", dpi=130)
    plt.close(fig)


# ----------------------------------------------------------------------------- benchmark
def benchmark():
    """Interactions per second for each rung of the performance ladder, N = 64 ... 65536."""
    rng = np.random.default_rng(1)
    e2 = EPS * EPS
    tiers = [  # name, largest N worth timing, make(x, m) -> zero-arg callable
        ("pure Python (lists)", 512,
         lambda x, m: lambda: accel(*x.tolist(), m.tolist(), e2, *([0.0] * len(m) for _ in range(3)))),
        ("NumPy broadcasting", 2048, lambda x, m: lambda: accel_numpy(x, m, e2)),
        ("Numba @njit", 8192, lambda x, m: lambda f=gravity(m, EPS, accel_jit): f(x)),
        ("+ fastmath (SIMD)", 16384, lambda x, m: lambda f=gravity(m, EPS, accel_simd): f(x)),
        ("+ parallel, 4 threads", 32768, lambda x, m: lambda f=gravity(m, EPS, accel_par): f(x)),
        ("+ float32 (production)", 65536, lambda x, m: lambda f=gravity(m, EPS, accel_par, np.float32): f(x)),
        ("parallel, AoS (N,3) layout", 32768,
         lambda x, m: lambda xa=np.ascontiguousarray(x.T): accel_aos(xa, m, e2, np.empty_like(xa))),
    ]
    Ns = [2**k for k in range(6, 17)]
    rates = {}
    for name, nmax, make in tiers:
        rates[name] = []
        for N in [N for N in Ns if N <= nmax]:
            x, m = rng.normal(size=(3, N)), np.full(N, 1.0 / N)
            f = make(x, m)
            t0 = time.perf_counter()
            f()  # warm-up (includes JIT compile on the very first call)
            if N == Ns[0] and "Python" not in name and "NumPy" not in name:
                print(f"[bench] {name}: first call incl. JIT compile {time.perf_counter() - t0:.2f} s")
            n = max(1, int(0.05 / timeit.timeit(f, number=1)))  # batch tiny calls to >= 50 ms per sample
            t = min(timeit.repeat(f, number=n, repeat=3)) / n  # best of 3: least disturbed by other jobs
            rates[name].append((N, N * N / t))
        N, r = rates[name][-1]
        print(f"[bench] {name:28s} N={N:6d}  {r:10.3e} interactions/s  ({N * N / r * 1e3:9.2f} ms per force eval)")

    base = rates["pure Python (lists)"][-1][1]
    print("\n| rung | largest N timed | interactions/s | speed-up vs pure Python |\n|---|---|---|---|")
    for name, rr in rates.items():
        print(f"| {name} | {rr[-1][0]} | {rr[-1][1]:.2e} | {rr[-1][1] / base:,.0f}x |")
    print()

    fig, (a1, a2) = plt.subplots(1, 2, figsize=(13, 5.2))
    colors = ["#adb5bd", "#f4a261", "#2a9d8f", "#4cc9f0", "#7b2cbf", "#f72585", "#7b2cbf"]
    for (name, rr), c in zip(rates.items(), colors):
        N, r = np.array(rr).T
        ls = "--" if "AoS" in name else "-"
        a1.loglog(N, r, ls, marker="o", ms=4, color=c, label=name)
        a2.loglog(N, N * N / r, ls, marker="o", ms=4, color=c, label=name)
    Nr = np.array([64, 65536])
    a2.loglog(Nr, 5e-9 * Nr**2, ":", color="w", alpha=0.6, label=r"$\propto N^2$")
    a1.set(xlabel="N (particles)", ylabel="pair interactions per second",
           title="Throughput: each rung of the ladder (4 threads max)")
    a2.set(xlabel="N (particles)", ylabel="wall time per force evaluation (s)",
           title="Cost of one force evaluation: direct sum is O(N$^2$)")
    for a in (a1, a2):
        a.grid(alpha=0.25)
        a.grid(alpha=0.07, which="minor")
    a1.legend(fontsize=8, loc="lower right")
    fig.tight_layout()
    fig.savefig(OUT / "benchmark.png", dpi=130)
    plt.close(fig)
    return rates


# ----------------------------------------------------------------------------- initial conditions
def unit_vectors(n, rng):
    u = rng.normal(size=(3, n))
    return u / np.linalg.norm(u, axis=0)


def plummer(n, M, b, rng):
    """Isotropic Plummer sphere, Aarseth, Henon & Wielen (1974) recipe."""
    r = b / np.sqrt(rng.uniform(0, 0.99, n) ** (-2 / 3) - 1)  # ponytail: outer 1% of mass cut off
    q = np.empty(0)
    while q.size < n:  # rejection-sample q = v / v_esc from g(q) ~ q^2 (1 - q^2)^3.5
        u, w = rng.random(n), 0.1 * rng.random(n)
        q = np.r_[q, u[w < u**2 * (1 - u**2) ** 3.5]]
    v = q[:n] * np.sqrt(2 * M / np.sqrt(r * r + b * b))
    return r * unit_vectors(n, rng), v * unit_vectors(n, rng)


def galaxy(n_disk, n_halo, rng, Md=1.0, Rd=1.0, z0=0.1, Rmax=6.0, Mh=5.0, b=2.0, Q=1.3):
    """Exponential stellar disk (sech^2 vertical) inside a live Plummer dark-matter halo."""
    Rg = np.linspace(0, Rmax, 2000)
    cum = 1 - (1 + Rg / Rd) * np.exp(-Rg / Rd)  # enclosed mass of an exponential disk
    R = np.interp(rng.random(n_disk) * cum[-1], cum, Rg)
    phi = rng.uniform(0, 2 * np.pi, n_disk)
    z = z0 * np.arctanh(rng.uniform(-1, 1, n_disk))
    xd = np.array([R * np.cos(phi), R * np.sin(phi), z])
    xh, vh = plummer(n_halo, Mh, b, rng)
    x = np.c_[xd, xh]
    m = np.r_[np.full(n_disk, Md / n_disk), np.full(n_halo, Mh / n_halo)]

    # Rotation curve measured from the particles themselves: mean radial pull on massless test rings
    rr, ang = np.meshgrid(np.linspace(0.05, Rmax, 60), np.linspace(0, 2 * np.pi, 32, endpoint=False))
    ring = np.array([rr * np.cos(ang), rr * np.sin(ang), 0 * rr]).reshape(3, -1)
    a = gravity(np.r_[m, np.zeros(ring.shape[1])], EPS)(np.c_[x, ring])[:, -ring.shape[1]:]
    vc2 = np.maximum(-(a[0] * ring[0] + a[1] * ring[1]).reshape(rr.shape).mean(0), 1e-6)
    r1 = rr[0]
    Om2 = vc2 / r1**2
    kappa = np.sqrt(np.maximum(r1 * np.gradient(Om2, r1) + 4 * Om2, 1e-6))  # epicyclic frequency
    Sig = Md / (2 * np.pi * Rd**2) * np.exp(-r1 / Rd)
    sR = Q * 3.36 * Sig / kappa  # Toomre Q = sigma_R kappa / (3.36 G Sigma)
    sphi = sR * kappa / (2 * np.sqrt(Om2))  # epicyclic ratio
    sz = np.sqrt(np.pi * Sig * z0)  # isothermal sheet
    vc, sR, sphi, sz = (np.interp(R, r1, f) for f in (np.sqrt(vc2), sR, sphi, sz))
    # ponytail: no asymmetric-drift correction, the disk relaxes a little in the first orbit
    vR = sR * rng.normal(size=n_disk)
    vphi = vc + sphi * rng.normal(size=n_disk)
    vd = np.array([vR * np.cos(phi) - vphi * np.sin(phi), vR * np.sin(phi) + vphi * np.cos(phi),
                   sz * rng.normal(size=n_disk)])
    v = np.c_[vd, vh]
    return x - x @ m[:, None] / m.sum(), v - v @ m[:, None] / m.sum(), m  # centre-of-mass frame


def collision_ics(rng, rp=RP, d0=D0, tilt1=(0, 0), tilt2=(60, 30)):
    """Two identical galaxies on a parabolic orbit (Toomre & Toomre 1972). tilt = (inclination, azimuth)
    in degrees of each disk's spin relative to the orbital angular momentum (+z). (0, 0) = prograde."""
    xs, vs, ms, gal, disk, core = [], [], [], [], [], []
    for g, tilt in enumerate((tilt1, tilt2)):
        x, v, m = galaxy(N_DISK, N_HALO, rng)
        rot = Rotation.from_euler("xz", tilt, degrees=True)
        xs.append(rot.apply(x.T).T)
        vs.append(rot.apply(v.T).T)
        ms.append(m)
        gal.append(np.full(len(m), g))
        disk.append(np.arange(len(m)) < N_DISK)
        core.append(disk[-1] & (np.linalg.norm(x, axis=0) < 0.5))  # nucleus tracers
    M1, M2 = ms[0].sum(), ms[1].sum()
    p = 2 * rp  # semi-latus rectum of a parabola
    nu = -np.arccos(p / d0 - 1)  # true anomaly at separation d0, before pericentre
    rel_x = d0 * np.array([np.cos(nu), np.sin(nu), 0])
    rel_v = np.sqrt((M1 + M2) / p) * np.array([-np.sin(nu), 1 + np.cos(nu), 0])
    for i, s in ((0, -M2), (1, M1)):  # both galaxies about the common centre of mass
        xs[i] += (s / (M1 + M2) * rel_x)[:, None]
        vs[i] += (s / (M1 + M2) * rel_v)[:, None]
    t_peri = 0.5 * np.sqrt(p**3 / (M1 + M2)) * (np.tan(-nu / 2) + np.tan(-nu / 2) ** 3 / 3)  # Barker
    return np.hstack(xs), np.hstack(vs), np.hstack(ms), np.hstack(gal), np.hstack(disk), np.hstack(core), t_peri


# ----------------------------------------------------------------------------- the collision
def collide():
    rng = np.random.default_rng(42)
    x, v, m, gal, disk, core, t_peri = collision_ics(rng)
    N = len(m)
    nsteps = int(round(T_END / DT))
    core = [np.flatnonzero(core & (gal == g)) for g in (0, 1)]
    force = gravity(m, EPS, accel_par, np.float32)
    frames, t_frames, t_E, E, sep = [], [], [0.0], [energy(x, v, m, EPS)], []
    t0 = time.perf_counter()
    for step, x, v in leapfrog(x, v, force, DT, nsteps):
        if step % FRAME_EVERY == 0:
            frames.append(x[:, disk].astype(np.float32))
            t_frames.append(step * DT)
            c = [np.median(x[:, idx], axis=1) for idx in core]
            sep.append(np.linalg.norm(c[1] - c[0]))
        if step % 50 == 0:
            t_E.append(step * DT)
            E.append(energy(x, v, m, EPS))
        if step % 500 == 0:
            el = time.perf_counter() - t0
            print(f"[collide] step {step}/{nsteps}  t = {step * DT * T_MYR:6.0f} Myr  "
                  f"{el:6.1f} s  ({el / step * 1e3:.1f} ms/step)")
    wall = time.perf_counter() - t0
    E = np.array(E)
    dE = np.abs(E / E[0] - 1)
    print(f"[collide] N = {N}, {nsteps} steps in {wall:.1f} s ({wall / nsteps * 1e3:.1f} ms/step, "
          f"{N * N * nsteps / wall:.2e} interactions/s)")
    print(f"[collide] energy error max |dE/E| = {dE.max():.2e}, final {dE[-1]:.2e}")
    return dict(frames=np.array(frames), t=np.array(t_frames), gal=gal[disk], sep=np.array(sep),
                t_E=np.array(t_E), dE=dE, t_peri=t_peri, N=N, m_disk=m[0], M=m.sum(), wall=wall)


# ----------------------------------------------------------------------------- rendering
COLORS = np.array([[1.0, 0.60, 0.25], [0.30, 0.62, 1.0]], np.float32)  # galaxy 1 gold, galaxy 2 azure
VIEW = Rotation.from_euler("zx", [90, -30], degrees=True)  # turn tails along the wide axis, tilt 30 deg


def render(p, gal, m_disk, half_w, h, w):
    """Surface density on black with an asinh (log-like) stretch, one hue per galaxy -> (h, w, 3) uint8.

    Two layers per galaxy: diffuse light with SPH-style adaptive smoothing (Gaussian width set by the
    distance to the 8th-nearest neighbour, so sparse tails blur into streams and dense cores stay
    sharp), plus every particle as a sharp 'star'. Pixel scales are relative to 720p.
    """
    q = VIEW.apply(p.T).T
    half_h, s = half_w * h / w, h / 720
    rc = np.c_[(q[1] / half_h + 1) * h / 2, (q[0] / half_w + 1) * w / 2]  # (row, col) in pixels
    to_sigma = m_disk / (2 * half_w / w) ** 2  # particles per pixel -> surface density (N-body units)
    stretch = lambda S, S0: np.arcsinh(S * to_sigma / S0) / np.arcsinh(0.3 / S0)  # 1 at Sigma = 0.3
    img = np.zeros((h, w, 3), np.float32)
    for g, c in enumerate(COLORS):
        pg = rc[gal == g]
        hist = lambda sel: np.histogram2d(pg[sel, 0], pg[sel, 1], bins=(h, w),
                                          range=((0, h), (0, w)))[0].astype(np.float32)
        d8 = cKDTree(pg).query(pg, k=8)[0][:, -1]
        width = np.clip(np.round(np.log2(np.maximum(0.5 * d8 / s, 1))), 1, 3)  # classes: 2, 4, 8 px
        glow = sum(gaussian_filter(hist(width == k), s * 2.0**k, truncate=3) for k in (1, 2, 3))
        stars = gaussian_filter(hist(slice(None)), 0.7 * s)
        img += (stretch(glow, 2e-3) + 0.6 * stretch(stars, 1e-2))[..., None] * c
    return (255 * (1 - np.exp(-1.6 * img))).astype(np.uint8)  # tone map: bright cores go white


def zoom_schedule(frames):
    """Half-width of the view per frame: 1.8 x the 90th-percentile extent, smoothed, kept in [16, 34]."""
    ext = [np.percentile(np.maximum(abs(q[0]), 16 / 9 * abs(q[1])), 90) for q in (VIEW.apply(f.T).T for f in frames)]
    return np.clip(gaussian_filter1d(1.8 * np.array(ext), 30, mode="nearest"), 16, 34)


def _render_args(a):
    return render(*a)


def make_movie(run, path, w=1280, h=720, fps=30):
    """Frames are rendered by 4 worker processes; the main process pipes them to ffmpeg in order."""
    frames, t, gal, m_disk = run["frames"], run["t"], run["gal"], run["m_disk"]
    hw = zoom_schedule(frames)
    fig = plt.figure(figsize=(w / 100, h / 100), dpi=100, facecolor="k")
    ax = fig.add_axes((0, 0, 1, 1))
    ax.set_axis_off()
    im = ax.imshow(np.zeros((h, w, 3), np.uint8), origin="lower", extent=(0, w, 0, h), interpolation="none")
    ax.set(xlim=(0, w), ylim=(0, h))
    txt = ax.text(24, h - 42, "", color="w", fontsize=15, family="monospace")
    ax.text(24, h - 66, "gold: galaxy 1, disk in the orbital plane   azure: galaxy 2, disk tilted 60 deg",
            color="#8d99ae", fontsize=10)
    ax.text(w - 24, 20, f"N = {run['N']:,} bodies  |  direct summation + leapfrog  |  Numba, 4 threads",
            ha="right", color="#8d99ae", fontsize=10)
    bar, = ax.plot([], [], color="w", lw=2)
    ax.text(24, 32, "20 kpc", color="w", fontsize=10)
    writer = FFMpegWriter(fps=fps, bitrate=5000, extra_args=["-pix_fmt", "yuv420p"])
    with Pool(4) as pool, writer.saving(fig, str(path), dpi=100):
        jobs = ((f, gal, m_disk, s, h, w) for f, s in zip(frames, hw))
        for img, tt, s in zip(pool.imap(_render_args, jobs, chunksize=2), t, hw):
            im.set_data(img)
            txt.set_text(f"t = {tt * T_MYR / 1000:4.2f} Gyr")
            bar.set_data([24, 24 + 20 / L_KPC / (2 * s) * w], [24, 24])  # 20 kpc in pixels
            writer.grab_frame(facecolor="k")
    plt.close(fig)


def make_stills(run):
    frames, t, gal, m_disk, sep = run["frames"], run["t"], run["gal"], run["m_disk"], run["sep"]
    hw = zoom_schedule(frames)
    # hero: ~230 Myr after first pericentre, when the tails and the bridge are at their most striking
    i = min(int(np.argmin(sep[: len(sep) // 2])) + int(13 / (DT * FRAME_EVERY)), len(frames) - 1)
    W, H = 2400, 1350
    fig = plt.figure(figsize=(W / 150, H / 150), dpi=150, facecolor="k")
    ax = fig.add_axes((0, 0, 1, 1))
    ax.set_axis_off()
    ax.imshow(render(frames[i], gal, m_disk, 1.2 * hw[i], H, W), origin="lower", interpolation="none")
    ax.text(0.02, 0.945, "Galactic collision", transform=ax.transAxes, color="w", fontsize=22)
    ax.text(0.02, 0.905, f"two disk galaxies in live dark-matter halos, {run['N']:,} bodies, "
            f"t = {t[i] * T_MYR / 1000:.2f} Gyr", transform=ax.transAxes, color="#8d99ae", fontsize=11)
    fig.savefig(OUT / "hero.png", dpi=150, facecolor="k")
    plt.close(fig)

    fig, axs = plt.subplots(2, 3, figsize=(15, 6.6), facecolor="k", layout="constrained")
    for a, j in zip(axs.flat, np.linspace(0, len(frames) - 1, 6).astype(int)):
        s = hw[j] * L_KPC
        a.imshow(render(frames[j], gal, m_disk, hw[j], 360, 640), origin="lower",
                 extent=(-s, s, -s * 9 / 16, s * 9 / 16), interpolation="none")
        a.set_title(f"t = {t[j] * T_MYR / 1000:.2f} Gyr", color="w", fontsize=11)
        a.tick_params(colors="#8d99ae", labelsize=8)
        a.set_xlabel("x (kpc)", color="#8d99ae", fontsize=9)
        a.set_ylabel("y (kpc)", color="#8d99ae", fontsize=9)
    fig.suptitle("Approach, first pericentre, tidal tails + bridge, apocentre, second passage, merger remnant",
                 color="w")
    fig.savefig(OUT / "sequence.png", dpi=110, facecolor="k")
    plt.close(fig)
    return i


def diagnostics(run):
    """Nucleus separation vs the point-mass (Keplerian) parabola, and energy conservation of the run."""
    t, sep, M, t_peri = run["t"], run["sep"], run["M"], run["t_peri"]
    rp, p = RP, 2 * RP
    D = np.linspace(-np.tan(np.arccos(p / D0 - 1) / 2), 6, 2000)  # Barker's equation for the parabola
    tk = t_peri + 0.5 * np.sqrt(p**3 / M) * (D + D**3 / 3)
    i = int(np.argmin(sep[: len(sep) // 2]))
    j = i + int(np.argmax(sep[i:]))
    k = int(np.flatnonzero(sep > 2 / L_KPC)[-1]) + 1  # nuclei stay within 2 kpc after this
    merged = f"from t = {t[k] * T_MYR / 1e3:.2f} Gyr" if k < len(t) else "not yet (run too short)"
    print(f"[diag] first pericentre t = {t[i] * T_MYR / 1e3:.2f} Gyr at {sep[i] * L_KPC:.1f} kpc "
          f"(Kepler parabola: {t_peri * T_MYR / 1e3:.2f} Gyr at {RP * L_KPC:.1f} kpc); "
          f"apocentre {sep[j] * L_KPC:.0f} kpc at {t[j] * T_MYR / 1e3:.2f} Gyr; "
          f"nuclei within 2 kpc for good: {merged}")
    fig, (a1, a2) = plt.subplots(1, 2, figsize=(13, 4.8))
    a1.plot(t * T_MYR / 1000, sep * L_KPC, color="#f4a261", lw=2, label="N-body: live halos")
    ok = tk <= t[-1]
    a1.plot(tk[ok] * T_MYR / 1000, rp * (1 + D[ok] ** 2) * L_KPC, "--", color="#4cc9f0",
            label="point masses (Kepler parabola)")
    a1.set(xlabel="time (Gyr)", ylabel="nucleus separation (kpc)",
           title="Dynamical friction: halos soak up orbital energy -> merger")
    a1.legend()
    a2.semilogy(run["t_E"][1:] * T_MYR / 1000, run["dE"][1:], color="#f72585")
    a2.set(xlabel="time (Gyr)", ylabel="|E(t)/E0 - 1|",
           title=f"Energy conservation, KDK leapfrog, dt = {DT * T_MYR:.2f} Myr, float32 forces")
    for a in (a1, a2):
        a.grid(alpha=0.2)
    fig.tight_layout()
    fig.savefig(OUT / "collision_diagnostics.png", dpi=130)
    plt.close(fig)


def main():
    OUT.mkdir(parents=True, exist_ok=True)
    plt.style.use("dark_background")
    T0 = time.perf_counter()
    stamp = lambda what: print(f"[time] {what} done, {time.perf_counter() - T0:.0f} s elapsed")
    print(f"units: length {L_KPC} kpc, velocity {V_KMS:.0f} km/s, time {T_MYR:.2f} Myr, mass {M_SUN:.0e} Msun")
    if not QUICK:
        benchmark()  # first, so the JIT compile times it prints are real
        stamp("benchmark")
    verify()
    stamp("verification")
    run = collide()
    stamp("collision")
    diagnostics(run)
    i = make_stills(run)
    print(f"[render] hero frame t = {run['t'][i] * T_MYR / 1000:.2f} Gyr")
    make_movie(run, OUT / "collision.mp4")
    print(f"[render] MP4: {len(run['frames'])} frames, {(OUT / 'collision.mp4').stat().st_size / 1e6:.1f} MB")
    stamp("rendering (total runtime)")


if __name__ == "__main__":
    main()