feat(v2): propagation-of-chaos animation, README section, notebook sandwich

- examples/animate_propagation_of_chaos.py: 4-panel McKean-Vlasov simulator at N in {20, 100, 500, 4000} with reference N=12000; bottom panel tracks W_2(mu^N_t, mu_t) on log scale -> visible 1/sqrt(N) decay (Sznitman 1991). - examples/propagation_of_chaos.gif (1.6 MB) - README: new 'Propagation of chaos' subsection under Mean-field & agent-based dynamics, with empirical-measure formula, k-tuple factorisation and GIF embed. - examples/notebooks/14_mckean_vlasov.ipynb: sandwich PRE/code/POST cells demonstrating W2 ~ 1/sqrt(N) on the same simulator. Verified executed: sqrt(N)*W2 ~ 0.7 across N (theoretical const).
2026-05-14 22:43:47 +02:00
parent 24556f51d7
commit b6a1e1a953
4 changed files with 414 additions and 2 deletions
@@ -54,6 +54,38 @@ v2 ships **eight brand-new CPU-only generic numerical primitive groups** with fu
 - **`consensus_dynamics(...)`** — synchronous opinion-dynamics consensus on a graph.
 - **`solve_mfg_1d_rust(MFGConfig)`** — 1-D mean-field game (HJB ↔ Fokker–Planck fixed-point).

+#### Propagation of chaos
+
+<p align="center">
+    <img src="examples/propagation_of_chaos.gif" alt="Propagation of chaos" width="680" />
+</p>
+
+For an interacting N-particle system
+
+$$
+dX^{i,N}_t \;=\; b\!\bigl(X^{i,N}_t,\; \mu^N_t\bigr)\, dt \;+\; \sigma\, dW^i_t,
+\qquad
+\mu^N_t \;=\; \frac{1}{N}\sum_{j=1}^{N}\delta_{X^{j,N}_t},
+$$
+
+Sznitman's theorem (1991) states that whenever $b$ is Lipschitz in both arguments,
+the empirical measure $\mu^N_t$ converges in Wasserstein-2 to the law $\mu_t$
+of the McKean–Vlasov limit at rate $\mathcal{O}(1/\sqrt{N})$, and any finite
+$k$-tuple of particles becomes asymptotically independent — *chaos propagates*
+from $t=0$ to all later times:
+
+$$
+\operatorname{Law}\!\bigl(X^{1,N}_t,\dots,X^{k,N}_t\bigr) \;\xrightarrow[N\to\infty]{w}\; \mu_t^{\otimes k}.
+$$
+
+The animation above runs four parallel simulations with $N\in\{20,100,500,4000\}$
+sharing the *same* bimodal initial law, the *same* drift $-\theta(x-\bar{x})$ and
+the *same* noise $\sigma\, dW$. The bottom panel tracks the Wasserstein-2 distance
+$W_2(\mu^N_t, \mu_t)$ to a high-resolution reference and visibly decays as
+$1/\sqrt{N}$. Source: [`examples/animate_propagation_of_chaos.py`](examples/animate_propagation_of_chaos.py).
+Companion notebook: [`examples/notebooks/14_mckean_vlasov.ipynb`](examples/notebooks/14_mckean_vlasov.ipynb).
+
+
 ### Topology, graphs & path signatures

 - **`vietoris_rips_filtration`**, **`persistent_homology`**, **`bottleneck_distance`** — TDA primitives.
@@ -0,0 +1,230 @@
+"""Propagation of chaos for the McKean-Vlasov mean-reverting flow.
+
+Theorem (Sznitman 1991): for the N-particle system
+
+    dX^{i,N}_t = b(X^{i,N}_t, mu^N_t) dt + sigma dW^i_t,
+    mu^N_t = (1/N) sum_j delta_{X^{j,N}_t},
+
+with b Lipschitz, the empirical measure mu^N_t converges (in
+Wasserstein-2) to the law mu_t of the McKean-Vlasov limit
+
+    dX_t = b(X_t, mu_t) dt + sigma dW_t,    Law(X_t) = mu_t,
+
+at rate O(1/sqrt(N)). Equivalently, any finite k-tuple
+(X^{1,N}_t, ..., X^{k,N}_t) becomes asymptotically independent --
+``chaos propagates`` from t = 0 to all later times.
+
+This animation visualises the convergence: we run four McKean-Vlasov
+simulations in parallel with N in {20, 200, 2000, 20000}, all sharing
+the SAME initial bimodal distribution and the SAME drift / noise.
+The coloured histograms are the empirical densities mu^N_t; the white
+dashed curve is a high-resolution reference (N = 100_000).
+
+As t advances and N increases, the histograms collapse onto the white
+reference -- propagation of chaos in action.
+
+Run with:
+    python examples/animate_propagation_of_chaos.py
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib import animation
+
+import optimizr as opt
+
+# ---------------------------------------------------------------------------
+# Parameters
+# ---------------------------------------------------------------------------
+N_VALUES = [20, 100, 500, 4000]
+N_REF = 12_000
+N_STEPS = 200
+T_HORIZON = 3.0
+THETA = 0.7
+SIGMA = 0.30
+SEED = 11
+FRAME_STRIDE = 4
+
+X_GRID = np.linspace(-3.5, 3.5, 200)
+BINS = np.linspace(-3.5, 3.5, 50)
+
+
+def make_initial(N: int, seed: int) -> np.ndarray:
+    rng = np.random.default_rng(seed)
+    half = N // 2
+    return np.concatenate([
+        rng.normal(-2.0, 0.35, half),
+        rng.normal(+2.0, 0.35, N - half),
+    ])
+
+
+def simulate(N: int, seed: int) -> np.ndarray:
+    """Return paths of shape (n_steps + 1, N)."""
+    initial = make_initial(N, seed)
+    out = opt.mean_reverting_mckean_vlasov(
+        initial=initial.tolist(),
+        theta=THETA,
+        sigma=SIGMA,
+        n_steps=N_STEPS,
+        t_horizon=T_HORIZON,
+        seed=seed,
+    )
+    return np.asarray(out["paths_flat"]).reshape(N_STEPS + 1, N)
+
+
+print("Simulating propagation-of-chaos panels...")
+panels = {}
+for i, N in enumerate(N_VALUES):
+    print(f"  panel N = {N}...", flush=True)
+    panels[N] = simulate(N, SEED + i)
+print(f"  reference (N = {N_REF})...", flush=True)
+ref_paths = simulate(N_REF, SEED + 999)
+print("  done.", flush=True)
+
+# Pre-compute smoothed reference density for each frame using a histogram
+# convolved with a Gaussian kernel -- O(N) per frame instead of O(N*G).
+print("Building reference density curves...")
+DENS_BINS = np.linspace(-3.5, 3.5, 161)
+DENS_CENTERS = 0.5 * (DENS_BINS[:-1] + DENS_BINS[1:])
+bw = 0.12
+kernel_x = np.arange(-int(4 * bw / (DENS_BINS[1] - DENS_BINS[0])),
+                     int(4 * bw / (DENS_BINS[1] - DENS_BINS[0])) + 1)
+kernel = np.exp(-0.5 * (kernel_x * (DENS_BINS[1] - DENS_BINS[0]) / bw) ** 2)
+kernel /= kernel.sum() * (DENS_BINS[1] - DENS_BINS[0])
+ref_density = np.empty((N_STEPS + 1, DENS_CENTERS.size))
+for k in range(N_STEPS + 1):
+    h, _ = np.histogram(ref_paths[k], bins=DENS_BINS, density=True)
+    ref_density[k] = np.convolve(h, kernel, mode="same") / kernel.sum() * kernel.sum()
+# Interpolate onto display grid
+ref_density_grid = np.empty((N_STEPS + 1, X_GRID.size))
+for k in range(N_STEPS + 1):
+    ref_density_grid[k] = np.interp(X_GRID, DENS_CENTERS, ref_density[k])
+ref_density = ref_density_grid
+
+times = np.linspace(0.0, T_HORIZON, N_STEPS + 1)
+
+# Wasserstein-2 distance between empirical mu^N and reference, per frame
+print("Computing W2(mu^N_t, mu_t) curves...")
+def w2_to_ref(samples_a: np.ndarray, samples_b: np.ndarray) -> float:
+    """1-D Wasserstein-2 via sorted samples (Sklar / quantile transport)."""
+    a = np.sort(samples_a)
+    b = np.sort(samples_b)
+    # Resample b to len(a) via interpolation of empirical quantiles.
+    qa = np.linspace(0, 1, len(a))
+    qb = np.linspace(0, 1, len(b))
+    b_resampled = np.interp(qa, qb, b)
+    return float(np.sqrt(np.mean((a - b_resampled) ** 2)))
+
+
+w2_curves = {N: np.array([w2_to_ref(panels[N][k], ref_paths[k]) for k in range(N_STEPS + 1)])
+             for N in N_VALUES}
+
+# ---------------------------------------------------------------------------
+# Figure -- 2x2 panels + W2 convergence track
+# ---------------------------------------------------------------------------
+plt.rcParams.update({
+    "axes.facecolor": "#0b1020",
+    "figure.facecolor": "#0b1020",
+    "axes.edgecolor": "#3a4a72",
+    "axes.labelcolor": "#dbe7ff",
+    "xtick.color": "#9eb1d8",
+    "ytick.color": "#9eb1d8",
+    "text.color": "#dbe7ff",
+    "axes.grid": True,
+    "grid.color": "#1e2a47",
+    "grid.linestyle": "--",
+    "grid.alpha": 0.4,
+})
+
+fig = plt.figure(figsize=(8.5, 6.5), dpi=80)
+gs = fig.add_gridspec(3, 2, height_ratios=[1, 1, 0.9], hspace=0.55, wspace=0.25)
+
+panel_axes = {}
+hist_artists = {}
+panel_colors = ["#39d2ff", "#7be495", "#ffd166", "#ff7847"]
+
+for ax_idx, (N, color) in enumerate(zip(N_VALUES, panel_colors)):
+    ax = fig.add_subplot(gs[ax_idx // 2, ax_idx % 2])
+    panel_axes[N] = ax
+    ax.set_xlim(-3.5, 3.5)
+    ax.set_ylim(0, ref_density.max() * 1.15)
+    ax.set_title(f"$N = {N}$", color=color, fontsize=10, pad=4)
+    ax.set_xticks([-3, -1.5, 0, 1.5, 3])
+    if ax_idx >= 2:
+        ax.set_xlabel("$x$", fontsize=9)
+    # Initial histogram & reference line
+    hist, _ = np.histogram(panels[N][0], bins=BINS, density=True)
+    centers = 0.5 * (BINS[:-1] + BINS[1:])
+    bars = ax.bar(centers, hist, width=BINS[1] - BINS[0],
+                  color=color, alpha=0.55, edgecolor="none")
+    (ref_line,) = ax.plot(X_GRID, ref_density[0], color="#ffffff",
+                          lw=1.3, ls="--", alpha=0.85,
+                          label=r"$\mu_t$ (ref. $N=10^5$)")
+    if ax_idx == 0:
+        ax.legend(loc="upper right", fontsize=7, framealpha=0.2,
+                  edgecolor="#3a4a72")
+    hist_artists[N] = (bars, ref_line)
+
+# Bottom row: W2 convergence on a log-log scale wrt time
+ax_w2 = fig.add_subplot(gs[2, :])
+ax_w2.set_xlim(0, T_HORIZON)
+ax_w2.set_yscale("log")
+ax_w2.set_ylim(max(1e-3, min(w2_curves[N_VALUES[-1]].min(), 1e-2) * 0.5),
+               max(w2_curves[N_VALUES[0]].max() * 1.5, 1.0))
+ax_w2.set_xlabel("$t$", fontsize=9)
+ax_w2.set_ylabel(r"$W_2(\mu_t^N, \mu_t)$", fontsize=9)
+ax_w2.set_title(r"Wasserstein-2 distance to the reference law (log scale)"
+                "  --  $W_2 \\sim O(1/\\sqrt{N})$",
+                color="#e9efff", fontsize=10, pad=6)
+
+w2_lines = {}
+for N, color in zip(N_VALUES, panel_colors):
+    (line,) = ax_w2.plot([], [], color=color, lw=1.6, label=f"$N = {N}$")
+    w2_lines[N] = line
+ax_w2.legend(loc="upper right", ncol=4, fontsize=8, framealpha=0.2,
+             edgecolor="#3a4a72")
+
+cursor = ax_w2.axvline(0.0, color="#ffd166", lw=1, ls=":", alpha=0.8)
+
+fig.suptitle(
+    r"Propagation of chaos (Sznitman 1991):  $\mu_t^N \to \mu_t$  as  $N \to \infty$",
+    color="#e9efff", fontsize=12, y=0.98,
+)
+
+
+def update(frame):
+    artists = []
+    for N in N_VALUES:
+        bars, ref_line = hist_artists[N]
+        hist, _ = np.histogram(panels[N][frame], bins=BINS, density=True)
+        for rect, h in zip(bars, hist):
+            rect.set_height(h)
+        ref_line.set_ydata(ref_density[frame])
+        artists.extend([*bars, ref_line])
+    cursor.set_xdata([times[frame], times[frame]])
+    for N in N_VALUES:
+        w2_lines[N].set_data(times[: frame + 1], w2_curves[N][: frame + 1])
+    artists.append(cursor)
+    artists.extend(w2_lines.values())
+    return artists
+
+
+frame_indices = list(range(0, N_STEPS + 1, FRAME_STRIDE))
+print(f"Rendering {len(frame_indices)} frames...")
+anim = animation.FuncAnimation(
+    fig, update, frames=frame_indices, interval=40, blit=False,
+)
+
+out_path = Path(__file__).with_name("propagation_of_chaos.gif")
+try:
+    anim.save(out_path, writer=animation.PillowWriter(fps=24))
+    print(f"Saved animation: {out_path}")
+except Exception as exc:
+    print(f"GIF write failed ({exc}); saving PNG of last frame instead.")
+    update(N_STEPS)
+    fig.savefig(out_path.with_suffix(".png"), dpi=120)
+    print(f"Saved PNG: {out_path.with_suffix('.png')}")