From cc4949d7c972d018ec215511c3ba79d004cd900f Mon Sep 17 00:00:00 2001
From: Hydrodynamical <joseph.kent.miller@gmail.com>
Date: Tue, 20 Oct 2026 20:02:23 -0400
Subject: [PATCH 1/2] Extract reusable graphon control core and validate
 numerical workflows

---
 .github/workflows/tests.yml                   |   21 +
 README.md                                     |   11 +-
 collect_paper_figures.sh                      |   36 +-
 docs/DEVELOPMENT.md                           |   97 ++
 docs/NUMERICS.md                              |   75 ++
 e1_kuramoto.py                                |  782 +-----------
 e1_msweep.py                                  |  377 +++---
 e1_noise_figure.py                            |  159 ++-
 e2_bc.py                                      |  895 +-------------
 graphon_control/__init__.py                   |   15 +
 graphon_control/artifacts.py                  |  278 +++++
 graphon_control/cli.py                        |   54 +
 graphon_control/controls.py                   |  200 +++
 graphon_control/diagnostics.py                |   33 +
 graphon_control/dynamics.py                   |  125 ++
 graphon_control/experiments/__init__.py       |    1 +
 .../experiments/bounded_confidence.py         |  537 ++++++++
 .../experiments/bounded_confidence_plots.py   |  254 ++++
 graphon_control/experiments/generic.py        |  256 ++++
 graphon_control/experiments/kuramoto.py       |  463 +++++++
 graphon_control/experiments/kuramoto_plots.py |  214 ++++
 graphon_control/initial_conditions.py         |   72 ++
 graphon_control/interactions.py               |   24 +
 graphon_control/objectives.py                 |  179 +++
 graphon_control/plotting.py                   |  508 ++++++++
 graphon_control/training.py                   |   76 ++
 graphon_control/validation.py                 |   83 ++
 model_1.py                                    | 1083 +----------------
 model_1_viz.py                                |  500 +-------
 pyproject.toml                                |   31 +
 requirements-tested.txt                       |    7 +
 tests/conftest.py                             |    7 +
 tests/test_artifacts.py                       |  228 ++++
 tests/test_controls.py                        |  194 +++
 tests/test_experiments.py                     |  148 +++
 tests/test_numerics.py                        |  211 ++++
 tests/test_package.py                         |   95 ++
 tests/test_plotting.py                        |   80 ++
 tests/test_training.py                        |   65 +
 tests/test_workflows.py                       |  133 ++
 40 files changed, 5140 insertions(+), 3467 deletions(-)
 create mode 100644 .github/workflows/tests.yml
 create mode 100644 docs/DEVELOPMENT.md
 create mode 100644 docs/NUMERICS.md
 create mode 100644 graphon_control/__init__.py
 create mode 100644 graphon_control/artifacts.py
 create mode 100644 graphon_control/cli.py
 create mode 100644 graphon_control/controls.py
 create mode 100644 graphon_control/diagnostics.py
 create mode 100644 graphon_control/dynamics.py
 create mode 100644 graphon_control/experiments/__init__.py
 create mode 100644 graphon_control/experiments/bounded_confidence.py
 create mode 100644 graphon_control/experiments/bounded_confidence_plots.py
 create mode 100644 graphon_control/experiments/generic.py
 create mode 100644 graphon_control/experiments/kuramoto.py
 create mode 100644 graphon_control/experiments/kuramoto_plots.py
 create mode 100644 graphon_control/initial_conditions.py
 create mode 100644 graphon_control/interactions.py
 create mode 100644 graphon_control/objectives.py
 create mode 100644 graphon_control/plotting.py
 create mode 100644 graphon_control/training.py
 create mode 100644 graphon_control/validation.py
 create mode 100644 pyproject.toml
 create mode 100644 requirements-tested.txt
 create mode 100644 tests/conftest.py
 create mode 100644 tests/test_artifacts.py
 create mode 100644 tests/test_controls.py
 create mode 100644 tests/test_experiments.py
 create mode 100644 tests/test_numerics.py
 create mode 100644 tests/test_package.py
 create mode 100644 tests/test_plotting.py
 create mode 100644 tests/test_training.py
 create mode 100644 tests/test_workflows.py

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
new file mode 100644
index 0000000..112bcd8
--- /dev/null
+++ b/.github/workflows/tests.yml
@@ -0,0 +1,21 @@
+name: Numerical checks
+on: [push, pull_request]
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    env:
+      MPLBACKEND: Agg
+      OMP_NUM_THREADS: '1'
+      MKL_NUM_THREADS: '1'
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+          cache: pip
+      - run: python -m pip install --index-url https://download.pytorch.org/whl/cpu torch==2.14.1
+      - run: python -m pip install -r requirements-tested.txt
+      - run: python -m pip install --no-deps -e .
+      - run: python -m pytest
+      - run: python e1_kuramoto.py --smoke --no-plots --no-mirror --out-root smoke-results
+      - run: python e2_bc.py --smoke --no-plots --no-mirror --out-root smoke-results
diff --git a/README.md b/README.md
index 8b1d114..9df3ad6 100644
--- a/README.md
+++ b/README.md
@@ -4,7 +4,7 @@
 
 A PyTorch research implementation that learns how agents influence one another over time. It combines neural interaction kernels, differentiable simulation, and stochastic optimization.
 
-Companion code for *Graphon Controlled Opinion Dynamics*, by **Joseph K. Miller and Brandon Su**.
+Companion code for *Optimal Graphon Control of Collective Dynamics*, by **Joseph K. Miller and Brandon Su**.
 
 ## Experiments
 
@@ -26,22 +26,23 @@ From the repository root:
 ```bash
 python3 -m venv .venv
 source .venv/bin/activate
-python -m pip install torch numpy matplotlib tqdm pillow
+python -m pip install -e .
 python e1_kuramoto.py --smoke --device cpu --no-mirror
 ```
 
-This runs a 20-iteration workflow check. Results appear under `out_1/`, including `untwist.gif`, `kernel_time.gif`, and `metrics.json`.
+This runs a 20-iteration workflow check on a small grid and network. Results appear under `out_1/`, including `untwist.gif`, `kernel_time.gif`, and `metrics.json`.
 
 ## Code
 
 | File | Contents |
 | --- | --- |
-| `model_1.py` | Neural controller, simulation, and training |
+| `graphon_control/` | Reusable controller, dynamics, objectives, and training |
+| `model_1.py` | General experiment entry point |
 | `e1_kuramoto.py` | Oscillator experiment |
 | `e2_bc.py` | Opinion dynamics experiment |
 | `model_1_viz.py` | Visualization utilities |
 
-Use `python <script> --help` for experiment options.
+Use `python <script> --help` for experiment options. See [development](docs/DEVELOPMENT.md) for testing and reproduction, and [numerical conventions](docs/NUMERICS.md) for the connection to the paper.
 
 ## Research scope
 
diff --git a/collect_paper_figures.sh b/collect_paper_figures.sh
index 1f91403..23ed967 100755
--- a/collect_paper_figures.sh
+++ b/collect_paper_figures.sh
@@ -1,28 +1,10 @@
 #!/bin/sh
-# Collect the figures used by main5.tex into the flat folder figures/paper/
-# (the folder to drag into Overleaf). Re-run after regenerating any experiment
-# figures; the experiment mirror dirs (figures/e1_*, figures/e2_*) stay the
-# archival source of truth.
-set -e
-cd "$(dirname "$0")"
-mkdir -p figures/paper
-
-# schematic / example figures (names unchanged)
-for f in x0-histogram w1 w2 graphon-N graphon-2N; do
-  cp "figures/$f.pdf" "figures/paper/$f.pdf"
-done
-
-# bounded-confidence experiment (subsec:exp-bc)
-cp figures/e2_eps1_T10/opinion_traj.pdf     figures/paper/bc_opinion_traj.pdf
-cp figures/e2_eps1_T10/state_evolution.pdf  figures/paper/bc_state_evolution.pdf
-cp figures/e2_eps1_T10/kernel_strip.pdf     figures/paper/bc_kernel_strip.pdf
-
-# Kuramoto experiment (subsec:exp-kuramoto)
-cp figures/e1_k3_T20/r_panels.pdf           figures/paper/kur_r_panels.pdf
-cp figures/e1_k3_T20/state_evolution.pdf    figures/paper/kur_state_evolution.pdf
-cp figures/e1_k3_T20/kernel_strip.pdf       figures/paper/kur_kernel_strip.pdf
-cp figures/e1_k3_T20/winding.pdf            figures/paper/kur_winding.pdf
-cp figures/e1_k2_T12_sigma0p3/r_noise_panel.pdf figures/paper/kur_noise_panel.pdf
-
-echo "figures/paper/ refreshed:"
-ls figures/paper/
+# Select source runs and schematics explicitly in a JSON manifest. All inputs
+# are checked before any destination is written; existing exports need --overwrite.
+set -eu
+if [ "$#" -eq 0 ]; then
+  echo "Usage: $0 MANIFEST.json [--output-dir PATH] [--overwrite]" >&2
+  exit 2
+fi
+script_dir=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
+exec "${PYTHON:-python3}" "$script_dir/graphon_control/artifacts.py" collect "$@"
diff --git a/docs/DEVELOPMENT.md b/docs/DEVELOPMENT.md
new file mode 100644
index 0000000..840316d
--- /dev/null
+++ b/docs/DEVELOPMENT.md
@@ -0,0 +1,97 @@
+# Developing and reproducing experiments
+
+Install the package and development tools from the repository root:
+
+```bash
+python -m pip install -e '.[dev]'
+python -m pytest
+```
+
+`requirements-tested.txt` records the direct dependency versions of the
+verified Python 3.12 environment. It is a tested snapshot, not a full transitive
+or cross-platform lockfile. CPU tests run in GitHub Actions; accelerator tests
+run when compatible hardware is available.
+
+## Structure
+
+- `graphon_control/controls.py`: neural time/label kernel.
+- `dynamics.py`, `objectives.py`, `diagnostics.py`: shared scientific operations.
+- `training.py`: Adam loop with finite-loss and finite-gradient guards.
+- `artifacts.py`: new run directories, provenance, checkpoints and figure export.
+- `interactions.py`, `initial_conditions.py`: reusable example ingredients.
+- `experiments/`: model-specific configurations, baselines and orchestration.
+- `plotting.py`: visualization, separate from the importable numerical core.
+- Root Python scripts preserve the previous command and import entry points.
+
+To add a model, supply a differentiable interaction function, an initial-state
+sampler and an objective to `train_controller`; reuse the shared rollout and
+cost. Keep model-specific diagnostics beside the experiment. Numerical
+conventions and the relationship to the paper are in [NUMERICS.md](NUMERICS.md).
+
+## Run and inspect
+
+```bash
+python e1_kuramoto.py --smoke --no-mirror
+python e2_bc.py --smoke --no-mirror
+python model_1.py --smoke --no-plots
+```
+
+Smoke runs perform 20 updates on a smaller grid/network and are workflow checks,
+not reproductions of manuscript results. `--no-plots` skips rendering only;
+configuration, weights, trajectories and diagnostics are still saved.
+
+For the full paper configurations:
+
+```bash
+python e2_bc.py --tag paper_bc
+python e1_kuramoto.py --k 3 --T 20 --steps 3000 --tag paper_kuramoto
+```
+
+Run tags reserve new directories and existing tags are refused. By default,
+figures stay in their run directory. `--mirror-figures` also writes into a
+new figure directory named for the run, never into a directory selected only
+by model parameters.
+
+Every run records a complete `config.json`, `provenance.json` (code revision,
+working-tree state, dependency versions and training/evaluation seeds), raw
+controller `model.pt`, and numerical artifacts. `complete.json` is written only
+after the requested outputs succeed. Training and evaluation have independent
+seeds; reproducibility across different hardware/library versions is not assumed.
+
+## Load settings and regenerate plots
+
+```bash
+python e1_kuramoto.py --config out_1/e1_paper_kuramoto/config.json --tag rerun
+python e1_kuramoto.py --replot out_1/e1_paper_kuramoto
+python e2_bc.py --replot out_1/e2_paper_bc
+```
+
+Explicit scientific CLI flags override saved settings. Replotting uses saved
+arrays and never retrains. It replaces plots within the selected run; original
+weights, arrays and metrics are preserved. Generic runs also archive sampled
+interaction values so replotting does not need to reconstruct arbitrary Python
+callables. Existing raw state_dict checkpoints remain readable through
+`artifacts.load_controller`.
+
+## Collect manuscript figures
+
+Choose exact artifacts in a JSON manifest, with paths relative to the manifest:
+
+```json
+{
+  "output_dir": "figures/paper",
+  "figures": {
+    "kur_winding.pdf": "out_1/e1_paper_kuramoto/winding.pdf",
+    "bc_opinion_traj.pdf": "out_1/e2_paper_bc/opinion_traj.pdf"
+  }
+}
+```
+
+```bash
+./collect_paper_figures.sh paper-figures.json
+```
+
+Collection checks every input and destination before copying. Existing exports
+require `--overwrite`. Add schematic PDFs explicitly once available; the public
+checkout does not contain or generate those manuscript assets. The noise panel
+helper checks compatible noisy/noiseless run settings before comparing them.
diff --git a/docs/NUMERICS.md b/docs/NUMERICS.md
new file mode 100644
index 0000000..e082eb5
--- /dev/null
+++ b/docs/NUMERICS.md
@@ -0,0 +1,75 @@
+# Numerical conventions
+
+The reference is *Optimal Graphon Control of Collective Dynamics*, Joseph K.
+Miller and Brandon Su, October 20, 2026, especially Appendix 5. The manuscript
+and its schematic PDFs are maintained separately from this code repository.
+
+## Kernel representation
+
+Rows are receivers, columns are senders. On the uniform grid of N agents,
+`W[i,j] = wbar[i,j]/N` (equation 64). The directed controller has nonnegative
+weights and unit row sums (equation 67); self-interactions are permitted.
+The uniform graphon `w = 1` is `W[i,j] = 1/N`, including the diagonal.
+
+Consequently equation 59 is implemented as
+
+```text
+x_next[i] = x[i] + dt * sum_j W[i,j] * phi(x[j] - x[i]) + sigma * delta_B[i]
+```
+
+There is no additional 1/N in the drift. Brownian increments have variance dt.
+`rollout_with_Ws(..., brownian_increments=...)` accepts physical increments,
+not standard normal draws. A supplied generator controls automatic draws;
+deterministic rollouts consume no random numbers.
+
+Shapes: initial states `(batch,N)`, sampled kernels `(NT,N,N)`, trajectories
+`(NT+1,batch,N)`. Scalar time returns `(N,N)`; a time vector returns `(times,N,N)`,
+including a vector with one entry. The horizon must satisfy `T/dt` integral;
+invalid grids are rejected instead of silently rounding the endpoint.
+
+## Objective
+
+Equation 60 uses the population variance (`unbiased=False`), the left endpoints
+of the time grid, and the kernel density penalty
+
+```text
+L2(wbar)^2 = (1/N^2) * sum_ij wbar[i,j]^2 = sum_ij W[i,j]^2
+J = sum_n dt * (Var(x[n]) + beta * sum_ij W[n,i,j]^2)
+    + gamma_target * (mean(x[-1]) - x_target)^2
+```
+
+`trajectory_cost` is the single implementation; `paper_cost` adapts it to the
+experiment configuration. The generic trainer uses the same beta coefficient.
+The old generic prototype used beta/2: to retain that historical regularization
+strength, pass half the former beta. The paper experiment coefficients and
+defaults are unchanged.
+
+The generic trainer also exposes named extensions: terminal variance,
+uniform-kernel penalty (eta_uniform/2), scheduled mean targets and ramped
+running variance. These are disabled or configured explicitly. Checkpoint
+targets are snapped to the nearest grid point, and the resolved times should
+be used when interpreting them.
+
+## Constraints and diagnostics
+
+Directed row softmax enforces nonnegativity and unit row sums. With
+`score_bound=S`, bounded scores give `N*W <= exp(2*S)`; the default `None`
+preserves the original unconstrained-score experiment architecture. A fixed
+attention cap is therefore not claimed for default runs across refinements.
+
+Receiver variation is measured by equation 14 on the interval:
+
+```text
+sum_i sqrt(mean_j((N*W[i+1,j] - N*W[i,j])^2))
+```
+
+It is not a training constraint. A finite trained result does not certify a
+continuum optimum or the theorem's fixed BV class. The optional symmetric
+controller normalizes the average row sum, not each row, and is an exploratory
+extension outside the paper's directed row-stochastic class.
+
+The compact-support opinion law is equation 68. The Gaussian fallback has
+nonzero tails, so its gap-threshold cluster count is descriptive and does not
+certify exact isolation or freezing. The Kuramoto experiments use identical
+natural frequencies in a rotating frame; order-parameter alignment alone does
+not distinguish phase differences that are integer multiples of 2*pi.
diff --git a/e1_kuramoto.py b/e1_kuramoto.py
index eb713ac..2687c06 100644
--- a/e1_kuramoto.py
+++ b/e1_kuramoto.py
@@ -1,776 +1,14 @@
-# e1_kuramoto.py
-#
-# Experiment E1: untwisting a twisted state (Kuramoto).
-# Paper reference: main5.tex, subsec:exp-kuramoto (protocol in subsec:exp-protocol).
-#
-#   dynamics    : dx_i/dt = (1/N) Σ_j wbar_ij(t) sin(x_j - x_i),  wbar = N * W row-softmax
-#   initial data: k-twisted state  x(0, ξ) = 2πk(ξ - 1/2)  (+ jitter during training)
-#   baselines   : (i) mean-field  w ≡ 1  (W_ij = 1/N),
-#                 (ii) smooth ring bump of range r (mollified — the sharp ring
-#                      indicator has infinite receiver variation V_ξ and is NOT
-#                      in U_BV; see §5.2)
-#   claim       : twisted states are stationary for every symmetric convolution
-#                 kernel (odd integrand vs. even kernel), so both baselines stay
-#                 pinned at order parameter r ≈ 0; the trained control must break
-#                 the translational symmetry to untwist.
-#
-# Outputs (all under out_1/e1_<tag>/, paper-ready PDFs mirrored to
-# figures/e1_k{k}_T{T}[_sigma*]/). All three arms are always evaluated and
-# archived in trajectories.npz/metrics.json; the FIGURES show the arms in
-# PAPER_ARMS (mean-field vs trained control, per the paper's two-arm design).
-#   config.json, metrics.json, run_summary.txt, model.pt, history.json
-#   trajectories.npz     — x(t) per arm, W(t) for trained, r/Var/V_ξ series
-#                          (+ r_noisy_*/winding_noisy_* ensemble stats if σ > 0)
-#   r_panels.pdf         — r(t), one panel per paper arm
-#   r_combined.pdf       — same data, single axis
-#   variance.pdf         — Var(x(t)) per arm (log scale)
-#   winding.pdf          — winding number W(t) per arm: untwisting happens in
-#                          integer quanta (each quantum = 2π of real transport
-#                          for the outer agents; sin is blind to exact 2π gaps,
-#                          and r(t)=1 only certifies consensus mod 2π)
-#   vxi.pdf              — receiver variation V_ξ(ŵ(t)) of the trained kernel
-#   kernel_strip.pdf     — trained kernel heatmaps at 6 times, shared color scale
-#   state_evolution.pdf  — x(t, ξ) heatmaps per arm
-#   training_history.pdf — loss components over training
-#   untwist.gif          — agents as phases on the circle, all three arms animated
-#   kernel_time.gif      — trained kernel w(t, ξ, ζ) animated
-#
-# Usage:
-#   python e1_kuramoto.py                     # full run (k=1, N=64, 2000 steps)
-#   python e1_kuramoto.py --k 2 --T 12        # deeper twist, longer horizon
-#   python e1_kuramoto.py --k 3 --T 20 --steps 3000   # the paper's headline run
-#   python e1_kuramoto.py --k 2 --T 12 --sigma 0.3    # noisy robustness run
-#                                             # (adds 32-path noisy-ensemble eval)
-#   python e1_kuramoto.py --smoke             # 20-step smoke test
-#   python e1_kuramoto.py --device mps        # Apple GPU
+"""Compatibility entry point; implementation lives in graphon_control.experiments.kuramoto."""
+from graphon_control.experiments.kuramoto import *  # noqa: F401,F403
+from graphon_control.experiments.kuramoto import main
 
-import argparse
-import dataclasses
-import datetime
-import json
-import math
-import os
-from dataclasses import dataclass, field
-from typing import Callable, Dict, List
+_LEGACY_PLOT_NAMES = ('fig_r_panels', 'fig_r_combined', 'fig_winding', 'fig_variance', 'fig_vxi', 'fig_kernel_strip', 'fig_state_evolution', 'fig_training_history', 'save_untwist_gif', 'ARM_LABELS', 'ARM_COLORS', 'PAPER_ARMS')
 
-import numpy as np
-import torch
-import matplotlib
-
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-from tqdm import tqdm
-
-from model_1 import TimeLabelGraphonControl, _random_run_tag
-from model_1_viz import imshow_graphon, save_graphon_gif
-
-
-# ============================================================
-# Configuration (paper protocol, §5.1)
-# ============================================================
-
-@dataclass
-class E1Config:
-    # experiment
-    k: int = 1                 # winding number of the twisted state
-    jitter: float = 0.1        # i.i.d. jitter on the training initial data (resampled each batch)
-    ring_range: float = 0.1    # range r of the smooth ring baseline
-    x0_level: int = 0          # if > 0: use the coarse datum P_M(twist) with M = x0_level
-                               # (the common datum of lem:AN-monotone; requires M | N)
-
-    # discretization
-    N: int = 64
-    T: float = 5.0
-    dt: float = 0.05
-    sigma: float = 0.0         # deterministic dynamics (theory of §3 is deterministic)
-
-    # cost — PAPER convention (main5.tex, eqn: JN definition / eqn:J-cost-1):
-    #   J = Σ_n Δt [ Var(x^n) + β ||w(t_n)||²_{L²(I²)} ] + γ (mean(x^T) - x_target)²
-    # (β, not β/2, and NO terminal-variance term; lambda_terminal kept as an
-    #  optional knob but defaults to 0 to match the paper)
-    beta: float = 10.0
-    lambda_terminal: float = 0.0
-    gamma_target: float = 100.0
-    x_target: float = 0.0
-
-    # optimization
-    batch_size: int = 16
-    num_steps: int = 2000
-    lr: float = 1e-3
-    grad_clip: float = 1.0
-
-    # architecture (model_1.py defaults)
-    num_fourier_freqs: int = 8
-    embed_dim: int = 64
-    hidden_embed: int = 128
-    hidden_score: int = 128
-
-    # bookkeeping
-    seed: int = 0
-    device: str = "cpu"
-    time_chunk: int = 50       # chunk size for the batched W(t) precomputation
-
-    @property
-    def NT(self) -> int:
-        return int(round(self.T / self.dt))
-
-
-# ============================================================
-# E1 ingredients
-# ============================================================
-
-def phi_sin(y: torch.Tensor) -> torch.Tensor:
-    """Kuramoto interaction φ(z) = sin z (bounded, Lipschitz, odd; φ(0)=0)."""
-    return torch.sin(y)
-
-
-def twisted_state(xi: torch.Tensor, k: int) -> torch.Tensor:
-    """k-twisted state x(0, ξ) = 2πk(ξ - 1/2)."""
-    return 2.0 * math.pi * k * (xi - 0.5)
-
-
-def coarse_twisted_state(xi: torch.Tensor, k: int, M: int) -> torch.Tensor:
-    """
-    P_M projection of the k-twisted state: block average over M cells.
-    Since the twist is linear in ξ, the block average equals the twist at the
-    M cell midpoints, repeated N/M times — piecewise constant on the level-M
-    partition, hence exactly representable at every dyadic N ≥ M. This is the
-    common initial datum x_M(0) of lem:AN-monotone / thm:main (main5.tex).
-    """
-    N = xi.numel()
-    if N % M != 0:
-        raise ValueError(f"N={N} must be a multiple of M={M}")
-    mid = (torch.arange(M, device=xi.device, dtype=xi.dtype) + 0.5) / M
-    vals = twisted_state(mid, k)                     # (M,)
-    return vals.repeat_interleave(N // M)            # (N,)
-
-
-def mean_field_kernel(N: int, device, dtype=torch.float32) -> torch.Tensor:
-    """Row-stochastic representation of w ≡ 1: W_ij = 1/N."""
-    return torch.full((N, N), 1.0 / N, device=device, dtype=dtype)
-
-
-def smooth_ring_kernel(xi: torch.Tensor, r: float) -> torch.Tensor:
-    """
-    Mollified ring kernel: smooth bump of range r in torus distance,
-    row-normalized to a row-stochastic W. Smooth in ξ, hence of finite
-    receiver variation (unlike the sharp indicator (1/2r) 1_{[-r,r]}).
-    """
-    d = (xi.unsqueeze(0) - xi.unsqueeze(1)).abs()
-    d = torch.minimum(d, 1.0 - d)  # torus distance
-    u = (d / r).clamp(max=1.0 - 1e-6)
-    bump = torch.where(d < r, torch.exp(-1.0 / (1.0 - u ** 2)), torch.zeros_like(d))
-    return bump / bump.sum(dim=1, keepdim=True)
-
-
-class ConstantKernel:
-    """Baseline wrapper exposing the same call signature as the trained model."""
-
-    def __init__(self, W0: torch.Tensor):
-        self.W0 = W0
-
-    def __call__(self, t: torch.Tensor) -> torch.Tensor:
-        if t.ndim == 0:
-            return self.W0
-        return self.W0.unsqueeze(0).expand(t.shape[0], *self.W0.shape)
-
-
-# ============================================================
-# Fast rollout: W depends only on t, so precompute all W(t_n)
-# in a few batched forward passes, then run the Euler loop.
-# ============================================================
-
-def precompute_Ws(model, NT: int, dt: float, device, chunk: int) -> torch.Tensor:
-    """Returns (NT, N, N); differentiable w.r.t. model parameters."""
-    t_vals = torch.arange(NT, dtype=torch.float32, device=device) * dt
-    outs = [model(t_vals[a:a + chunk]) for a in range(0, NT, chunk)]
-    outs = [o.unsqueeze(0) if o.ndim == 2 else o for o in outs]  # chunk of 1 returns (N, N)
-    return torch.cat(outs, dim=0)
-
-
-def rollout_with_Ws(
-    Ws: torch.Tensor,
-    x0: torch.Tensor,
-    dt: float,
-    phi: Callable[[torch.Tensor], torch.Tensor],
-    sigma: float = 0.0,
-) -> torch.Tensor:
-    """
-    Euler(-Maruyama) rollout with precomputed kernels.
-    Ws: (NT, N, N) row-stochastic; x0: (B, N).
-    Sign convention (paper, eqn FN definition): diff[b, i, j] = x_j - x_i.
-    Returns xs: (NT + 1, B, N).
-    """
-    B, N = x0.shape
-    NT = Ws.shape[0]
-    sqrt_dt = math.sqrt(dt)
-    x = x0
-    xs = [x]
-    for n in range(NT):
-        diff = x.unsqueeze(1) - x.unsqueeze(2)          # (B, N, N): x_j - x_i
-        drift = (Ws[n].unsqueeze(0) * phi(diff)).sum(-1)  # (B, N)
-        x = x + dt * drift
-        if sigma > 0.0:
-            x = x + sigma * sqrt_dt * torch.randn_like(x)
-        xs.append(x)
-    return torch.stack(xs, dim=0)
-
-
-def e1_cost(xs: torch.Tensor, Ws: torch.Tensor, cfg: E1Config) -> Dict[str, torch.Tensor]:
-    """
-    Discretized cost J^{N,Δt}, PAPER convention (eqn: JN definition / eqn:J-cost-1):
-      Σ_n Δt [ Var(x^n) + β ||w(t_n)||²_{L²(I²)} ] + γ (mean(x^T) - x_target)²
-    where ||w||²_{L²(I²)} = Σ_ij W_ij² for the row-stochastic W (wbar = N·W).
-    lambda_terminal · Var(x^T) is an optional extra knob, 0 in the paper cost.
-    xs: (NT+1, B, N); Ws: (NT, N, N).
-    """
-    var_t = xs.var(dim=2, unbiased=False).mean(dim=1)        # (NT+1,) batch-averaged Var
-    running_state = cfg.dt * var_t[:-1].sum()
-    running_reg = cfg.dt * cfg.beta * (Ws ** 2).sum()
-    terminal_var = cfg.lambda_terminal * var_t[-1]
-    mean_T = xs[-1].mean(dim=1)                              # (B,)
-    terminal_target = cfg.gamma_target * ((mean_T - cfg.x_target) ** 2).mean()
-    total = running_state + running_reg + terminal_var + terminal_target
-    return {
-        "total": total,
-        "state_cost": running_state,
-        "l2_reg_cost": running_reg,
-        "terminal_var": var_t[-1],
-        "terminal_target": terminal_target,
-        "terminal_mean": mean_T.mean(),
-    }
-
-
-# ============================================================
-# Diagnostics
-# ============================================================
-
-def order_parameter(xs: torch.Tensor) -> torch.Tensor:
-    """
-    Kuramoto order parameter r(t) = | (1/N) Σ_j e^{i x_j(t)} |.
-    xs: (..., N) -> (...)
-    """
-    return torch.sqrt(torch.cos(xs).mean(-1) ** 2 + torch.sin(xs).mean(-1) ** 2)
-
-
-def winding_number(xs: torch.Tensor) -> torch.Tensor:
-    """
-    Winding number of the state along the label circle:
-      W(t) = (1/2π) [ Σ_i wrap(x_{i+1} - x_i) + wrap(x_1 - x_N) ],
-    with wrap(d) = ((d + π) mod 2π) - π. The k-twisted state has W = k;
-    consensus has W = 0. Untwisting proceeds in integer quanta: each unit of
-    winding requires transporting the outer agents a full 2π of real phase
-    (the sine coupling is blind to exact 2π gaps), reeled in through the
-    intermediate agents. xs: (..., N) -> (...).
-    """
-    gaps = xs[..., 1:] - xs[..., :-1]
-    close = (xs[..., :1] - xs[..., -1:])
-    all_gaps = torch.cat([gaps, close], dim=-1)
-    wrapped = torch.remainder(all_gaps + math.pi, 2.0 * math.pi) - math.pi
-    return wrapped.sum(dim=-1) / (2.0 * math.pi)
-
-
-def vxi_discrete(W: torch.Tensor) -> torch.Tensor:
-    """
-    Discrete receiver variation (paper eqn:receiver-tv-discrete) of the graphon
-    values wbar = N * W:
-      V_ξ = Σ_{i=1}^{N-1} sqrt( (1/N) Σ_j (wbar_{i+1,j} - wbar_{i,j})^2 ).
-    Accepts (N, N) or batched (..., N, N).
-    """
-    wbar = W.shape[-1] * W
-    d = wbar[..., 1:, :] - wbar[..., :-1, :]
-    return torch.sqrt((d ** 2).mean(dim=-1)).sum(dim=-1)
-
-
-# ============================================================
-# Figures
-# ============================================================
-
-ARM_LABELS = {"meanfield": r"mean-field  $w \equiv 1$",
-              "ring": "smooth ring",
-              "trained": "trained control"}
-ARM_COLORS = {"meanfield": "tab:gray", "ring": "tab:orange", "trained": "tab:blue"}
-
-
-PAPER_ARMS = ["meanfield", "trained"]   # arms shown in paper figures (ring is
-                                        # still evaluated and archived in npz)
-
-
-def fig_r_panels(t_grid, r_series: Dict[str, np.ndarray], k: int, path: str,
-                 arms: List[str] = None):
-    arms = arms or ["meanfield", "ring", "trained"]
-    fig, axes = plt.subplots(1, len(arms), figsize=(4.4 * len(arms), 3.6),
-                             sharey=True, squeeze=False)
-    for ax, arm in zip(axes[0], arms):
-        for other in arms:
-            if other != arm:
-                ax.plot(t_grid, r_series[other], color="0.85", lw=1.0)
-        ax.plot(t_grid, r_series[arm], color=ARM_COLORS[arm], lw=2.2)
-        ax.set_title(ARM_LABELS[arm])
-        ax.set_xlabel(r"$t$")
-        ax.set_ylim(-0.03, 1.03)
-        ax.grid(alpha=0.25)
-    axes[0][0].set_ylabel(r"order parameter $r(t)$")
-    fig.suptitle(rf"Untwisting the $k={k}$ twisted state", y=1.02)
-    fig.tight_layout()
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_r_combined(t_grid, r_series, k: int, path: str, arms: List[str] = None):
-    fig, ax = plt.subplots(figsize=(6.4, 4.0))
-    for arm in (arms or ["meanfield", "ring", "trained"]):
-        ax.plot(t_grid, r_series[arm], color=ARM_COLORS[arm], lw=2.0, label=ARM_LABELS[arm])
-    ax.set_xlabel(r"$t$")
-    ax.set_ylabel(r"$r(t)$")
-    ax.set_ylim(-0.03, 1.03)
-    ax.grid(alpha=0.25)
-    ax.legend(loc="center right")
-    ax.set_title(rf"Order parameter, $k={k}$ twisted state")
-    fig.tight_layout()
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_winding(t_grid, wind_series: Dict[str, np.ndarray], k: int, path: str,
-                arms: List[str] = None):
-    fig, ax = plt.subplots(figsize=(6.4, 4.0))
-    for arm in (arms or ["meanfield", "ring", "trained"]):
-        ax.plot(t_grid, wind_series[arm], color=ARM_COLORS[arm], lw=2.0,
-                label=ARM_LABELS[arm])
-    for q in range(k + 1):
-        ax.axhline(q, color="0.85", lw=0.8, zorder=0)
-    ax.set_xlabel(r"$t$")
-    ax.set_ylabel(r"winding number $\mathcal{W}(t)$")
-    ax.set_ylim(-0.25, k + 0.25)
-    ax.grid(alpha=0.15)
-    ax.legend()
-    ax.set_title("Untwisting proceeds in integer quanta")
-    fig.tight_layout()
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_variance(t_grid, var_series: Dict[str, np.ndarray], path: str,
-                 arms: List[str] = None):
-    fig, ax = plt.subplots(figsize=(6.4, 4.0))
-    for arm in (arms or ["meanfield", "ring", "trained"]):
-        ax.semilogy(t_grid, np.maximum(var_series[arm], 1e-12),
-                    color=ARM_COLORS[arm], lw=2.0, label=ARM_LABELS[arm])
-    ax.set_xlabel(r"$t$")
-    ax.set_ylabel(r"$\mathrm{Var}(x(t))$")
-    ax.grid(alpha=0.25, which="both")
-    ax.legend()
-    ax.set_title("Variance decay")
-    fig.tight_layout()
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_vxi(t_grid_W, vxi_trained: np.ndarray, vxi_ring: float, path: str):
-    fig, ax = plt.subplots(figsize=(6.4, 4.0))
-    ax.plot(t_grid_W, vxi_trained, color=ARM_COLORS["trained"], lw=2.0,
-            label=r"trained $V_\xi(\hat{w}(t))$")
-    ax.axhline(0.0, color=ARM_COLORS["meanfield"], lw=1.5, ls="--",
-               label=r"mean-field ($V_\xi = 0$)")
-    ax.axhline(vxi_ring, color=ARM_COLORS["ring"], lw=1.5, ls="--",
-               label=rf"smooth ring ($V_\xi = {vxi_ring:.1f}$)")
-    ax.set_xlabel(r"$t$")
-    ax.set_ylabel(r"$V_\xi$")
-    ax.grid(alpha=0.25)
-    ax.legend()
-    ax.set_title("Receiver variation of the kernel (measured, not enforced)")
-    fig.tight_layout()
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_kernel_strip(Ws: torch.Tensor, dt: float, path: str, n_panels: int = 6):
-    NT = Ws.shape[0]
-    idxs = np.linspace(0, NT - 1, n_panels).round().astype(int)
-    wbars = [Ws.shape[-1] * Ws[i].cpu().numpy() for i in idxs]
-    vmin = min(w.min() for w in wbars)
-    vmax = max(w.max() for w in wbars)
-    fig, axes = plt.subplots(1, n_panels, figsize=(2.6 * n_panels, 2.9))
-    for ax, i, w in zip(axes, idxs, wbars):
-        im = ax.imshow(w, origin="lower", extent=(0, 1, 0, 1), aspect="equal",
-                       cmap="viridis", interpolation="bicubic", vmin=vmin, vmax=vmax)
-        ax.set_title(rf"$t = {i * dt:.1f}$", fontsize=10)
-        ax.set_xticks([0, 1])
-        ax.set_yticks([0, 1])
-        if ax is axes[0]:
-            ax.set_ylabel(r"$\xi$")
-        ax.set_xlabel(r"$\zeta$")
-    fig.colorbar(im, ax=axes, fraction=0.02, pad=0.01, label=r"$\hat{w}(t,\xi,\zeta)$")
-    fig.suptitle("Trained kernel over time", y=1.04)
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_state_evolution(t_grid, xi, xs_arms: Dict[str, np.ndarray], k: int, path: str,
-                        arms: List[str] = None):
-    arms = arms or ["meanfield", "ring", "trained"]
-    lim = math.pi * max(k, 1) * 1.1
-    fig, axes = plt.subplots(1, len(arms), figsize=(4.4 * len(arms), 3.6),
-                             sharey=True, squeeze=False)
-    for ax, arm in zip(axes[0], arms):
-        X = xs_arms[arm]  # (NT+1, N)
-        im = ax.imshow(X.T, origin="lower", aspect="auto",
-                       extent=(t_grid[0], t_grid[-1], 0.0, 1.0),
-                       cmap="twilight_shifted", vmin=-lim, vmax=lim)
-        ax.set_title(ARM_LABELS[arm])
-        ax.set_xlabel(r"$t$")
-    axes = axes[0]
-    axes[0].set_ylabel(r"label $\xi$")
-    fig.colorbar(im, ax=axes, fraction=0.02, pad=0.01, label=r"$x(t,\xi)$")
-    fig.suptitle("State evolution", y=1.04)
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_training_history(history: Dict[str, List[float]], path: str):
-    fig, axes = plt.subplots(1, 2, figsize=(11, 3.8))
-    axes[0].plot(history["loss"], lw=1.0, color="tab:blue")
-    axes[0].set_title("total loss")
-    axes[0].set_xlabel("step")
-    axes[0].grid(alpha=0.25)
-    for key, color in [("state_cost", "tab:green"), ("l2_reg_cost", "tab:red"),
-                       ("terminal_var", "tab:purple"), ("terminal_target", "tab:brown")]:
-        axes[1].semilogy(np.maximum(np.array(history[key]), 1e-10), lw=1.0, label=key, color=color)
-    axes[1].set_title("components (log scale)")
-    axes[1].set_xlabel("step")
-    axes[1].grid(alpha=0.25, which="both")
-    axes[1].legend(fontsize=8)
-    fig.tight_layout()
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def save_untwist_gif(t_grid, xi, xs_arms: Dict[str, np.ndarray],
-                     r_series: Dict[str, np.ndarray], k: int, path: str,
-                     stride: int = 2, fps: int = 12, dpi: int = 80,
-                     arms: List[str] = None):
-    """
-    Animated circle view: each agent is a dot at angle x_i(t) on the unit circle,
-    colored by its label ξ. One panel per arm.
-    """
-    from PIL import Image
-    from io import BytesIO
-
-    arms = arms or ["meanfield", "ring", "trained"]
-    colors = plt.cm.hsv(xi)
-    frames = list(range(0, len(t_grid), stride))
-    if frames[-1] != len(t_grid) - 1:
-        frames.append(len(t_grid) - 1)
-
-    pil_imgs = []
-    theta = np.linspace(0, 2 * np.pi, 200)
-    for n in frames:
-        fig, axes = plt.subplots(1, len(arms), figsize=(3.5 * len(arms), 3.8),
-                                 squeeze=False)
-        for ax, arm in zip(axes[0], arms):
-            x_n = xs_arms[arm][n]  # (N,)
-            ax.plot(np.cos(theta), np.sin(theta), color="0.85", lw=1.0, zorder=1)
-            ax.scatter(np.cos(x_n), np.sin(x_n), c=colors, s=28,
-                       edgecolors="none", zorder=2)
-            ax.set_xlim(-1.25, 1.25)
-            ax.set_ylim(-1.25, 1.25)
-            ax.set_aspect("equal")
-            ax.axis("off")
-            ax.set_title(f"{ARM_LABELS[arm]}\n$r = {r_series[arm][n]:.3f}$", fontsize=10)
-        fig.suptitle(rf"$k={k}$ twist:  $t = {t_grid[n]:.2f}$", y=0.99)
-        fig.tight_layout()
-        buf = BytesIO()
-        fig.savefig(buf, format="png", dpi=dpi)
-        plt.close(fig)
-        buf.seek(0)
-        pil_imgs.append(Image.open(buf).copy())
-
-    pil_imgs[0].save(path, save_all=True, append_images=pil_imgs[1:], loop=0,
-                     duration=int(1000 / fps), optimize=False)
-
-
-# ============================================================
-# Main experiment
-# ============================================================
-
-def run_e1(cfg: E1Config, tag: str, out_root: str = "out_1", fig_root: str = "figures",
-           mirror_figs: bool = True):
-    torch.manual_seed(cfg.seed)
-    np.random.seed(cfg.seed)
-    device = torch.device(cfg.device)
-    dtype = torch.float32
-
-    run_dir = os.path.join(out_root, f"e1_{tag}")
-    # sigma > 0 runs mirror to their own figure dir so they never overwrite the
-    # noiseless run at the same (k, T); dot-free suffix keeps LaTeX paths safe
-    sigma_suffix = f"_sigma{cfg.sigma:g}".replace(".", "p") if cfg.sigma > 0 else ""
-    fig_dir = os.path.join(fig_root, f"e1_k{cfg.k}_T{cfg.T:g}{sigma_suffix}")
-    os.makedirs(run_dir, exist_ok=True)
-    os.makedirs(fig_dir, exist_ok=True)
-
-    xi = (torch.arange(cfg.N, device=device, dtype=dtype) + 0.5) / cfg.N
-    if cfg.x0_level > 0:
-        base = coarse_twisted_state(xi, cfg.k, cfg.x0_level)
-    else:
-        base = twisted_state(xi, cfg.k)
-    NT = cfg.NT
-
-    # Sanity assert: the twisted datum (smooth or coarse) must be exactly
-    # stationary for the mean-field baseline (roots-of-unity cancellation);
-    # if this fails, the "pinned baseline" claim of §5.2 is wrong for this grid.
-    with torch.no_grad():
-        drift0 = (mean_field_kernel(cfg.N, device) * phi_sin(
-            base.unsqueeze(0) - base.unsqueeze(1))).sum(-1)
-        pin_err = float(drift0.abs().max())
-    if pin_err > 1e-5:
-        raise RuntimeError(f"baseline pinning violated: |drift|={pin_err:.2e} "
-                           f"(N={cfg.N}, k={cfg.k}, x0_level={cfg.x0_level})")
-
-    print("=" * 64)
-    print(f"E1 Kuramoto untwisting | k={cfg.k} N={cfg.N} T={cfg.T} dt={cfg.dt}")
-    print(f"steps={cfg.num_steps} batch={cfg.batch_size} lr={cfg.lr} "
-          f"beta={cfg.beta} gamma={cfg.gamma_target} lambda={cfg.lambda_terminal}")
-    print(f"device={cfg.device} tag={tag}")
-    print("=" * 64)
-
-    # ---------- model ----------
-    model = TimeLabelGraphonControl(
-        xi=xi,
-        num_fourier_freqs=cfg.num_fourier_freqs,
-        embed_dim=cfg.embed_dim,
-        hidden_embed=cfg.hidden_embed,
-        hidden_score=cfg.hidden_score,
-        time_scale=max(cfg.T, 1e-8),
-        enforce_symmetry=False,
-    ).to(device)
-    optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr)
-
-    # ---------- training ----------
-    history: Dict[str, List[float]] = {kk: [] for kk in
-                                       ["loss", "state_cost", "l2_reg_cost",
-                                        "terminal_var", "terminal_target", "terminal_mean"]}
-    pbar = tqdm(range(cfg.num_steps), desc="E1 training", unit="step")
-    for _ in pbar:
-        model.train()
-        optimizer.zero_grad()
-        Ws = precompute_Ws(model, NT, cfg.dt, device, cfg.time_chunk)
-        x0 = base.unsqueeze(0) + cfg.jitter * torch.randn(cfg.batch_size, cfg.N,
-                                                          device=device, dtype=dtype)
-        xs = rollout_with_Ws(Ws, x0, cfg.dt, phi_sin, sigma=cfg.sigma)
-        parts = e1_cost(xs, Ws, cfg)
-        parts["total"].backward()
-        if cfg.grad_clip is not None:
-            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
-        optimizer.step()
-
-        for key in ["loss", "state_cost", "l2_reg_cost",
-                    "terminal_var", "terminal_target", "terminal_mean"]:
-            src = "total" if key == "loss" else key
-            history[key].append(float(parts[src].detach()))
-        pbar.set_postfix(loss=f"{history['loss'][-1]:.3f}",
-                         term_var=f"{history['terminal_var'][-1]:.4f}")
-
-    # ---------- evaluation: deterministic, exact twist, all three arms ----------
-    model.eval()
-    arms: Dict[str, object] = {
-        "meanfield": ConstantKernel(mean_field_kernel(cfg.N, device)),
-        "ring": ConstantKernel(smooth_ring_kernel(xi, cfg.ring_range)),
-        "trained": model,
-    }
-    x0_eval = base.unsqueeze(0)  # (1, N), no jitter: the exact stationary point
-
-    xs_arms_np: Dict[str, np.ndarray] = {}
-    Ws_arms: Dict[str, torch.Tensor] = {}
-    r_series: Dict[str, np.ndarray] = {}
-    var_series: Dict[str, np.ndarray] = {}
-    wind_series: Dict[str, np.ndarray] = {}
-    cost_summary: Dict[str, Dict[str, float]] = {}
-
-    with torch.no_grad():
-        for arm, kernel in arms.items():
-            Ws_eval = precompute_Ws(kernel, NT, cfg.dt, device, cfg.time_chunk)
-            xs_eval = rollout_with_Ws(Ws_eval, x0_eval, cfg.dt, phi_sin, sigma=0.0)
-            parts = e1_cost(xs_eval, Ws_eval, cfg)
-            X = xs_eval[:, 0, :].cpu().numpy()          # (NT+1, N)
-            xs_arms_np[arm] = X
-            Ws_arms[arm] = Ws_eval
-            r_series[arm] = order_parameter(xs_eval[:, 0, :]).cpu().numpy()
-            var_series[arm] = xs_eval[:, 0, :].var(dim=1, unbiased=False).cpu().numpy()
-            wind_series[arm] = winding_number(xs_eval[:, 0, :]).cpu().numpy()
-            cost_summary[arm] = {kk: float(vv) for kk, vv in parts.items()}
-
-    vxi_trained = vxi_discrete(Ws_arms["trained"]).cpu().numpy()   # (NT,)
-    vxi_ring = float(vxi_discrete(Ws_arms["ring"][0]))
-    vxi_mf = float(vxi_discrete(Ws_arms["meanfield"][0]))
-
-    t_grid = np.arange(NT + 1) * cfg.dt
-    t_grid_W = np.arange(NT) * cfg.dt
-
-    # ---------- noisy ensemble evaluation (robustness runs, sigma > 0) ----------
-    # The deterministic evaluation above supplies the paper's table values (the
-    # noisy training loss carries a noise floor ~ sigma^2 T^2 / 2); here we
-    # additionally roll the trained control on the noisy dynamics it was trained
-    # for, over an ensemble of Brownian paths from the exact twist.
-    noisy_eval: Dict[str, np.ndarray] = {}
-    if cfg.sigma > 0.0:
-        n_rep = 32
-        with torch.no_grad():
-            xs_noisy = rollout_with_Ws(Ws_arms["trained"], x0_eval.expand(n_rep, -1),
-                                       cfg.dt, phi_sin, sigma=cfg.sigma)  # (NT+1, n_rep, N)
-            r_rep = order_parameter(xs_noisy).cpu().numpy()               # (NT+1, n_rep)
-            w_rep = winding_number(xs_noisy).cpu().numpy()
-        noisy_eval = {
-            "r_noisy_mean": r_rep.mean(axis=1),
-            "r_noisy_q10": np.quantile(r_rep, 0.10, axis=1),
-            "r_noisy_q90": np.quantile(r_rep, 0.90, axis=1),
-            "winding_noisy_mean": w_rep.mean(axis=1),
-            "winding_noisy_q10": np.quantile(w_rep, 0.10, axis=1),
-            "winding_noisy_q90": np.quantile(w_rep, 0.90, axis=1),
-        }
-
-    # ---------- persistence ----------
-    torch.save(model.state_dict(), os.path.join(run_dir, "model.pt"))
-    with open(os.path.join(run_dir, "config.json"), "w") as f:
-        json.dump(dataclasses.asdict(cfg), f, indent=2)
-    with open(os.path.join(run_dir, "history.json"), "w") as f:
-        json.dump(history, f, indent=2)
-
-    np.savez_compressed(
-        os.path.join(run_dir, "trajectories.npz"),
-        t=t_grid, t_W=t_grid_W, xi=xi.cpu().numpy(),
-        x_meanfield=xs_arms_np["meanfield"], x_ring=xs_arms_np["ring"],
-        x_trained=xs_arms_np["trained"],
-        W_trained=Ws_arms["trained"].cpu().numpy().astype(np.float32),
-        W_ring=Ws_arms["ring"][0].cpu().numpy().astype(np.float32),
-        r_meanfield=r_series["meanfield"], r_ring=r_series["ring"],
-        r_trained=r_series["trained"],
-        var_meanfield=var_series["meanfield"], var_ring=var_series["ring"],
-        var_trained=var_series["trained"],
-        winding_meanfield=wind_series["meanfield"], winding_ring=wind_series["ring"],
-        winding_trained=wind_series["trained"],
-        vxi_trained=vxi_trained,
-        **noisy_eval,
-    )
-
-    metrics = {
-        "k": cfg.k,
-        "costs": cost_summary,
-        "r_initial": float(r_series["trained"][0]),
-        "r_final": {arm: float(r_series[arm][-1]) for arm in arms},
-        "var_final": {arm: float(var_series[arm][-1]) for arm in arms},
-        "winding_final": {arm: float(wind_series[arm][-1]) for arm in arms},
-        "terminal_mean_trained": cost_summary["trained"]["terminal_mean"],
-        "vxi": {"trained_max": float(vxi_trained.max()),
-                "trained_final": float(vxi_trained[-1]),
-                "ring": vxi_ring, "meanfield": vxi_mf},
-        "value_gap_meanfield_minus_trained":
-            cost_summary["meanfield"]["total"] - cost_summary["trained"]["total"],
-    }
-    if noisy_eval:
-        metrics["noisy_eval"] = {
-            "sigma": cfg.sigma,
-            "n_replicates": 32,
-            "r_final_mean": float(noisy_eval["r_noisy_mean"][-1]),
-            "r_final_q10": float(noisy_eval["r_noisy_q10"][-1]),
-            "r_final_q90": float(noisy_eval["r_noisy_q90"][-1]),
-            "winding_final_mean": float(noisy_eval["winding_noisy_mean"][-1]),
-        }
-    with open(os.path.join(run_dir, "metrics.json"), "w") as f:
-        json.dump(metrics, f, indent=2)
-
-    # ---------- figures ----------
-    def both(name: str) -> List[str]:
-        if not mirror_figs:
-            return [os.path.join(run_dir, name)]
-        return [os.path.join(run_dir, name), os.path.join(fig_dir, name)]
-
-    for p in both("r_panels.pdf"):
-        fig_r_panels(t_grid, r_series, cfg.k, p, arms=PAPER_ARMS)
-    for p in both("r_combined.pdf"):
-        fig_r_combined(t_grid, r_series, cfg.k, p, arms=PAPER_ARMS)
-    for p in both("variance.pdf"):
-        fig_variance(t_grid, var_series, p, arms=PAPER_ARMS)
-    for p in both("winding.pdf"):
-        fig_winding(t_grid, wind_series, cfg.k, p, arms=PAPER_ARMS)
-    for p in both("vxi.pdf"):
-        fig_vxi(t_grid_W, vxi_trained, vxi_ring, p)
-    for p in both("kernel_strip.pdf"):
-        fig_kernel_strip(Ws_arms["trained"], cfg.dt, p)
-    for p in both("state_evolution.pdf"):
-        fig_state_evolution(t_grid, xi.cpu().numpy(), xs_arms_np, cfg.k, p,
-                            arms=PAPER_ARMS)
-    fig_training_history(history, os.path.join(run_dir, "training_history.pdf"))
-
-    # ---------- GIFs ----------
-    save_untwist_gif(t_grid, xi.cpu().numpy(), xs_arms_np, r_series, cfg.k,
-                     os.path.join(run_dir, "untwist.gif"), arms=PAPER_ARMS)
-    n_gif = min(60, NT)
-    gif_idx = np.linspace(0, NT - 1, n_gif).round().astype(int)
-    gif_frames = [Ws_arms["trained"][i].cpu() * cfg.N for i in gif_idx]
-    save_graphon_gif(gif_frames, os.path.join(run_dir, "kernel_time.gif"),
-                     fps=10, t_vals=[float(i * cfg.dt) for i in gif_idx])
-
-    # ---------- summary ----------
-    lines = [
-        "=== E1 Kuramoto untwisting: run summary ===",
-        f"timestamp : {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
-        f"run_dir   : {run_dir}",
-        f"k = {cfg.k}, N = {cfg.N}, T = {cfg.T}, dt = {cfg.dt}, steps = {cfg.num_steps}",
-        "",
-        "arm         J_total      r(T)      Var(T)    winding(T)",
-    ]
-    for arm in ["meanfield", "ring", "trained"]:
-        lines.append(f"{arm:<11s} {cost_summary[arm]['total']:>9.4f}  "
-                     f"{metrics['r_final'][arm]:>8.4f}  {metrics['var_final'][arm]:>10.6f}  "
-                     f"{metrics['winding_final'][arm]:>+9.3f}")
-    lines += [
-        "",
-        f"V_xi trained: max {metrics['vxi']['trained_max']:.3f}, "
-        f"final {metrics['vxi']['trained_final']:.3f} "
-        f"(ring {vxi_ring:.3f}, mean-field {vxi_mf:.3f})",
-        f"value gap (mean-field - trained): "
-        f"{metrics['value_gap_meanfield_minus_trained']:.4f}",
-    ]
-    summary = "\n".join(lines)
-    with open(os.path.join(run_dir, "run_summary.txt"), "w") as f:
-        f.write(summary + "\n")
-    print(summary)
-    print(f"\nAll outputs in: {run_dir}\nPaper figures mirrored to: {fig_dir}")
-    return metrics
-
-
-# ============================================================
-# CLI
-# ============================================================
+def __getattr__(name):
+    if name in _LEGACY_PLOT_NAMES:
+        from importlib import import_module
+        return getattr(import_module("graphon_control.experiments.kuramoto_plots"), name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="E1: Kuramoto twisted-state untwisting")
-    parser.add_argument("--k", type=int, default=1, help="winding number of the twist")
-    parser.add_argument("--T", type=float, default=5.0, help="time horizon")
-    parser.add_argument("--N", type=int, default=64)
-    parser.add_argument("--x0-level", type=int, default=0, dest="x0_level",
-                        help="if > 0, use the coarse datum P_M(twist) with M = this level")
-    parser.add_argument("--no-mirror", action="store_true",
-                        help="do not mirror figures into figures/ (sweep runs)")
-    parser.add_argument("--steps", type=int, default=2000)
-    parser.add_argument("--batch", type=int, default=16)
-    parser.add_argument("--sigma", type=float, default=0.0,
-                        help="Euler-Maruyama noise level (training dynamics and "
-                             "the noisy ensemble evaluation)")
-    parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument("--device", type=str, default="cpu", choices=["cpu", "mps", "cuda"])
-    parser.add_argument("--tag", type=str, default="")
-    parser.add_argument("--smoke", action="store_true", help="20-step smoke test")
-    args = parser.parse_args()
-
-    cfg = E1Config(k=args.k, T=args.T, N=args.N, num_steps=args.steps,
-                   batch_size=args.batch, sigma=args.sigma, seed=args.seed,
-                   device=args.device, x0_level=args.x0_level)
-    if args.smoke:
-        cfg.num_steps = 20
-
-    tag = args.tag or f"k{args.k}_{_random_run_tag()}"
-    if args.smoke:
-        tag = f"smoke_{tag}"
-    run_e1(cfg, tag=tag, mirror_figs=not args.no_mirror)
+    main()
diff --git a/e1_msweep.py b/e1_msweep.py
index 0d7f0bd..cb21769 100644
--- a/e1_msweep.py
+++ b/e1_msweep.py
@@ -1,180 +1,255 @@
-# e1_msweep.py
-#
-# M-sweep of experiment E1: the numerical portrait of thm:main's OUTER limit
-# and of cor:value-stability (main5.tex, §5.3).
-#
-# Fix N = 64 and sweep the DATUM resolution: initial datum P_M(k=1 twist) for
-# M in {2, 4, 8, 16, 32, 64}, three seeds each. cor:value-stability predicts
-#   |A_inf(x_M(0)) - A_inf(x(0))| <= C ||P_M x0 - x0||_{L2} = C * 2pi/(sqrt(12) M),
-# a genuine 1/M convergence curve. Closed-form anchor at M = 2: the two blocks
-# sit at +-pi/2, every pairwise difference lies in {0, pi} where sin vanishes,
-# so the drift is identically zero FOR EVERY admissible kernel: the datum is
-# uncontrollable and
-#   A(M=2) = Var * T + beta * ||w=1||^2 * T = (pi^2/4) * 5 + 10 * 5 = 62.337.
-# (The optimizer should recover this by driving the kernel to uniform.)
-#
-# Reused cells (identical problems already trained):
-#   M=8  s0  <- out_1/e1_sweep_N64_s0   (N-sweep run, x0_level=8)
-#   M=8  s1  <- out_1/e1_sweep_N64_s1
-#   M=64 s0  <- out_1/e1_k1_T5          (P_64 twist == smooth twist on the N=64 grid)
-#
-# Usage:
-#   .venv312/bin/python e1_msweep.py                 # run + aggregate
-#   .venv312/bin/python e1_msweep.py --aggregate-only
+"""Run and summarize the Kuramoto initial-datum resolution study.
+
+Each cell fixes N=64, k=1, T=5 and changes the block-averaged initial datum P_M x0.
+The plotted values are optimizer outputs, not certified value functions or a
+proof of the manuscript's continuum limit. At M=2 the sine drift vanishes for
+all kernels, giving the exact uniform-kernel anchor recorded in the summary.
+
+Existing cells are reused only after validating config, provenance, and metrics.
+Incomplete cells are never overwritten: select a new --sweep-tag to retry.
+"""
 
 import argparse
+import hashlib
 import json
 import math
-import os
-import shutil
+from pathlib import Path
 import subprocess
 import sys
 import time
 
 import numpy as np
 
+from graphon_control.artifacts import (
+    SCHEMA_VERSION, atomic_write_json, create_run_directory, generated_run_tag,
+)
+
+
 MS = [64, 32, 16, 8, 4, 2]
 SEEDS = [0, 1, 2]
 N = 64
 T = 5.0
 BETA = 10.0
 LANES = 3
-PY = sys.executable
-
-REUSE = {
-    (8, 0): "out_1/e1_sweep_N64_s0",
-    (8, 1): "out_1/e1_sweep_N64_s1",
-    (64, 0): "out_1/e1_k1_T5",
-}
-
-J_M2_EXACT = (math.pi ** 2 / 4.0) * T + BETA * 1.0 * T  # uncontrollable anchor
-
-
-def cell_dir(M, s):
-    return os.path.join("out_1", f"e1_msweep_M{M}_s{s}")
-
-
-def done(M, s):
-    return os.path.exists(os.path.join(cell_dir(M, s), "metrics.json"))
-
-
-def stage_reuse():
-    for (M, s), src in REUSE.items():
-        if done(M, s):
-            continue
-        if os.path.exists(os.path.join(src, "metrics.json")):
-            shutil.copytree(src, cell_dir(M, s))
-            print(f"reused M={M} s={s} <- {src}")
-
-
-def launch_all(steps: int):
-    jobs = [(M, s) for s in SEEDS for M in MS if not done(M, s)]
-    print(f"{len(jobs)} runs to do")
+REPO_ROOT = Path(__file__).resolve().parent
+DEFAULT_OUT_ROOT = REPO_ROOT / "out_1"
+J_M2_EXACT = (math.pi ** 2 / 4.0) * T + BETA * T
+
+
+def cell_tag(M, seed, sweep_tag="msweep"):
+    letters_and_digits = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+    if (not sweep_tag or sweep_tag[0] not in letters_and_digits
+            or any(c not in letters_and_digits + "_.-" for c in sweep_tag)):
+        raise ValueError("sweep_tag must be a nonempty safe path component")
+    return f"{sweep_tag}_M{M}_s{seed}"
+
+
+def cell_dir(M, seed, out_root=DEFAULT_OUT_ROOT, sweep_tag="msweep"):
+    return Path(out_root) / f"e1_{cell_tag(M, seed, sweep_tag)}"
+
+
+def _expected_config(M, seed, steps):
+    return {"N": N, "k": 1, "T": T, "dt": 0.05, "sigma": 0.0,
+            "beta": BETA, "x0_level": M, "seed": seed, "num_steps": steps}
+
+
+def read_cell(M, seed, *, steps=2000, out_root=DEFAULT_OUT_ROOT, sweep_tag="msweep"):
+    """Return verified metrics, None if absent, or fail on an incompatible cell."""
+    directory = cell_dir(M, seed, out_root, sweep_tag)
+    if not directory.exists():
+        return None
+    if not directory.is_dir():
+        raise ValueError(f"Cell path is not a directory: {directory}")
+    try:
+        with (directory / "config.json").open(encoding="utf-8") as stream:
+            config = json.load(stream)
+        with (directory / "provenance.json").open(encoding="utf-8") as stream:
+            provenance = json.load(stream)
+        with (directory / "metrics.json").open(encoding="utf-8") as stream:
+            metrics = json.load(stream)
+        with (directory / "complete.json").open(encoding="utf-8") as stream:
+            completion = json.load(stream)
+        if not isinstance(completion, dict) or completion.get("status") != "complete":
+            raise ValueError("run has no successful completion marker")
+        if not (directory / "model.pt").is_file():
+            raise ValueError("missing model.pt")
+        for key, value in _expected_config(M, seed, steps).items():
+            if config.get(key) != value:
+                raise ValueError(f"config {key}: expected {value!r}, got {config.get(key)!r}")
+        encoded = json.dumps(config, sort_keys=True, allow_nan=False).encode("utf-8")
+        if provenance.get("schema_version") != SCHEMA_VERSION:
+            raise ValueError("missing or unsupported provenance schema")
+        if provenance.get("training_seed") != seed:
+            raise ValueError("provenance training seed mismatch")
+        if provenance.get("config_sha256") != hashlib.sha256(encoded).hexdigest():
+            raise ValueError("configuration does not match saved provenance")
+        values = [metrics["costs"]["trained"]["total"], metrics["costs"]["meanfield"]["total"],
+                  metrics["r_final"]["trained"], metrics["var_final"]["trained"],
+                  metrics["vxi"]["trained_max"]]
+        if not all(np.isfinite(v) for v in values):
+            raise ValueError("nonfinite metrics")
+    except (OSError, ValueError, KeyError, TypeError) as exc:
+        raise ValueError(
+            f"Incomplete or incompatible sweep cell {directory}: {exc}. "
+            "Use a new --sweep-tag instead of overwriting it."
+        ) from exc
+    return metrics
+
+
+def done(M, seed, *, steps=2000, out_root=DEFAULT_OUT_ROOT, sweep_tag="msweep"):
+    return read_cell(M, seed, steps=steps, out_root=out_root, sweep_tag=sweep_tag) is not None
+
+
+def launch_all(steps, *, out_root=DEFAULT_OUT_ROOT, lanes=LANES, sweep_tag="msweep", device="cpu"):
+    if steps < 1 or lanes < 1:
+        raise ValueError("steps and lanes must be positive")
+    root = Path(out_root).expanduser().resolve()
+    # Preflight existing cells before starting any subprocess.
+    jobs = [(M, seed) for seed in SEEDS for M in MS
+            if not done(M, seed, steps=steps, out_root=root, sweep_tag=sweep_tag)]
+    root.mkdir(parents=True, exist_ok=True)
+    print(f"{len(jobs)} runs to do", flush=True)
     active = []
-    while jobs or active:
-        while jobs and len(active) < LANES:
-            M, s = jobs.pop(0)
-            log = open(os.path.join("out_1", f"e1_msweep_M{M}_s{s}.log"), "w")
-            p = subprocess.Popen(
-                [PY, "e1_kuramoto.py", "--k", "1", "--T", str(T), "--N", str(N),
-                 "--seed", str(s), "--x0-level", str(M), "--steps", str(steps),
-                 "--tag", f"msweep_M{M}_s{s}", "--no-mirror"],
-                stdout=log, stderr=subprocess.STDOUT)
-            active.append((p, M, s, log))
-            print(f"launched M={M} seed={s} (pid {p.pid})", flush=True)
-        for t in list(active):
-            if t[0].poll() is not None:
-                active.remove(t)
-                t[3].close()
-                status = "ok" if t[0].returncode == 0 else f"FAILED rc={t[0].returncode}"
-                print(f"finished M={t[1]} seed={t[2]}: {status}", flush=True)
-        time.sleep(10)
-
-
-def aggregate():
+    try:
+        while jobs or active:
+            while jobs and len(active) < lanes:
+                M, seed = jobs.pop(0)
+                log_path = root / f"e1_{cell_tag(M, seed, sweep_tag)}.log"
+                log = log_path.open("x", encoding="utf-8")
+                command = [sys.executable, str(REPO_ROOT / "e1_kuramoto.py"),
+                           "--k", "1", "--T", str(T), "--N", str(N), "--seed", str(seed),
+                           "--x0-level", str(M), "--steps", str(steps), "--device", device,
+                           "--out-root", str(root), "--tag", cell_tag(M, seed, sweep_tag),
+                           "--no-mirror"]
+                try:
+                    process = subprocess.Popen(command, cwd=REPO_ROOT, stdout=log, stderr=subprocess.STDOUT)
+                except BaseException:
+                    log.close()
+                    raise
+                active.append((process, M, seed, log, log_path))
+                print(f"launched M={M} seed={seed} (pid {process.pid})", flush=True)
+            for job in list(active):
+                process, M, seed, log, log_path = job
+                if process.poll() is None:
+                    continue
+                active.remove(job)
+                log.close()
+                if process.returncode:
+                    raise RuntimeError(f"M={M} seed={seed} failed ({process.returncode}); see {log_path}")
+                if not done(M, seed, steps=steps, out_root=root, sweep_tag=sweep_tag):
+                    raise RuntimeError(f"M={M} seed={seed} exited without a complete run; see {log_path}")
+                print(f"finished M={M} seed={seed}", flush=True)
+            if active:
+                time.sleep(1)
+    finally:
+        for process, _, _, log, _ in active:
+            if process.poll() is None:
+                process.terminate()
+            try:
+                process.wait(timeout=10)
+            except subprocess.TimeoutExpired:
+                process.kill()
+                process.wait()
+            log.close()
+
+
+def aggregate(*, steps=2000, out_root=DEFAULT_OUT_ROOT, sweep_tag="msweep", allow_partial=False):
+    rows, missing = [], []
+    for M in sorted(MS):
+        for seed in SEEDS:
+            metrics = read_cell(M, seed, steps=steps, out_root=out_root, sweep_tag=sweep_tag)
+            if metrics is None:
+                missing.append({"M": M, "seed": seed})
+                continue
+            rows.append({"M": M, "seed": seed,
+                         "J_trained": metrics["costs"]["trained"]["total"],
+                         "J_meanfield": metrics["costs"]["meanfield"]["total"],
+                         "r_final": metrics["r_final"]["trained"],
+                         "var_final": metrics["var_final"]["trained"],
+                         "vxi_max": metrics["vxi"]["trained_max"]})
+    if not rows:
+        raise ValueError("No complete sweep cells are available; no summary was written")
+    if missing and not allow_partial:
+        raise ValueError(f"Sweep is incomplete ({len(missing)} missing cells); use --allow-partial to summarize explicitly")
+
     import matplotlib
     matplotlib.use("Agg")
     import matplotlib.pyplot as plt
 
-    rows = []
-    for M in sorted(MS):
-        for s in SEEDS:
-            d = cell_dir(M, s)
-            if not os.path.exists(os.path.join(d, "metrics.json")):
-                print(f"missing: M={M} seed={s}")
-                continue
-            with open(os.path.join(d, "metrics.json")) as f:
-                m = json.load(f)
-            rows.append({"M": M, "seed": s,
-                         "J_trained": m["costs"]["trained"]["total"],
-                         "J_meanfield": m["costs"]["meanfield"]["total"],
-                         "r_final": m["r_final"]["trained"],
-                         "var_final": m["var_final"]["trained"],
-                         "vxi_max": m["vxi"]["trained_max"]})
-    os.makedirs("out_1/e1_msweep", exist_ok=True)
-    with open("out_1/e1_msweep/msweep_summary.json", "w") as f:
-        json.dump({"rows": rows, "J_M2_exact": J_M2_EXACT}, f, indent=2)
-
-    Ms = sorted({r["M"] for r in rows})
-    jmin = {M: min(r["J_trained"] for r in rows if r["M"] == M) for M in Ms}
-    jlo = {M: min(r["J_trained"] for r in rows if r["M"] == M) for M in Ms}
-    jhi = {M: max(r["J_trained"] for r in rows if r["M"] == M) for M in Ms}
-    jmf = {M: np.median([r["J_meanfield"] for r in rows if r["M"] == M]) for M in Ms}
-
+    levels = sorted({r["M"] for r in rows})
+    jmin = {M: min(r["J_trained"] for r in rows if r["M"] == M) for M in levels}
+    jhi = {M: max(r["J_trained"] for r in rows if r["M"] == M) for M in levels}
+    jmf = {M: np.median([r["J_meanfield"] for r in rows if r["M"] == M]) for M in levels}
     fig, axes = plt.subplots(1, 2, figsize=(11, 4.0))
     ax = axes[0]
-    ax.fill_between(Ms, [jlo[M] for M in Ms], [jhi[M] for M in Ms],
-                    color="tab:blue", alpha=0.18, label="seed range")
-    ax.plot(Ms, [jmin[M] for M in Ms], "o-", color="tab:blue", lw=2,
-            label=r"trained $\widehat{J}(M)$ (min over seeds)")
-    ax.plot(Ms, [jmf[M] for M in Ms], "s--", color="tab:gray", lw=1.5,
-            label=r"mean-field baseline")
-    ax.plot([2], [J_M2_EXACT], marker="*", ms=16, color="tab:red", ls="none",
-            label=r"$M=2$: exact (uncontrollable)")
+    ax.fill_between(levels, [jmin[M] for M in levels], [jhi[M] for M in levels],
+                    color="tab:blue", alpha=0.18, label="available seed range")
+    ax.plot(levels, [jmin[M] for M in levels], "o-", color="tab:blue", lw=2,
+            label=r"trained $\widehat{J}(M)$ (min over available seeds)")
+    ax.plot(levels, [jmf[M] for M in levels], "s--", color="tab:gray", label="mean-field baseline")
+    ax.plot([2], [J_M2_EXACT], "*", ms=16, color="tab:red", label=r"$M=2$: exact")
     ax.set_xscale("log", base=2)
-    ax.set_xticks(Ms); ax.set_xticklabels([str(m) for m in Ms])
-    ax.set_xlabel(r"datum resolution $M$"); ax.set_ylabel(r"$\widehat{J}$")
-    ax.set_title(r"Values vs.\ datum resolution ($N = 64$)")
-    ax.grid(alpha=0.25); ax.legend(fontsize=9)
-
+    ax.set_xticks(levels, [str(M) for M in levels])
+    ax.set(xlabel=r"datum resolution $M$", ylabel=r"$\widehat{J}$",
+           title=f"Optimizer values vs. datum resolution (N={N})")
+    ax.grid(alpha=0.25)
+    ax.legend(fontsize=8)
     ax = axes[1]
-    Mref = max(Ms)
-    Mplot = [M for M in Ms if M != Mref]
-    gaps = [abs(jmin[M] - jmin[Mref]) for M in Mplot]
-    ax.loglog(Mplot, gaps, "o-", color="tab:blue", lw=2,
-              label=r"$|\widehat{J}(M) - \widehat{J}(%d)|$" % Mref)
-    guide1 = [gaps[1] * Mplot[1] / M for M in Mplot]
-    guide2 = [gaps[1] * (Mplot[1] / M) ** 2 for M in Mplot]
-    ax.loglog(Mplot, guide1, "--", color="0.6",
-              label=r"$1/M$ (bound of \Cref{cor:value-stability})"
-              if False else r"$1/M$ (corollary bound)")
-    ax.loglog(Mplot, guide2, ":", color="0.4", label=r"$1/M^2$ (observed)")
-    ax.set_xticks(Mplot); ax.set_xticklabels([str(m) for m in Mplot])
-    ax.set_xlabel(r"$M$"); ax.set_ylabel("value gap")
-    ax.set_title(r"Convergence in the datum resolution")
-    ax.grid(alpha=0.25, which="both"); ax.legend(fontsize=9)
+    reference = max(levels)
+    plotted = [M for M in levels if M != reference]
+    gaps = [abs(jmin[M] - jmin[reference]) for M in plotted]
+    if plotted:
+        ax.loglog(plotted, gaps, "o-", color="tab:blue", lw=2,
+                  label=rf"$|\widehat{{J}}(M)-\widehat{{J}}({reference})|$")
+        positive = [(M, gap) for M, gap in zip(plotted, gaps) if gap > 0]
+        if positive:
+            anchor_M, anchor_gap = positive[0]
+            ax.loglog(plotted, [anchor_gap * anchor_M / M for M in plotted], "--",
+                      color="0.6", label=r"$1/M$ reference slope")
+        else:
+            ax.text(0.5, 0.5, "All observed gaps are zero", transform=ax.transAxes, ha="center")
+        ax.set_xticks(plotted, [str(M) for M in plotted])
+        ax.legend(fontsize=8)
+    else:
+        ax.text(0.5, 0.5, "Need at least two resolutions for gaps", transform=ax.transAxes, ha="center")
+    ax.set(xlabel=r"$M$", ylabel="optimizer value gap", title="Exploratory resolution comparison")
+    ax.grid(alpha=0.25, which="both")
+    if missing:
+        fig.suptitle(f"PARTIAL SWEEP: {len(rows)} of {len(MS) * len(SEEDS)} cells")
     fig.tight_layout()
-    for outp in ["out_1/e1_msweep/convergence_M.pdf",
-                 "figures/e1_msweep/convergence_M.pdf"]:
-        os.makedirs(os.path.dirname(outp), exist_ok=True)
-        fig.savefig(outp, bbox_inches="tight")
-    plt.close(fig)
-
-    print(f"aggregated {len(rows)} cells -> convergence_M.pdf")
-    for M in Ms:
-        rs = [r for r in rows if r["M"] == M]
-        print(f"  M={M:3d}  J min {jmin[M]:8.4f}  range [{jlo[M]:.4f}, {jhi[M]:.4f}]  "
-              f"baseline {jmf[M]:8.4f}  r(T) max {max(r['r_final'] for r in rs):.4f}")
-    print(f"  M=2 exact anchor: {J_M2_EXACT:.4f}")
-    print("MSWEEP_COMPLETE")
+    output = create_run_directory(out_root, "e1", f"{sweep_tag}_summary_{generated_run_tag()}")
+    try:
+        fig.savefig(output / "convergence_M.pdf", bbox_inches="tight")
+    finally:
+        plt.close(fig)
+    atomic_write_json(output / "msweep_summary.json", {
+        "rows": rows, "J_M2_exact": J_M2_EXACT, "complete": not missing,
+        "missing": missing, "steps": steps, "sweep_tag": sweep_tag,
+        "interpretation": "Optimizer upper-bound estimates; no certified convergence rate.",
+    })
+    print(f"aggregated {len(rows)} cells -> {output}")
+    return output
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--steps", type=int, default=2000)
+    parser.add_argument("--aggregate-only", action="store_true")
+    parser.add_argument("--allow-partial", action="store_true")
+    parser.add_argument("--out-root", type=Path, default=DEFAULT_OUT_ROOT)
+    parser.add_argument("--sweep-tag", default="msweep")
+    parser.add_argument("--lanes", type=int, default=LANES)
+    parser.add_argument("--device", choices=["cpu", "mps", "cuda"], default="cpu")
+    args = parser.parse_args(argv)
+    try:
+        if not args.aggregate_only:
+            launch_all(args.steps, out_root=args.out_root, lanes=args.lanes,
+                       sweep_tag=args.sweep_tag, device=args.device)
+        aggregate(steps=args.steps, out_root=args.out_root, sweep_tag=args.sweep_tag,
+                  allow_partial=args.allow_partial)
+    except (OSError, ValueError, RuntimeError) as exc:
+        parser.exit(1, f"{exc}\n")
 
 
 if __name__ == "__main__":
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--steps", type=int, default=2000)
-    ap.add_argument("--aggregate-only", action="store_true")
-    args = ap.parse_args()
-    if not args.aggregate_only:
-        stage_reuse()
-        launch_all(steps=args.steps)
-    aggregate()
+    main()
diff --git a/e1_noise_figure.py b/e1_noise_figure.py
index cd04b87..36d7565 100644
--- a/e1_noise_figure.py
+++ b/e1_noise_figure.py
@@ -1,46 +1,112 @@
-# e1_noise_figure.py
-#
-# Robustness panel for the noisy Kuramoto runs (main5.tex, subsec:exp-kuramoto):
-# trained control evaluated on deterministic dynamics (dashed) against the
-# noisy-ensemble statistics saved by e1_kuramoto.py for sigma > 0 runs
-# (mean over 32 Brownian replicates, shaded 10-90% band), for both the order
-# parameter r(t) and the winding number W(t).
-#
-# Usage:
-#   python e1_noise_figure.py --noisy k2_T12_sig01 [--noiseless k2_T12] [--out ...]
+"""Compare compatible, separately trained clean and noisy Kuramoto runs.
+
+Use --noisy and --noiseless with run directories, or with legacy run tags under
+--out-root. No shared manuscript figure directory is inferred or modified.
+"""
 
 import argparse
 import json
-import os
+from pathlib import Path
 
 import numpy as np
-import matplotlib
 
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
+
+_REQUIRED_COMPARISON_FIELDS = ("N", "k", "T", "dt")
+_OPTIONAL_COMPARISON_FIELDS = (
+    "x0_level", "beta", "lambda_terminal", "gamma_target", "x_target", "jitter",
+    "ring_range", "batch_size", "num_steps", "lr", "grad_clip",
+    "num_fourier_freqs", "embed_dim", "hidden_embed", "hidden_score", "score_bound",
+)
+
+
+def validate_run_configs(noisy_config, clean_config):
+    """Refuse comparisons with different problems, grids, or training settings."""
+    for field in (*_REQUIRED_COMPARISON_FIELDS, "sigma"):
+        if field not in noisy_config or field not in clean_config:
+            raise ValueError(f"Both run configurations must record {field!r}")
+    noisy_sigma = float(noisy_config["sigma"])
+    clean_sigma = float(clean_config["sigma"])
+    if not np.isfinite(noisy_sigma) or noisy_sigma <= 0:
+        raise ValueError("The noisy run must have finite sigma > 0")
+    if clean_sigma != 0:
+        raise ValueError("The noiseless reference must have sigma = 0")
+    for field in (*_REQUIRED_COMPARISON_FIELDS, *_OPTIONAL_COMPARISON_FIELDS):
+        if field in _REQUIRED_COMPARISON_FIELDS or field in noisy_config or field in clean_config:
+            if noisy_config.get(field) != clean_config.get(field):
+                raise ValueError(
+                    f"Incompatible {field}: noisy={noisy_config.get(field)!r}, "
+                    f"noiseless={clean_config.get(field)!r}"
+                )
 
 
-def main():
-    parser = argparse.ArgumentParser(description="E1 noise robustness panel")
-    parser.add_argument("--noisy", type=str, default="k2_T12_sig01",
-                        help="tag of the sigma > 0 run (out_1/e1_<tag>)")
-    parser.add_argument("--noiseless", type=str, default="k2_T12",
-                        help="tag of the sigma = 0 reference run")
-    parser.add_argument("--out", type=str, default="",
-                        help="output path (default: run dir + figure mirror)")
-    args = parser.parse_args()
+def resolve_run_directory(value, out_root):
+    path = Path(value).expanduser()
+    if path.is_dir():
+        return path.resolve()
+    if path.name != str(path) or path.name in {".", ".."}:
+        raise FileNotFoundError(f"Run directory does not exist: {path}")
+    tagged = Path(out_root).expanduser() / f"e1_{value}"
+    if not tagged.is_dir():
+        raise FileNotFoundError(f"Run directory does not exist: {tagged}")
+    return tagged.resolve()
 
-    noisy_dir = os.path.join("out_1", f"e1_{args.noisy}")
-    clean_dir = os.path.join("out_1", f"e1_{args.noiseless}")
-    with open(os.path.join(noisy_dir, "config.json")) as f:
-        cfg = json.load(f)
-    sigma, k, T = cfg["sigma"], cfg["k"], cfg["T"]
-    if sigma <= 0:
-        raise SystemExit(f"run {args.noisy} has sigma = {sigma}; need a noisy run")
 
-    noisy = np.load(os.path.join(noisy_dir, "trajectories.npz"))
-    clean = np.load(os.path.join(clean_dir, "trajectories.npz"))
+def load_comparison(noisy_dir, clean_dir):
+    """Load and check figure inputs before creating any output."""
+    configs = []
+    for directory in (noisy_dir, clean_dir):
+        with (Path(directory) / "config.json").open(encoding="utf-8") as stream:
+            configs.append(json.load(stream))
+    noisy_config, clean_config = configs
+    validate_run_configs(noisy_config, clean_config)
+    with np.load(Path(noisy_dir) / "trajectories.npz", allow_pickle=False) as archive:
+        noisy = {key: archive[key] for key in archive.files}
+    with np.load(Path(clean_dir) / "trajectories.npz", allow_pickle=False) as archive:
+        clean = {key: archive[key] for key in archive.files}
+    for name, data in (("noisy", noisy), ("noiseless", clean)):
+        if "t" not in data or data["t"].ndim != 1 or len(data["t"]) < 2:
+            raise ValueError(f"Invalid time grid in {name} trajectories")
+        if not np.isfinite(data["t"]).all() or not np.all(np.diff(data["t"]) > 0):
+            raise ValueError(f"Invalid time grid in {name} trajectories")
+        for key in ("r", "winding"):
+            suffixes = ("noisy_mean", "noisy_q10", "noisy_q90") if name == "noisy" else ("trained",)
+            for suffix in suffixes:
+                series = data.get(f"{key}_{suffix}")
+                if series is None or series.shape != data["t"].shape or not np.isfinite(series).all():
+                    raise ValueError(f"Missing or invalid {key}_{suffix} in {name} trajectories")
+            if name == "noisy" and np.any(data[f"{key}_noisy_q10"] > data[f"{key}_noisy_q90"]):
+                raise ValueError(f"Invalid {key} quantile band in noisy trajectories")
+    if noisy["t"].shape != clean["t"].shape or not np.allclose(noisy["t"], clean["t"]):
+        raise ValueError("Saved trajectory time grids differ")
+    expected_t = np.arange(len(noisy["t"])) * noisy_config["dt"]
+    if not np.allclose(noisy["t"], expected_t) or not np.isclose(expected_t[-1], noisy_config["T"]):
+        raise ValueError("Trajectory time grid does not match config T and dt")
+    return noisy_config, noisy, clean
 
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--noisy", required=True, help="sigma > 0 run directory or tag")
+    parser.add_argument("--noiseless", required=True, help="sigma = 0 run directory or tag")
+    parser.add_argument("--out-root", type=Path, default=Path(__file__).resolve().parent / "out_1")
+    parser.add_argument("--out", type=Path, help="output PDF (default: noisy run/r_noise_panel.pdf)")
+    parser.add_argument("--overwrite", action="store_true", help="explicitly replace an existing panel")
+    args = parser.parse_args(argv)
+    try:
+        noisy_dir = resolve_run_directory(args.noisy, args.out_root)
+        clean_dir = resolve_run_directory(args.noiseless, args.out_root)
+        config, noisy, clean = load_comparison(noisy_dir, clean_dir)
+        output = args.out.expanduser() if args.out else noisy_dir / "r_noise_panel.pdf"
+        if output.exists() and not args.overwrite:
+            raise FileExistsError(f"Output exists: {output}. Pass --overwrite to replace it.")
+    except (OSError, ValueError) as exc:
+        parser.exit(1, f"{exc}\n")
+
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    sigma, k, T = config["sigma"], config["k"], config["T"]
     fig, axes = plt.subplots(1, 2, figsize=(11, 3.8))
     panels = [("r", r"order parameter $r(t)$", (-0.03, 1.03)),
               ("winding", r"winding number $\mathcal{W}(t)$", (-0.25, k + 0.25))]
@@ -48,33 +114,22 @@ def main():
         ax.fill_between(noisy["t"], noisy[f"{key}_noisy_q10"],
                         noisy[f"{key}_noisy_q90"], color="tab:blue", alpha=0.25,
                         label="noisy ensemble, 10–90%")
-        ax.plot(noisy["t"], noisy[f"{key}_noisy_mean"], color="tab:blue", lw=2.0,
-                label=rf"noisy ensemble mean ($\sigma = {sigma:g}$)")
+        ax.plot(noisy["t"], noisy[f"{key}_noisy_mean"], color="tab:blue", lw=2,
+                label=rf"noise-trained control ($\sigma = {sigma:g}$)")
         ax.plot(clean["t"], clean[f"{key}_trained"], color="0.25", lw=1.8, ls="--",
-                label=r"deterministic ($\sigma = 0$)")
+                label=r"separately trained reference ($\sigma = 0$)")
         if key == "winding":
             for q in range(k + 1):
                 ax.axhline(q, color="0.85", lw=0.8, zorder=0)
-        ax.set_xlabel(r"$t$")
-        ax.set_ylabel(ylabel)
-        ax.set_ylim(*ylim)
+        ax.set(xlabel=r"$t$", ylabel=ylabel, ylim=ylim)
         ax.grid(alpha=0.2)
-    axes[0].legend(loc="lower right", fontsize=9)
-    fig.suptitle(rf"Trained control under noise ($k = {k}$, $T = {T:g}$)", y=1.02)
+    axes[0].legend(loc="lower right", fontsize=8)
+    fig.suptitle(rf"Separate training with and without noise ($k = {k}$, $T = {T:g}$)", y=1.02)
     fig.tight_layout()
-
-    if args.out:
-        paths = [args.out]
-    else:
-        suffix = f"_sigma{sigma:g}".replace(".", "p")
-        fig_dir = os.path.join("figures", f"e1_k{k}_T{T:g}{suffix}")
-        os.makedirs(fig_dir, exist_ok=True)
-        paths = [os.path.join(noisy_dir, "r_noise_panel.pdf"),
-                 os.path.join(fig_dir, "r_noise_panel.pdf")]
-    for p in paths:
-        fig.savefig(p, bbox_inches="tight")
-        print(f"wrote {p}")
+    output.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(output, bbox_inches="tight")
     plt.close(fig)
+    print(f"wrote {output}")
 
 
 if __name__ == "__main__":
diff --git a/e2_bc.py b/e2_bc.py
index 1fc77ff..fb05ce0 100644
--- a/e2_bc.py
+++ b/e2_bc.py
@@ -1,889 +1,14 @@
-# e2_bc.py
-#
-# Experiment E2: preventing opinion fragmentation (bounded confidence).
-# Paper reference: main5.tex, subsec:exp-bc, protocol in subsec:exp-protocol.
-#
-#   dynamics    : dx_i/dt = (1/N) Σ_j wbar_ij(t) φ_ε(x_j - x_i),  wbar = N * W
-#   interaction : φ_ε(z) = a · z · (1 - z²/ε²)₊²  — a C^{1,1} compactly
-#                 supported mollification of the sharp Hegselmann-Krause
-#                 influence z · 1{|z| ≤ ε}. Odd, φ_ε(0) = 0, Lipschitz
-#                 constant a, ||φ_ε||_∞ = (16/25√5) aε, support exactly [-ε, ε].
-#   initial data: x(0, ξ) = s(ξ - 1/2) + A sin(2πmξ)  (+ jitter during training)
-#                 — monotone-dominant in the label (readable kernel heatmaps),
-#                 span s ≫ ε so the uncontrolled dynamics fragments into
-#                 several frozen opinion clusters.
-#   baselines   : (i) mean-field  w ≡ 1  (W_ij = 1/N),
-#                 (ii) label-band: mollified bump of range r in |ξ - ζ| on the
-#                      interval (no torus wrap; opinions are not periodic) —
-#                      the admissible stand-in for the local bounded-confidence
-#                      variant. Optional --hk-ref adds the classical sharp
-#                      state-dependent HK kernel as a NON-admissible reference.
-#   claim       : uncontrolled, the population fragments into clusters with all
-#                 inter-cluster gaps > ε; such profiles are stationary for
-#                 EVERY admissible kernel (the interaction has exact support ε),
-#                 so control must act early — prevention, not reversal.
-#
-# Outputs (all under out_1/e2_<tag>/, paper-ready PDFs mirrored to
-# figures/e2_eps{eps}_T{T}/). All arms are evaluated and archived in
-# trajectories.npz/metrics.json; the FIGURES show mean-field vs trained
-# control (the paper's two-arm design), unless --uncontrolled-only.
-#   config.json, metrics.json, run_summary.txt, model.pt, history.json
-#   trajectories.npz     — x(t) per arm, W(t) trained, Var/K/authority series
-#   state_evolution.pdf  — x(t, ξ) heatmaps per arm (stripes vs consensus)
-#   opinion_traj.pdf     — classical HK "cluster tree": x_i(t) per arm
-#   kernel_strip.pdf     — trained kernel heatmaps at 6 times
-#   cluster_count.pdf    — number of opinion clusters K(t) per arm
-#   variance.pdf         — Var(x(t)) per arm (log scale)
-#   authority.pdf        — control authority (1/N) Σ_i max_j |φ_ε(x_j - x_i)|
-#   vxi.pdf              — receiver variation V_ξ(ŵ(t)) of the trained kernel
-#   phi.pdf              — φ_ε against the sharp HK influence
-#   training_history.pdf, opinions.gif, kernel_time.gif
-#
-# Usage:
-#   python e2_bc.py                        # full run (2000 steps)
-#   python e2_bc.py --uncontrolled-only    # baselines only: verify fragmentation
-#   python e2_bc.py --smoke                # 20-step smoke test
-#   python e2_bc.py --device mps           # Apple GPU
+"""Compatibility entry point; implementation lives in graphon_control.experiments.bounded_confidence."""
+from graphon_control.experiments.bounded_confidence import *  # noqa: F401,F403
+from graphon_control.experiments.bounded_confidence import main
 
-import argparse
-import dataclasses
-import datetime
-import json
-import math
-import os
-from dataclasses import dataclass
-from typing import Callable, Dict, List, Optional
+_LEGACY_PLOT_NAMES = ('fig_state_evolution', 'fig_opinion_traj', 'fig_kernel_strip', 'fig_cluster_count', 'fig_variance', 'fig_authority', 'fig_vxi', 'fig_phi', 'fig_training_history', 'save_opinion_gif', 'ARM_LABELS', 'ARM_COLORS')
 
-import numpy as np
-import torch
-import matplotlib
-
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-from tqdm import tqdm
-
-from model_1 import TimeLabelGraphonControl, _random_run_tag
-from model_1_viz import save_graphon_gif
-
-
-# ============================================================
-# Configuration (paper protocol; deviations stated in subsec:exp-bc)
-# ============================================================
-
-@dataclass
-class E2Config:
-    # interaction
-    eps: float = 1.0           # confidence radius
-    gain: float = 4.0          # gain a (Lipschitz constant of φ_ε)
-    phi_form: str = "poly"     # "poly" (compact support) | "gauss" (fallback)
-
-    # initial profile x(0, ξ) = span (ξ - 1/2) + amp sin(2π modes ξ)
-    span: float = 6.0
-    amp: float = 0.5
-    modes: int = 4
-    jitter: float = 0.1        # i.i.d. jitter on the training initial data
-
-    # label-band baseline
-    band_range: float = 0.10
-
-    # discretization
-    N: int = 64
-    T: float = 10.0
-    dt: float = 0.05
-    sigma: float = 0.0
-
-    # cost — PAPER convention (eqn:J-cost-1), with β = 1 (stated deviation from
-    # the e1 protocol's β = 10: merging clusters needs concentrated attention,
-    # and the pilot decides the final value; see --beta)
-    beta: float = 1.0
-    lambda_terminal: float = 0.0
-    gamma_target: float = 100.0
-    x_target: float = 0.0
-
-    # optimization
-    batch_size: int = 16
-    num_steps: int = 2000
-    lr: float = 1e-3
-    grad_clip: float = 1.0
-
-    # architecture (model_1.py defaults, as in e1)
-    num_fourier_freqs: int = 8
-    embed_dim: int = 64
-    hidden_embed: int = 128
-    hidden_score: int = 128
-
-    # bookkeeping
-    seed: int = 0
-    device: str = "cpu"
-    time_chunk: int = 50
-    hk_ref: bool = False       # add the sharp state-dependent HK reference arm
-
-    @property
-    def NT(self) -> int:
-        return int(round(self.T / self.dt))
-
-
-# ============================================================
-# E2 ingredients
-# ============================================================
-
-def phi_bc_poly(y: torch.Tensor, eps: float, gain: float) -> torch.Tensor:
-    """
-    Mollified bounded-confidence interaction
-      φ_ε(z) = a z (1 - z²/ε²)₊²,
-    C^{1,1}, odd, Lipschitz constant a, support exactly [-ε, ε].
-    """
-    u = y / eps
-    return gain * y * torch.clamp(1.0 - u * u, min=0.0) ** 2
-
-
-def phi_bc_gauss(y: torch.Tensor, eps: float, gain: float) -> torch.Tensor:
-    """Gaussian-tail fallback φ_ε(z) = a z exp(-z²/ε²) (no dead gradients)."""
-    return gain * y * torch.exp(-(y / eps) ** 2)
-
-
-def make_phi(cfg: E2Config) -> Callable[[torch.Tensor], torch.Tensor]:
-    base = {"poly": phi_bc_poly, "gauss": phi_bc_gauss}[cfg.phi_form]
-    return lambda y: base(y, cfg.eps, cfg.gain)
-
-
-def bc_profile(xi: torch.Tensor, span: float, amp: float, modes: int) -> torch.Tensor:
-    """Initial opinion profile x(0, ξ) = span (ξ - 1/2) + amp sin(2π modes ξ)."""
-    return span * (xi - 0.5) + amp * torch.sin(2.0 * math.pi * modes * xi)
-
-
-def mean_field_kernel(N: int, device, dtype=torch.float32) -> torch.Tensor:
-    """Row-stochastic representation of w ≡ 1: W_ij = 1/N."""
-    return torch.full((N, N), 1.0 / N, device=device, dtype=dtype)
-
-
-def band_kernel(xi: torch.Tensor, r: float) -> torch.Tensor:
-    """
-    Label-band kernel: smooth bump of range r in |ξ - ζ| on the INTERVAL
-    (no torus wrap — opinions are not periodic; row normalization absorbs the
-    boundary), row-normalized to a row-stochastic W. Smooth in ξ, hence of
-    finite receiver variation, unlike the sharp band indicator.
-    """
-    d = (xi.unsqueeze(0) - xi.unsqueeze(1)).abs()
-    u = (d / r).clamp(max=1.0 - 1e-6)
-    bump = torch.where(d < r, torch.exp(-1.0 / (1.0 - u ** 2)), torch.zeros_like(d))
-    return bump / bump.sum(dim=1, keepdim=True)
-
-
-class ConstantKernel:
-    """Baseline wrapper exposing the same call signature as the trained model."""
-
-    def __init__(self, W0: torch.Tensor):
-        self.W0 = W0
-
-    def __call__(self, t: torch.Tensor) -> torch.Tensor:
-        if t.ndim == 0:
-            return self.W0
-        return self.W0.unsqueeze(0).expand(t.shape[0], *self.W0.shape)
-
-
-# ============================================================
-# Rollouts and cost (e1 conventions)
-# ============================================================
-
-def precompute_Ws(model, NT: int, dt: float, device, chunk: int) -> torch.Tensor:
-    """Returns (NT, N, N); differentiable w.r.t. model parameters."""
-    t_vals = torch.arange(NT, dtype=torch.float32, device=device) * dt
-    outs = [model(t_vals[a:a + chunk]) for a in range(0, NT, chunk)]
-    outs = [o.unsqueeze(0) if o.ndim == 2 else o for o in outs]
-    return torch.cat(outs, dim=0)
-
-
-def rollout_with_Ws(
-    Ws: torch.Tensor,
-    x0: torch.Tensor,
-    dt: float,
-    phi: Callable[[torch.Tensor], torch.Tensor],
-    sigma: float = 0.0,
-) -> torch.Tensor:
-    """
-    Euler(-Maruyama) rollout with precomputed kernels.
-    Ws: (NT, N, N) row-stochastic; x0: (B, N).
-    Sign convention (paper, eqn FN definition): diff[b, i, j] = x_j - x_i.
-    Returns xs: (NT + 1, B, N).
-    """
-    B, N = x0.shape
-    NT = Ws.shape[0]
-    sqrt_dt = math.sqrt(dt)
-    x = x0
-    xs = [x]
-    for n in range(NT):
-        diff = x.unsqueeze(1) - x.unsqueeze(2)            # (B, N, N): x_j - x_i
-        drift = (Ws[n].unsqueeze(0) * phi(diff)).sum(-1)  # (B, N)
-        x = x + dt * drift
-        if sigma > 0.0:
-            x = x + sigma * sqrt_dt * torch.randn_like(x)
-        xs.append(x)
-    return torch.stack(xs, dim=0)
-
-
-def rollout_hk_reference(x0: torch.Tensor, NT: int, dt: float, eps: float) -> torch.Tensor:
-    """
-    Classical sharp Hegselmann-Krause reference: state-dependent kernel
-    W_ij(t) ∝ 1{|x_i - x_j| ≤ ε}, row-normalized, with φ(z) = z within range.
-    NOT admissible in our control class (state-dependent, infinite V_ξ);
-    reported only as a reference portrait. x0: (B, N) -> (NT+1, B, N).
-    """
-    x = x0
-    xs = [x]
-    for _ in range(NT):
-        diff = x.unsqueeze(1) - x.unsqueeze(2)            # (B, N, N): x_j - x_i
-        adj = (diff.abs() <= eps).to(x.dtype)
-        W = adj / adj.sum(dim=-1, keepdim=True)           # diagonal is always in range
-        x = x + dt * (W * diff).sum(-1)
-        xs.append(x)
-    return torch.stack(xs, dim=0)
-
-
-def e2_cost(xs: torch.Tensor, Ws: torch.Tensor, cfg: E2Config) -> Dict[str, torch.Tensor]:
-    """
-    Discretized cost J^{N,Δt}, PAPER convention (eqn: JN definition / eqn:J-cost-1):
-      Σ_n Δt [ Var(x^n) + β ||w(t_n)||²_{L²(I²)} ] + γ (mean(x^T) - x_target)²
-    where ||w||²_{L²(I²)} = Σ_ij W_ij² for the row-stochastic W (wbar = N·W).
-    xs: (NT+1, B, N); Ws: (NT, N, N).
-    """
-    var_t = xs.var(dim=2, unbiased=False).mean(dim=1)
-    running_state = cfg.dt * var_t[:-1].sum()
-    running_reg = cfg.dt * cfg.beta * (Ws ** 2).sum()
-    terminal_var = cfg.lambda_terminal * var_t[-1]
-    mean_T = xs[-1].mean(dim=1)
-    terminal_target = cfg.gamma_target * ((mean_T - cfg.x_target) ** 2).mean()
-    total = running_state + running_reg + terminal_var + terminal_target
-    return {
-        "total": total,
-        "state_cost": running_state,
-        "l2_reg_cost": running_reg,
-        "terminal_var": var_t[-1],
-        "terminal_target": terminal_target,
-        "terminal_mean": mean_T.mean(),
-    }
-
-
-# ============================================================
-# Diagnostics
-# ============================================================
-
-def cluster_count(x: torch.Tensor, eps: float) -> torch.Tensor:
-    """
-    Number of opinion clusters: 1 + #{gaps > ε in the sorted state}. A gap
-    larger than ε is definitively frozen — the interaction has support ε, so
-    no admissible kernel can act across it. x: (..., N) -> (...) int.
-    """
-    xs, _ = torch.sort(x, dim=-1)
-    gaps = xs[..., 1:] - xs[..., :-1]
-    return 1 + (gaps > eps).sum(dim=-1)
-
-
-def cluster_inventory(x: torch.Tensor, eps: float) -> List[Dict[str, float]]:
-    """Cluster centers and mass fractions of a single state x: (N,)."""
-    xs, _ = torch.sort(x)
-    xs = xs.cpu().numpy()
-    breaks = np.where(np.diff(xs) > eps)[0]
-    out = []
-    start = 0
-    for b in list(breaks) + [len(xs) - 1]:
-        block = xs[start:b + 1]
-        out.append({"center": float(block.mean()), "mass": len(block) / len(xs)})
-        start = b + 1
-    return out
-
-
-def control_authority(x: torch.Tensor,
-                      phi: Callable[[torch.Tensor], torch.Tensor]) -> torch.Tensor:
-    """
-    Control authority (1/N) Σ_i max_j |φ_ε(x_j - x_i)|: the strongest force any
-    admissible kernel can still exert, averaged over receivers. Decays to 0 both
-    at fragmentation (window of controllability closed) and at consensus
-    (nothing left to do) — the two are distinguished by Var. x: (..., N) -> (...).
-    """
-    diff = x.unsqueeze(-2) - x.unsqueeze(-1)   # (..., N, N): [i, j] = x_j - x_i
-    return phi(diff).abs().amax(dim=-1).mean(dim=-1)
-
-
-def vxi_discrete(W: torch.Tensor) -> torch.Tensor:
-    """
-    Discrete receiver variation (paper eqn:receiver-tv-discrete) of the graphon
-    values wbar = N * W:
-      V_ξ = Σ_{i=1}^{N-1} sqrt( (1/N) Σ_j (wbar_{i+1,j} - wbar_{i,j})^2 ).
-    Accepts (N, N) or batched (..., N, N).
-    """
-    wbar = W.shape[-1] * W
-    d = wbar[..., 1:, :] - wbar[..., :-1, :]
-    return torch.sqrt((d ** 2).mean(dim=-1)).sum(dim=-1)
-
-
-# ============================================================
-# Figures
-# ============================================================
-
-ARM_LABELS = {"meanfield": r"mean-field  $w \equiv 1$",
-              "band": "label-band",
-              "trained": "trained control",
-              "hk": "sharp HK (reference)"}
-ARM_COLORS = {"meanfield": "tab:gray", "band": "tab:orange",
-              "trained": "tab:blue", "hk": "tab:green"}
-
-
-def fig_state_evolution(t_grid, xs_arms: Dict[str, np.ndarray], arms: List[str],
-                        cfg: E2Config, path: str):
-    lim = cfg.span / 2 + cfg.amp + 0.5
-    fig, axes = plt.subplots(1, len(arms), figsize=(4.3 * len(arms), 3.6),
-                             sharey=True, squeeze=False)
-    for ax, arm in zip(axes[0], arms):
-        X = xs_arms[arm]  # (NT+1, N)
-        im = ax.imshow(X.T, origin="lower", aspect="auto",
-                       extent=(t_grid[0], t_grid[-1], 0.0, 1.0),
-                       cmap="RdBu_r", vmin=-lim, vmax=lim)
-        ax.set_title(ARM_LABELS[arm])
-        ax.set_xlabel(r"$t$")
-    axes[0][0].set_ylabel(r"label $\xi$")
-    fig.colorbar(im, ax=axes, fraction=0.02, pad=0.01, label=r"$x(t,\xi)$")
-    fig.suptitle("State evolution", y=1.04)
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_opinion_traj(t_grid, xi, xs_arms: Dict[str, np.ndarray], arms: List[str],
-                     cfg: E2Config, path: str):
-    """The classical HK cluster-tree: opinion trajectories colored by label."""
-    colors = plt.cm.viridis(xi)
-    fig, axes = plt.subplots(1, len(arms), figsize=(4.3 * len(arms), 3.6),
-                             sharey=True, squeeze=False)
-    for ax, arm in zip(axes[0], arms):
-        X = xs_arms[arm]  # (NT+1, N)
-        for i in range(X.shape[1]):
-            ax.plot(t_grid, X[:, i], color=colors[i], lw=0.8, alpha=0.8)
-        ax.axhline(cfg.x_target, color="0.3", lw=1.0, ls="--")
-        ax.set_title(ARM_LABELS[arm])
-        ax.set_xlabel(r"$t$")
-        ax.grid(alpha=0.15)
-    # ε scale bar on the first panel: the interaction's dead-zone width
-    ax0 = axes[0][0]
-    t_bar = t_grid[0] + 0.05 * (t_grid[-1] - t_grid[0])
-    y0 = cfg.span / 2 + cfg.amp - cfg.eps
-    ax0.annotate("", xy=(t_bar, y0), xytext=(t_bar, y0 + cfg.eps),
-                 arrowprops=dict(arrowstyle="<->", color="0.2", lw=1.2))
-    ax0.text(t_bar + 0.02 * (t_grid[-1] - t_grid[0]), y0 + cfg.eps / 2,
-             r"$\varepsilon$", va="center", fontsize=11, color="0.2")
-    axes[0][0].set_ylabel(r"opinion $x_i(t)$")
-    fig.suptitle("Opinion trajectories", y=1.02)
-    fig.tight_layout()
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_kernel_strip(Ws: torch.Tensor, dt: float, path: str, n_panels: int = 6,
-                     times: List[float] = None):
-    """
-    Kernel heatmaps on a LOG color scale: the merge-phase attention peaks are
-    an order of magnitude above the uniform level wbar = 1, and a linear scale
-    would flatten the relax-to-uniform endpoint. Panel times default to a
-    uniform grid; pass `times` to concentrate panels in the merge window.
-    """
-    from matplotlib.colors import LogNorm
-    NT = Ws.shape[0]
-    if times is None:
-        idxs = np.linspace(0, NT - 1, n_panels).round().astype(int)
-    else:
-        idxs = np.clip(np.round(np.array(times) / dt).astype(int), 0, NT - 1)
-        n_panels = len(idxs)
-    wbars = [Ws.shape[-1] * Ws[i].cpu().numpy() for i in idxs]
-    vmax = max(w.max() for w in wbars)
-    vmin = max(min(w.min() for w in wbars), vmax * 1e-3)
-    fig, axes = plt.subplots(1, n_panels, figsize=(2.6 * n_panels, 2.9))
-    for ax, i, w in zip(axes, idxs, wbars):
-        im = ax.imshow(np.maximum(w, vmin), origin="lower", extent=(0, 1, 0, 1),
-                       aspect="equal", cmap="viridis", interpolation="bicubic",
-                       norm=LogNorm(vmin=vmin, vmax=vmax))
-        ax.set_title(rf"$t = {i * dt:.1f}$", fontsize=10)
-        ax.set_xticks([0, 1])
-        ax.set_yticks([0, 1])
-        if ax is axes[0]:
-            ax.set_ylabel(r"$\xi$")
-        ax.set_xlabel(r"$\zeta$")
-    fig.colorbar(im, ax=axes, fraction=0.02, pad=0.01, label=r"$\hat{w}(t,\xi,\zeta)$")
-    fig.suptitle("Trained kernel over time", y=1.04)
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_cluster_count(t_grid, K_series: Dict[str, np.ndarray], arms: List[str],
-                      path: str):
-    fig, ax = plt.subplots(figsize=(6.4, 4.0))
-    for arm in arms:
-        ax.step(t_grid, K_series[arm], where="post", color=ARM_COLORS[arm],
-                lw=2.0, label=ARM_LABELS[arm])
-    ax.axhline(1, color="0.85", lw=0.8, zorder=0)
-    ax.set_xlabel(r"$t$")
-    ax.set_ylabel(r"clusters $K(t)$")
-    Kmax = max(int(K_series[a].max()) for a in arms)
-    ax.set_yticks(range(1, Kmax + 1))
-    ax.set_ylim(0.75, Kmax + 0.25)
-    ax.grid(alpha=0.15)
-    ax.legend()
-    ax.set_title("Opinion cluster count")
-    fig.tight_layout()
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_variance(t_grid, var_series: Dict[str, np.ndarray], arms: List[str],
-                 path: str):
-    fig, ax = plt.subplots(figsize=(6.4, 4.0))
-    for arm in arms:
-        ax.semilogy(t_grid, np.maximum(var_series[arm], 1e-12),
-                    color=ARM_COLORS[arm], lw=2.0, label=ARM_LABELS[arm])
-    ax.set_xlabel(r"$t$")
-    ax.set_ylabel(r"$\mathrm{Var}(x(t))$")
-    ax.grid(alpha=0.25, which="both")
-    ax.legend()
-    ax.set_title("Variance: frozen plateau vs consensus")
-    fig.tight_layout()
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_authority(t_grid, auth_series: Dict[str, np.ndarray], arms: List[str],
-                  path: str):
-    fig, ax = plt.subplots(figsize=(6.4, 4.0))
-    for arm in arms:
-        ax.plot(t_grid, auth_series[arm], color=ARM_COLORS[arm], lw=2.0,
-                label=ARM_LABELS[arm])
-    ax.set_xlabel(r"$t$")
-    ax.set_ylabel(r"$\frac{1}{N}\sum_i \max_j |\varphi_\varepsilon(x_j - x_i)|$")
-    ax.grid(alpha=0.25)
-    ax.legend()
-    ax.set_title("Control authority (window of controllability)")
-    fig.tight_layout()
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_vxi(t_grid_W, vxi_trained: np.ndarray, vxi_band: float, path: str):
-    fig, ax = plt.subplots(figsize=(6.4, 4.0))
-    ax.plot(t_grid_W, vxi_trained, color=ARM_COLORS["trained"], lw=2.0,
-            label=r"trained $V_\xi(\hat{w}(t))$")
-    ax.axhline(0.0, color=ARM_COLORS["meanfield"], lw=1.5, ls="--",
-               label=r"mean-field ($V_\xi = 0$)")
-    ax.axhline(vxi_band, color=ARM_COLORS["band"], lw=1.5, ls="--",
-               label=rf"label-band ($V_\xi = {vxi_band:.1f}$)")
-    ax.set_xlabel(r"$t$")
-    ax.set_ylabel(r"$V_\xi$")
-    ax.grid(alpha=0.25)
-    ax.legend()
-    ax.set_title("Receiver variation of the kernel (measured, not enforced)")
-    fig.tight_layout()
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_phi(cfg: E2Config, path: str):
-    z = torch.linspace(-2.0 * cfg.eps, 2.0 * cfg.eps, 801)
-    phi = make_phi(cfg)
-    sharp = torch.where(z.abs() <= cfg.eps, cfg.gain * z, torch.zeros_like(z))
-    fig, ax = plt.subplots(figsize=(6.0, 3.6))
-    ax.plot(z.numpy(), sharp.numpy(), color="0.6", lw=1.5, ls="--",
-            label=r"sharp HK  $a\,z\,\mathbf{1}\{|z| \leq \varepsilon\}$")
-    ax.plot(z.numpy(), phi(z).numpy(), color="tab:blue", lw=2.2,
-            label=r"$\varphi_\varepsilon(z) = a\,z\,(1 - z^2/\varepsilon^2)_+^2$")
-    for s in (-1, 1):
-        ax.axvline(s * cfg.eps, color="0.85", lw=0.8, zorder=0)
-    ax.set_xlabel(r"$z$")
-    ax.set_ylabel(r"$\varphi(z)$")
-    ax.grid(alpha=0.2)
-    ax.legend(fontsize=9)
-    ax.set_title(rf"Bounded-confidence interaction ($\varepsilon = {cfg.eps:g}$, "
-                 rf"$a = {cfg.gain:g}$)")
-    fig.tight_layout()
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def fig_training_history(history: Dict[str, List[float]], path: str):
-    fig, axes = plt.subplots(1, 2, figsize=(11, 3.8))
-    axes[0].plot(history["loss"], lw=1.0, color="tab:blue")
-    axes[0].set_title("total loss")
-    axes[0].set_xlabel("step")
-    axes[0].grid(alpha=0.25)
-    for key, color in [("state_cost", "tab:green"), ("l2_reg_cost", "tab:red"),
-                       ("terminal_var", "tab:purple"), ("terminal_target", "tab:brown")]:
-        axes[1].semilogy(np.maximum(np.array(history[key]), 1e-10), lw=1.0,
-                         label=key, color=color)
-    axes[1].set_title("components (log scale)")
-    axes[1].set_xlabel("step")
-    axes[1].grid(alpha=0.25, which="both")
-    axes[1].legend(fontsize=8)
-    fig.tight_layout()
-    fig.savefig(path, bbox_inches="tight")
-    plt.close(fig)
-
-
-def save_opinion_gif(t_grid, xi, xs_arms: Dict[str, np.ndarray],
-                     K_series: Dict[str, np.ndarray], arms: List[str],
-                     cfg: E2Config, path: str,
-                     stride: int = 2, fps: int = 12, dpi: int = 80):
-    """Animated opinion axis: dot at (x_i(t), ξ_i), colored by label."""
-    from PIL import Image
-    from io import BytesIO
-
-    colors = plt.cm.viridis(xi)
-    lim = cfg.span / 2 + cfg.amp + 0.7
-    frames = list(range(0, len(t_grid), stride))
-    if frames[-1] != len(t_grid) - 1:
-        frames.append(len(t_grid) - 1)
-
-    pil_imgs = []
-    for n in frames:
-        fig, axes = plt.subplots(1, len(arms), figsize=(3.5 * len(arms), 3.6),
-                                 squeeze=False)
-        for ax, arm in zip(axes[0], arms):
-            x_n = xs_arms[arm][n]  # (N,)
-            ax.scatter(x_n, xi, c=colors, s=22, edgecolors="none")
-            ax.axvline(cfg.x_target, color="0.8", lw=1.0, ls="--", zorder=0)
-            ax.set_xlim(-lim, lim)
-            ax.set_ylim(-0.03, 1.03)
-            ax.set_xlabel(r"opinion $x$")
-            ax.set_title(f"{ARM_LABELS[arm]}\n$K = {int(K_series[arm][n])}$",
-                         fontsize=10)
-        axes[0][0].set_ylabel(r"label $\xi$")
-        fig.suptitle(rf"Bounded confidence:  $t = {t_grid[n]:.2f}$", y=0.99)
-        fig.tight_layout()
-        buf = BytesIO()
-        fig.savefig(buf, format="png", dpi=dpi)
-        plt.close(fig)
-        buf.seek(0)
-        pil_imgs.append(Image.open(buf).copy())
-
-    pil_imgs[0].save(path, save_all=True, append_images=pil_imgs[1:], loop=0,
-                     duration=int(1000 / fps), optimize=False)
-
-
-# ============================================================
-# Evaluation helper (shared by full and uncontrolled-only runs)
-# ============================================================
-
-def evaluate_arms(arm_specs: Dict[str, object], x0_eval: torch.Tensor,
-                  cfg: E2Config, phi, device):
-    """
-    arm_specs values: a kernel (ConstantKernel / trained model), or the string
-    "hk" for the sharp state-dependent reference. Deterministic, exact datum.
-    """
-    out = {}
-    with torch.no_grad():
-        for arm, spec in arm_specs.items():
-            if spec == "hk":
-                xs_eval = rollout_hk_reference(x0_eval, cfg.NT, cfg.dt, cfg.eps)
-                Ws_eval, parts = None, None
-            else:
-                Ws_eval = precompute_Ws(spec, cfg.NT, cfg.dt, device, cfg.time_chunk)
-                xs_eval = rollout_with_Ws(Ws_eval, x0_eval, cfg.dt, phi, sigma=0.0)
-                parts = e2_cost(xs_eval, Ws_eval, cfg)
-            x_t = xs_eval[:, 0, :]                               # (NT+1, N)
-            out[arm] = {
-                "X": x_t.cpu().numpy(),
-                "Ws": Ws_eval,
-                "var": x_t.var(dim=1, unbiased=False).cpu().numpy(),
-                "K": cluster_count(x_t, cfg.eps).cpu().numpy(),
-                "authority": control_authority(x_t, phi).cpu().numpy(),
-                "cost": None if parts is None
-                        else {kk: float(vv) for kk, vv in parts.items()},
-            }
-    return out
-
-
-def freeze_time(K: np.ndarray, dt: float) -> Optional[float]:
-    """First time after which K(t) never changes again (None if K constant)."""
-    changes = np.where(K[1:] != K[:-1])[0]
-    if len(changes) == 0:
-        return None
-    return float((changes[-1] + 1) * dt)
-
-
-# ============================================================
-# Main experiment
-# ============================================================
-
-def run_e2(cfg: E2Config, tag: str, out_root: str = "out_1", fig_root: str = "figures",
-           mirror_figs: bool = True, uncontrolled_only: bool = False):
-    torch.manual_seed(cfg.seed)
-    np.random.seed(cfg.seed)
-    device = torch.device(cfg.device)
-    dtype = torch.float32
-
-    run_dir = os.path.join(out_root, f"e2_{tag}")
-    fig_dir = os.path.join(fig_root, f"e2_eps{cfg.eps:g}_T{cfg.T:g}".replace(".", "p"))
-    os.makedirs(run_dir, exist_ok=True)
-    os.makedirs(fig_dir, exist_ok=True)
-
-    xi = (torch.arange(cfg.N, device=device, dtype=dtype) + 0.5) / cfg.N
-    base = bc_profile(xi, cfg.span, cfg.amp, cfg.modes)
-    NT = cfg.NT
-    phi = make_phi(cfg)
-
-    print("=" * 64)
-    print(f"E2 bounded confidence | eps={cfg.eps} gain={cfg.gain} "
-          f"span={cfg.span} N={cfg.N} T={cfg.T} dt={cfg.dt}")
-    print(f"steps={cfg.num_steps} batch={cfg.batch_size} lr={cfg.lr} "
-          f"beta={cfg.beta} gamma={cfg.gamma_target} phi={cfg.phi_form}")
-    print(f"device={cfg.device} tag={tag}")
-    print("=" * 64)
-
-    # ---------- sanity gate: the uncontrolled arm must fragment ----------
-    # (the narrative claim of subsec:exp-bc is false for this config otherwise)
-    with torch.no_grad():
-        W_mf = mean_field_kernel(cfg.N, device)
-        Ws_gate = W_mf.unsqueeze(0).expand(NT, -1, -1)
-        xs_gate = rollout_with_Ws(Ws_gate, base.unsqueeze(0), cfg.dt, phi)
-        K_gate = int(cluster_count(xs_gate[-1, 0], cfg.eps))
-    print(f"sanity gate: uncontrolled mean-field arm ends with K = {K_gate} clusters")
-    if K_gate < 2 and not uncontrolled_only:
-        raise RuntimeError(
-            f"uncontrolled arm does not fragment (K={K_gate}); "
-            f"raise --span, lower --eps, or raise --gain")
-
-    arm_specs: Dict[str, object] = {
-        "meanfield": ConstantKernel(W_mf),
-        "band": ConstantKernel(band_kernel(xi, cfg.band_range)),
-    }
-
-    # ---------- training ----------
-    model = None
-    history: Dict[str, List[float]] = {kk: [] for kk in
-                                       ["loss", "state_cost", "l2_reg_cost",
-                                        "terminal_var", "terminal_target",
-                                        "terminal_mean"]}
-    if not uncontrolled_only:
-        model = TimeLabelGraphonControl(
-            xi=xi,
-            num_fourier_freqs=cfg.num_fourier_freqs,
-            embed_dim=cfg.embed_dim,
-            hidden_embed=cfg.hidden_embed,
-            hidden_score=cfg.hidden_score,
-            time_scale=max(cfg.T, 1e-8),
-            enforce_symmetry=False,
-        ).to(device)
-        optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr)
-
-        pbar = tqdm(range(cfg.num_steps), desc="E2 training", unit="step")
-        for _ in pbar:
-            model.train()
-            optimizer.zero_grad()
-            Ws = precompute_Ws(model, NT, cfg.dt, device, cfg.time_chunk)
-            x0 = base.unsqueeze(0) + cfg.jitter * torch.randn(
-                cfg.batch_size, cfg.N, device=device, dtype=dtype)
-            xs = rollout_with_Ws(Ws, x0, cfg.dt, phi, sigma=cfg.sigma)
-            parts = e2_cost(xs, Ws, cfg)
-            parts["total"].backward()
-            if cfg.grad_clip is not None:
-                torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
-            optimizer.step()
-
-            for key in ["loss", "state_cost", "l2_reg_cost",
-                        "terminal_var", "terminal_target", "terminal_mean"]:
-                src = "total" if key == "loss" else key
-                history[key].append(float(parts[src].detach()))
-            pbar.set_postfix(loss=f"{history['loss'][-1]:.3f}",
-                             term_var=f"{history['terminal_var'][-1]:.4f}")
-
-        model.eval()
-        arm_specs["trained"] = model
-
-    if cfg.hk_ref:
-        arm_specs["hk"] = "hk"
-
-    # ---------- evaluation: deterministic, exact datum ----------
-    x0_eval = base.unsqueeze(0)
-    results = evaluate_arms(arm_specs, x0_eval, cfg, phi, device)
-    arms = list(arm_specs.keys())
-    cost_arms = [a for a in arms if results[a]["cost"] is not None]
-
-    vxi_band = float(vxi_discrete(results["band"]["Ws"][0]))
-    vxi_mf = float(vxi_discrete(results["meanfield"]["Ws"][0]))
-    vxi_trained = (vxi_discrete(results["trained"]["Ws"]).cpu().numpy()
-                   if "trained" in results else None)
-
-    t_grid = np.arange(NT + 1) * cfg.dt
-    t_grid_W = np.arange(NT) * cfg.dt
-
-    # ---------- persistence ----------
-    with open(os.path.join(run_dir, "config.json"), "w") as f:
-        json.dump(dataclasses.asdict(cfg), f, indent=2)
-    if model is not None:
-        torch.save(model.state_dict(), os.path.join(run_dir, "model.pt"))
-        with open(os.path.join(run_dir, "history.json"), "w") as f:
-            json.dump(history, f, indent=2)
-
-    npz = {"t": t_grid, "t_W": t_grid_W, "xi": xi.cpu().numpy(),
-           "W_band": results["band"]["Ws"][0].cpu().numpy().astype(np.float32)}
-    for arm in arms:
-        npz[f"x_{arm}"] = results[arm]["X"]
-        npz[f"var_{arm}"] = results[arm]["var"]
-        npz[f"K_{arm}"] = results[arm]["K"]
-        npz[f"authority_{arm}"] = results[arm]["authority"]
-    if "trained" in results:
-        npz["W_trained"] = results["trained"]["Ws"].cpu().numpy().astype(np.float32)
-        npz["vxi_trained"] = vxi_trained
-    np.savez_compressed(os.path.join(run_dir, "trajectories.npz"), **npz)
-
-    var_final = {arm: float(results[arm]["var"][-1]) for arm in arms}
-    K_final = {arm: int(results[arm]["K"][-1]) for arm in arms}
-    metrics = {
-        "eps": cfg.eps, "gain": cfg.gain, "span": cfg.span, "beta": cfg.beta,
-        "costs": {arm: results[arm]["cost"] for arm in cost_arms},
-        "var_final": var_final,
-        "K_final": K_final,
-        "clusters_meanfield_T": cluster_inventory(
-            torch.from_numpy(results["meanfield"]["X"][-1]), cfg.eps),
-        "freeze_time_meanfield": freeze_time(results["meanfield"]["K"], cfg.dt),
-        "vxi": {"band": vxi_band, "meanfield": vxi_mf},
-    }
-    if "trained" in results:
-        cons = np.where(results["trained"]["var"] < 1e-2)[0]
-        metrics["consensus_time_trained"] = (float(cons[0] * cfg.dt)
-                                             if len(cons) else None)
-        metrics["terminal_mean_trained"] = results["trained"]["cost"]["terminal_mean"]
-        metrics["vxi"]["trained_max"] = float(vxi_trained.max())
-        metrics["vxi"]["trained_final"] = float(vxi_trained[-1])
-        metrics["value_gap_meanfield_minus_trained"] = (
-            results["meanfield"]["cost"]["total"] - results["trained"]["cost"]["total"])
-    with open(os.path.join(run_dir, "metrics.json"), "w") as f:
-        json.dump(metrics, f, indent=2)
-
-    # ---------- figures ----------
-    def both(name: str) -> List[str]:
-        if not mirror_figs:
-            return [os.path.join(run_dir, name)]
-        return [os.path.join(run_dir, name), os.path.join(fig_dir, name)]
-
-    xs_arms_np = {arm: results[arm]["X"] for arm in arms}
-    K_series = {arm: results[arm]["K"] for arm in arms}
-    var_series = {arm: results[arm]["var"] for arm in arms}
-    auth_series = {arm: results[arm]["authority"] for arm in arms}
-
-    # paper figures show mean-field vs trained only; the band arm stays in the
-    # npz/metrics archive (uncontrolled-only runs keep whatever arms they have)
-    fig_arms = ([a for a in arms if a in ("meanfield", "trained")]
-                if "trained" in results else arms)
-
-    for p in both("state_evolution.pdf"):
-        fig_state_evolution(t_grid, xs_arms_np, fig_arms, cfg, p)
-    for p in both("opinion_traj.pdf"):
-        fig_opinion_traj(t_grid, xi.cpu().numpy(), xs_arms_np, fig_arms, cfg, p)
-    for p in both("cluster_count.pdf"):
-        fig_cluster_count(t_grid, K_series, fig_arms, p)
-    for p in both("variance.pdf"):
-        fig_variance(t_grid, var_series, fig_arms, p)
-    for p in both("authority.pdf"):
-        fig_authority(t_grid, auth_series, fig_arms, p)
-    for p in both("phi.pdf"):
-        fig_phi(cfg, p)
-    if "trained" in results:
-        strip_times = [0.0, 0.1 * cfg.T, 0.2 * cfg.T, 0.3 * cfg.T, 0.4 * cfg.T, cfg.T]
-        for p in both("kernel_strip.pdf"):
-            fig_kernel_strip(results["trained"]["Ws"], cfg.dt, p, times=strip_times)
-        for p in both("vxi.pdf"):
-            fig_vxi(t_grid_W, vxi_trained, vxi_band, p)
-        fig_training_history(history, os.path.join(run_dir, "training_history.pdf"))
-
-    # ---------- GIFs ----------
-    save_opinion_gif(t_grid, xi.cpu().numpy(), xs_arms_np, K_series, fig_arms, cfg,
-                     os.path.join(run_dir, "opinions.gif"))
-    if "trained" in results:
-        n_gif = min(60, NT)
-        gif_idx = np.linspace(0, NT - 1, n_gif).round().astype(int)
-        gif_frames = [results["trained"]["Ws"][i].cpu() * cfg.N for i in gif_idx]
-        save_graphon_gif(gif_frames, os.path.join(run_dir, "kernel_time.gif"),
-                         fps=10, t_vals=[float(i * cfg.dt) for i in gif_idx])
-
-    # ---------- summary ----------
-    lines = [
-        "=== E2 bounded confidence: run summary ===",
-        f"timestamp : {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
-        f"run_dir   : {run_dir}",
-        f"eps = {cfg.eps}, gain = {cfg.gain}, span = {cfg.span}, "
-        f"N = {cfg.N}, T = {cfg.T}, dt = {cfg.dt}, steps = {cfg.num_steps}, "
-        f"beta = {cfg.beta}",
-        "",
-        "arm         J_total      Var(T)     K(T)   mean(T)",
-    ]
-    for arm in arms:
-        j = (f"{results[arm]['cost']['total']:>9.4f}"
-             if results[arm]["cost"] is not None else "        —")
-        mean_T = float(results[arm]["X"][-1].mean())
-        lines.append(f"{arm:<11s} {j}  {var_final[arm]:>10.6f}   "
-                     f"{K_final[arm]:>3d}   {mean_T:>+8.4f}")
-    lines.append("")
-    lines.append(f"mean-field clusters at T: "
-                 f"{[(round(c['center'], 2), round(c['mass'], 3)) for c in metrics['clusters_meanfield_T']]}")
-    if metrics["freeze_time_meanfield"] is not None:
-        lines.append(f"mean-field freeze time: t = {metrics['freeze_time_meanfield']:.2f}")
-    if "trained" in results:
-        lines.append(
-            f"V_xi trained: max {metrics['vxi']['trained_max']:.3f}, "
-            f"final {metrics['vxi']['trained_final']:.3f} "
-            f"(band {vxi_band:.3f}, mean-field {vxi_mf:.3f})")
-        lines.append(f"value gap (mean-field - trained): "
-                     f"{metrics['value_gap_meanfield_minus_trained']:.4f}")
-        if metrics["consensus_time_trained"] is not None:
-            lines.append(f"trained consensus time (Var < 1e-2): "
-                         f"t = {metrics['consensus_time_trained']:.2f}")
-    summary = "\n".join(lines)
-    with open(os.path.join(run_dir, "run_summary.txt"), "w") as f:
-        f.write(summary + "\n")
-    print(summary)
-    print(f"\nAll outputs in: {run_dir}")
-    if mirror_figs:
-        print(f"Paper figures mirrored to: {fig_dir}")
-    return metrics
-
-
-# ============================================================
-# CLI
-# ============================================================
+def __getattr__(name):
+    if name in _LEGACY_PLOT_NAMES:
+        from importlib import import_module
+        return getattr(import_module("graphon_control.experiments.bounded_confidence_plots"), name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="E2: bounded-confidence cluster formation and prevention")
-    parser.add_argument("--eps", type=float, default=1.0, help="confidence radius")
-    parser.add_argument("--gain", type=float, default=4.0, help="interaction gain a")
-    parser.add_argument("--span", type=float, default=6.0, help="initial opinion span")
-    parser.add_argument("--amp", type=float, default=0.5, help="seed sine amplitude")
-    parser.add_argument("--modes", type=int, default=4, help="seed sine mode count")
-    parser.add_argument("--phi-form", type=str, default="poly", dest="phi_form",
-                        choices=["poly", "gauss"])
-    parser.add_argument("--T", type=float, default=10.0, help="time horizon")
-    parser.add_argument("--N", type=int, default=64)
-    parser.add_argument("--steps", type=int, default=2000)
-    parser.add_argument("--batch", type=int, default=16)
-    parser.add_argument("--beta", type=float, default=1.0)
-    parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument("--device", type=str, default="cpu",
-                        choices=["cpu", "mps", "cuda"])
-    parser.add_argument("--tag", type=str, default="")
-    parser.add_argument("--smoke", action="store_true", help="20-step smoke test")
-    parser.add_argument("--no-mirror", action="store_true",
-                        help="do not mirror figures into figures/ (pilot runs)")
-    parser.add_argument("--uncontrolled-only", action="store_true",
-                        dest="uncontrolled_only",
-                        help="baselines only: verify fragmentation, no training")
-    parser.add_argument("--hk-ref", action="store_true", dest="hk_ref",
-                        help="add the sharp state-dependent HK reference arm")
-    args = parser.parse_args()
-
-    cfg = E2Config(eps=args.eps, gain=args.gain, span=args.span, amp=args.amp,
-                   modes=args.modes, phi_form=args.phi_form, T=args.T, N=args.N,
-                   num_steps=args.steps, batch_size=args.batch, beta=args.beta,
-                   seed=args.seed, device=args.device, hk_ref=args.hk_ref)
-    if args.smoke:
-        cfg.num_steps = 20
-
-    tag = args.tag or f"eps{args.eps:g}_{_random_run_tag()}"
-    if args.smoke:
-        tag = f"smoke_{tag}"
-    run_e2(cfg, tag=tag, mirror_figs=not args.no_mirror,
-           uncontrolled_only=args.uncontrolled_only)
+    main()
diff --git a/graphon_control/__init__.py b/graphon_control/__init__.py
new file mode 100644
index 0000000..5e8d8c8
--- /dev/null
+++ b/graphon_control/__init__.py
@@ -0,0 +1,15 @@
+"""Differentiable time-and-label graphon control, following the paper's conventions.
+
+W is row-stochastic; the graphon density on a uniform N-grid is N*W.
+Importing the numerical package does not load plotting or run experiments.
+"""
+from .controls import FourierFeatures, MLP, TimeLabelGraphonControl
+from .dynamics import ConstantKernel, mean_field_kernel, precompute_Ws, rollout_with_Ws
+from .objectives import paper_cost, trajectory_cost
+from .diagnostics import vxi_discrete
+from .training import train_controller
+
+__all__ = ['FourierFeatures', 'MLP', 'TimeLabelGraphonControl', 'ConstantKernel',
+           'mean_field_kernel', 'precompute_Ws', 'rollout_with_Ws', 'paper_cost',
+           'trajectory_cost', 'vxi_discrete', 'train_controller']
+__version__ = '0.1.0'
diff --git a/graphon_control/artifacts.py b/graphon_control/artifacts.py
new file mode 100644
index 0000000..57cc0d1
--- /dev/null
+++ b/graphon_control/artifacts.py
@@ -0,0 +1,278 @@
+"""Immutable run directories, provenance, checkpoints, and explicit figure exports.
+
+Paths in a figure manifest are relative to the manifest file, so collection is
+independent of the caller's working directory. No experiment directory or
+shared figure directory is inferred from numerical parameter values.
+"""
+
+from __future__ import annotations
+
+import argparse
+from collections.abc import Mapping
+import dataclasses
+from datetime import datetime, timezone
+import hashlib
+from importlib import metadata
+import json
+import os
+from pathlib import Path
+import platform
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+import uuid
+
+
+SCHEMA_VERSION = 1
+_SAFE_COMPONENT = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]*\Z")
+
+
+def _safe_component(value: str, description: str) -> str:
+    if not isinstance(value, str) or not _SAFE_COMPONENT.fullmatch(value):
+        raise ValueError(f"{description} must be one safe path component: {value!r}")
+    if value in {".", ".."}:
+        raise ValueError(f"invalid {description}: {value!r}")
+    return value
+
+
+def generated_run_tag() -> str:
+    """Generate a collision-resistant tag without consulting experiment data."""
+    return uuid.uuid4().hex
+
+
+def create_run_directory(out_root, prefix: str, tag: str | None = None) -> Path:
+    """Reserve a new ``<prefix>_<tag>`` directory; existing runs are never reused."""
+    prefix = _safe_component(prefix, "run prefix")
+    tag = generated_run_tag() if tag is None else _safe_component(tag, "run tag")
+    root = Path(out_root).expanduser()
+    root.mkdir(parents=True, exist_ok=True)
+    run_dir = root / f"{prefix}_{tag}"
+    try:
+        run_dir.mkdir()
+    except FileExistsError as exc:
+        raise FileExistsError(f"Run already exists: {run_dir}. Choose a new tag.") from exc
+    return run_dir
+
+
+def atomic_write_json(path, data, *, overwrite: bool = False) -> Path:
+    """Write strict JSON atomically, rejecting accidental replacement by default."""
+    target = Path(path)
+    payload = json.dumps(data, indent=2, sort_keys=True, allow_nan=False) + "\n"
+    target.parent.mkdir(parents=True, exist_ok=True)
+    fd, temporary = tempfile.mkstemp(prefix=f".{target.name}.", dir=target.parent)
+    try:
+        with os.fdopen(fd, "w", encoding="utf-8") as stream:
+            stream.write(payload)
+            stream.flush()
+            os.fsync(stream.fileno())
+        if overwrite:
+            os.replace(temporary, target)
+        else:
+            # Linking publishes the complete file and rejects a concurrent writer.
+            os.link(temporary, target)
+        return target
+    finally:
+        Path(temporary).unlink(missing_ok=True)
+
+
+def _code_provenance() -> dict:
+    root = Path(__file__).resolve().parents[1]
+    try:
+        revision = subprocess.run(
+            ["git", "rev-parse", "HEAD"], cwd=root, check=True,
+            capture_output=True, text=True,
+        ).stdout.strip()
+        status = subprocess.run(
+            ["git", "status", "--porcelain"], cwd=root, check=True,
+            capture_output=True, text=True,
+        ).stdout.splitlines()
+        return {"commit": revision, "dirty": bool(status), "status": status}
+    except (OSError, subprocess.CalledProcessError):
+        return {"commit": None, "dirty": None, "status": None}
+
+
+def write_run_provenance(
+    run_dir, config, training_seed: int, evaluation_seed: int, *, device=None,
+) -> Path:
+    """Save the complete configuration and environment before training begins."""
+    if dataclasses.is_dataclass(config) and not isinstance(config, type):
+        config = dataclasses.asdict(config)
+    if not isinstance(config, Mapping):
+        raise TypeError("config must be a dataclass instance or mapping")
+    config = dict(config)
+    encoded = json.dumps(config, sort_keys=True, allow_nan=False).encode("utf-8")
+    versions = {}
+    for package in ("torch", "numpy", "matplotlib", "tqdm", "pillow"):
+        try:
+            versions[package] = metadata.version(package)
+        except metadata.PackageNotFoundError:
+            versions[package] = None
+    requested_device = str(device if device is not None else config.get("device", "cpu"))
+    import torch
+
+    provenance = {
+        "schema_version": SCHEMA_VERSION,
+        "created_at": datetime.now(timezone.utc).isoformat(),
+        "code": _code_provenance(),
+        "config_sha256": hashlib.sha256(encoded).hexdigest(),
+        "training_seed": training_seed,
+        "evaluation_seed": evaluation_seed,
+        "python": sys.version,
+        "platform": platform.platform(),
+        "libraries": versions,
+        "device": {
+            "requested": requested_device,
+            "cuda_build": torch.version.cuda,
+            "cuda_available": torch.cuda.is_available(),
+            "mps_available": bool(
+                hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+            ),
+        },
+    }
+    directory = Path(run_dir)
+    if not directory.is_dir():
+        raise FileNotFoundError(f"Reserve the run directory first: {directory}")
+    # Validate both payloads before publishing either one.
+    json.dumps(provenance, allow_nan=False)
+    for name in ("config.json", "provenance.json"):
+        if (directory / name).exists():
+            raise FileExistsError(f"Run metadata already exists: {directory / name}")
+    atomic_write_json(directory / "config.json", config)
+    return atomic_write_json(directory / "provenance.json", provenance)
+
+
+def save_checkpoint(path, controller, *, metadata=None) -> Path:
+    """Atomically save a raw state_dict, preserving existing checkpoint readers."""
+    import torch
+
+    target = Path(path)
+    if target.exists():
+        raise FileExistsError(f"Checkpoint already exists: {target}")
+    sidecar = target.with_suffix(target.suffix + ".metadata.json")
+    if metadata is not None:
+        if not isinstance(metadata, Mapping):
+            raise TypeError("checkpoint metadata must be a mapping")
+        json.dumps(metadata, allow_nan=False)
+        if sidecar.exists():
+            raise FileExistsError(f"Checkpoint metadata already exists: {sidecar}")
+    target.parent.mkdir(parents=True, exist_ok=True)
+    fd, temporary = tempfile.mkstemp(prefix=f".{target.name}.", dir=target.parent)
+    os.close(fd)
+    try:
+        torch.save(controller.state_dict(), temporary)
+        os.link(temporary, target)
+    finally:
+        Path(temporary).unlink(missing_ok=True)
+    if metadata is not None:
+        atomic_write_json(sidecar, {**metadata, "schema_version": SCHEMA_VERSION})
+    return target
+
+
+def load_controller(path, controller, *, map_location="cpu", strict: bool = True) -> dict:
+    """Load raw state_dicts and schema-wrapped checkpoints into an existing model."""
+    import torch
+
+    target = Path(path)
+    state = torch.load(target, map_location=map_location, weights_only=True)
+    details = {}
+    if isinstance(state, Mapping) and "state_dict" in state:
+        if state.get("schema_version", SCHEMA_VERSION) != SCHEMA_VERSION:
+            raise ValueError("Unsupported checkpoint schema_version")
+        details = dict(state.get("metadata", {}))
+        state = state["state_dict"]
+    controller.load_state_dict(state, strict=strict)
+    sidecar = target.with_suffix(target.suffix + ".metadata.json")
+    if sidecar.is_file():
+        with sidecar.open(encoding="utf-8") as stream:
+            sidecar_data = json.load(stream)
+        if not isinstance(sidecar_data, dict) or sidecar_data.get("schema_version") != SCHEMA_VERSION:
+            raise ValueError("Unsupported checkpoint metadata schema_version")
+        details.update(sidecar_data)
+    return details
+
+
+def collect_paper_figures(manifest_path, output_dir=None, *, overwrite=False) -> list[Path]:
+    """Copy an explicitly selected set of figures after preflighting every input.
+
+    Example manifest::
+
+        {"output_dir": "figures/paper", "figures": {
+          "kur_winding.pdf": "out_1/e1_selected/winding.pdf",
+          "graphon-N.pdf": "schematics/graphon-N.pdf"
+        }}
+
+    Replacing existing exports requires ``overwrite=True``. Source runs remain
+    untouched. Relative output overrides are relative to the current directory.
+    """
+    manifest_path = Path(manifest_path).expanduser().resolve()
+    with manifest_path.open(encoding="utf-8") as stream:
+        manifest = json.load(stream)
+    if not isinstance(manifest, dict):
+        raise ValueError("Manifest must be a JSON object")
+    figures = manifest.get("figures")
+    if not isinstance(figures, dict) or not figures:
+        raise ValueError("Manifest must contain a nonempty 'figures' mapping")
+    if output_dir is None:
+        configured = manifest.get("output_dir", "figures/paper")
+        if not isinstance(configured, str) or not configured:
+            raise ValueError("Manifest output_dir must be a nonempty path")
+        directory = (manifest_path.parent / configured).resolve()
+    else:
+        directory = Path(output_dir).expanduser().resolve()
+    sources = []
+    errors = []
+    if directory.exists() and not directory.is_dir():
+        errors.append(f"Output path is not a directory: {directory}")
+    for name, source_value in figures.items():
+        _safe_component(name, "output filename")
+        if not isinstance(source_value, str) or not source_value:
+            raise ValueError(f"Source for {name!r} must be a nonempty path")
+        source = (manifest_path.parent / source_value).resolve()
+        target = directory / name
+        if not source.is_file():
+            errors.append(f"Missing figure: {source}")
+        if source == target:
+            errors.append(f"Source and destination are identical: {source}")
+        if target.exists() and not target.is_file():
+            errors.append(f"Output filename is not a regular file: {target}")
+        if (target.exists() or target.is_symlink()) and not overwrite:
+            errors.append(f"Export exists (pass --overwrite to replace): {target}")
+        sources.append((source, target))
+    if errors:
+        raise ValueError("Figure collection failed before writing:\n" + "\n".join(errors))
+
+    directory.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.TemporaryDirectory(prefix=".paper-figures-", dir=directory.parent) as stage:
+        # Copy all inputs before creating or modifying the destination directory.
+        for source, target in sources:
+            shutil.copy2(source, Path(stage) / target.name)
+        directory.mkdir(exist_ok=True)
+        for _, target in sources:
+            staged = Path(stage) / target.name
+            if overwrite:
+                os.replace(staged, target)
+            else:
+                os.link(staged, target)
+    return [target for _, target in sources]
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(description=__doc__)
+    commands = parser.add_subparsers(dest="command", required=True)
+    collector = commands.add_parser("collect", help="collect figures from an explicit manifest")
+    collector.add_argument("manifest", type=Path)
+    collector.add_argument("--output-dir", type=Path)
+    collector.add_argument("--overwrite", action="store_true")
+    args = parser.parse_args(argv)
+    try:
+        paths = collect_paper_figures(args.manifest, args.output_dir, overwrite=args.overwrite)
+    except (OSError, ValueError) as exc:
+        parser.exit(1, f"{exc}\n")
+    for path in paths:
+        print(path)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/graphon_control/cli.py b/graphon_control/cli.py
new file mode 100644
index 0000000..e496d77
--- /dev/null
+++ b/graphon_control/cli.py
@@ -0,0 +1,54 @@
+"""Shared CLI configuration loading; numerical defaults live in dataclasses."""
+import dataclasses
+import json
+from pathlib import Path
+
+
+def add_common_arguments(parser):
+    parser.add_argument('--config', type=Path, help='load a saved config.json; explicit flags override it')
+    for flag, dest, kind in [
+        ('N', 'N', int), ('T', 'T', float), ('dt', 'dt', float),
+        ('steps', 'num_steps', int), ('batch', 'batch_size', int),
+        ('beta', 'beta', float), ('gamma-target', 'gamma_target', float),
+        ('x-target', 'x_target', float), ('seed', 'seed', int),
+        ('evaluation-seed', 'evaluation_seed', int), ('score-bound', 'score_bound', float),
+        ('sigma', 'sigma', float),
+    ]:
+        parser.add_argument(f'--{flag}', dest=dest, type=kind)
+    parser.add_argument('--device', choices=['cpu', 'mps', 'cuda'])
+    parser.add_argument('--tag', default='', help='new run tag; existing directories are refused')
+    parser.add_argument('--out-root', default='out_1')
+    parser.add_argument('--fig-root', default='figures')
+    parser.add_argument('--smoke', action='store_true', help='20 iterations on a small network and grid')
+    parser.add_argument('--no-plots', action='store_true', help='save numerical artifacts without rendering figures')
+    parser.add_argument('--replot', type=Path, help='render an archived run without retraining')
+    mirror = parser.add_mutually_exclusive_group()
+    mirror.add_argument('--mirror-figures', action='store_true', help='copy plots to a new directory named for this run')
+    mirror.add_argument('--no-mirror', action='store_false', dest='mirror_figures', help='compatibility flag; mirroring is off by default')
+    parser.set_defaults(mirror_figures=False)
+
+
+def configuration_from_args(config_class, args):
+    values = {}
+    if args.config is not None:
+        try:
+            values = json.loads(args.config.read_text())
+        except (OSError, json.JSONDecodeError) as exc:
+            raise ValueError(f'cannot read configuration: {exc}') from exc
+        if not isinstance(values, dict):
+            raise ValueError('configuration must be a JSON object')
+    fields = {field.name for field in dataclasses.fields(config_class)}
+    unknown = set(values) - fields
+    if unknown:
+        raise ValueError(f'unknown configuration fields: {sorted(unknown)}')
+    for name in fields:
+        value = getattr(args, name, None)
+        if value is not None:
+            values[name] = value
+    if args.smoke:
+        defaults = config_class()
+        values['num_steps'] = 20
+        for name, cap in [('N', 16), ('num_fourier_freqs', 2), ('embed_dim', 8),
+                          ('hidden_embed', 16), ('hidden_score', 16)]:
+            values[name] = min(values.get(name, getattr(defaults, name)), cap)
+    return config_class(**values)
diff --git a/graphon_control/controls.py b/graphon_control/controls.py
new file mode 100644
index 0000000..9335c93
--- /dev/null
+++ b/graphon_control/controls.py
@@ -0,0 +1,200 @@
+"""Time-and-label neural interaction controls, independent of experiment plotting.
+
+The directed controller implements the manuscript's row-softmax parameterization:
+``W[i, j]`` is the weight that receiver i assigns sender j.  The quadrature
+factor is already included in W, so the associated graphon density is ``N * W``.
+Self-interactions are permitted.  No prescribed receiver-variation bound is
+enforced by this module.
+"""
+
+from __future__ import annotations
+
+import math
+from numbers import Real
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def _positive_int(name: str, value: int) -> None:
+    if isinstance(value, bool) or not isinstance(value, int) or value < 1:
+        raise ValueError(f"{name} must be a positive integer")
+
+
+def _finite_scalar(name: str, value: Real, *, positive: bool = False) -> float:
+    if isinstance(value, bool) or not isinstance(value, Real):
+        raise ValueError(f"{name} must be a finite real number")
+    value = float(value)
+    if not math.isfinite(value) or (positive and value <= 0):
+        requirement = "finite and positive" if positive else "finite"
+        raise ValueError(f"{name} must be {requirement}")
+    return value
+
+
+class FourierFeatures(nn.Module):
+    """Dyadic sine/cosine features, with all sine terms followed by cosine terms."""
+
+    def __init__(self, num_freqs: int):
+        super().__init__()
+        _positive_int("num_freqs", num_freqs)
+        self.num_freqs = num_freqs
+        freqs = 2.0 ** torch.arange(num_freqs, dtype=torch.float32)
+        if not torch.isfinite(freqs).all():
+            raise ValueError("num_freqs produces nonfinite Fourier frequencies")
+        self.register_buffer("freqs", freqs)
+
+    @property
+    def out_dim(self) -> int:
+        return 2 * self.num_freqs
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Map labels of shape ``(...)`` to features of shape ``(..., 2*m)``."""
+        angles = 2.0 * math.pi * x.unsqueeze(-1) * self.freqs
+        return torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1)
+
+
+class MLP(nn.Module):
+    """A ReLU multilayer perceptron with a linear output layer."""
+
+    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, depth: int = 2):
+        super().__init__()
+        for name, value in (("in_dim", in_dim), ("hidden_dim", hidden_dim), ("out_dim", out_dim)):
+            _positive_int(name, value)
+        if isinstance(depth, bool) or not isinstance(depth, int) or depth < 0:
+            raise ValueError("depth must be a nonnegative integer")
+        layers = []
+        d = in_dim
+        for _ in range(depth):
+            layers.extend((nn.Linear(d, hidden_dim), nn.ReLU()))
+            d = hidden_dim
+        layers.append(nn.Linear(d, out_dim))
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class TimeLabelGraphonControl(nn.Module):
+    """Learn an open-loop interaction schedule from time and fixed agent labels.
+
+    ``xi`` is a nonempty, nondecreasing floating tensor in [0, 1].  Model
+    parameters and buffers inherit its device and dtype.  A scalar time returns
+    ``(N, N)``; a time vector of length B returns ``(B, N, N)``, including B=1.
+
+    Directed mode is row-stochastic, as in manuscript equations (64)--(67).
+    Optional ``score_bound=S`` uses ``S*tanh(raw_score)`` and guarantees
+    ``N*W[i,j] <= exp(2*S)``.  The default leaves scores unconstrained, preserving
+    existing experiment behavior.  Neither setting imposes a prescribed BV bound.
+
+    The legacy symmetric alternative uses one shared embedding, symmetric
+    positive weights, and mean-row normalization.  Its individual row sums may
+    differ from one, so it lies outside the manuscript's row-stochastic class.
+    ``positive_map`` affects only that alternative.
+
+    Module/buffer names and parameter ordering match legacy checkpoints from
+    ``model_1.TimeLabelGraphonControl``.  ``score_bound`` is configuration, not a
+    learned parameter; restore it alongside other constructor settings.
+    """
+
+    def __init__(
+        self,
+        xi: torch.Tensor,
+        num_fourier_freqs: int,
+        embed_dim: int,
+        hidden_embed: int,
+        hidden_score: int,
+        time_scale: float = 1.0,
+        enforce_symmetry: bool = False,
+        positive_map: str = "softplus",
+        score_bound: float | None = None,
+    ):
+        super().__init__()
+        if not isinstance(xi, torch.Tensor) or not xi.is_floating_point():
+            raise ValueError("xi must be a floating-point tensor")
+        if xi.ndim != 1 or xi.numel() == 0:
+            raise ValueError("xi must be a nonempty one-dimensional label grid")
+        if not torch.isfinite(xi).all() or torch.any((xi < 0) | (xi > 1)):
+            raise ValueError("xi must contain finite labels in [0, 1]")
+        if torch.any(xi[1:] < xi[:-1]):
+            raise ValueError("xi must be sorted in nondecreasing order")
+        for name, value in (("num_fourier_freqs", num_fourier_freqs), ("embed_dim", embed_dim),
+                            ("hidden_embed", hidden_embed), ("hidden_score", hidden_score)):
+            _positive_int(name, value)
+        self.time_scale = _finite_scalar("time_scale", time_scale, positive=True)
+        if not isinstance(enforce_symmetry, bool):
+            raise ValueError("enforce_symmetry must be a boolean")
+        if positive_map not in ("softplus", "exp"):
+            raise ValueError("positive_map must be 'softplus' or 'exp'")
+        if score_bound is not None:
+            score_bound = _finite_scalar("score_bound", score_bound)
+            if score_bound < 0:
+                raise ValueError("score_bound must be nonnegative")
+        self.score_bound = score_bound
+        self.enforce_symmetry = enforce_symmetry
+        self.positive_map = positive_map
+        # Labels are fixed model data; clone so callers cannot invalidate cached features.
+        self.register_buffer("xi", xi.detach().clone())
+        self.N = xi.numel()
+
+        self.ff = FourierFeatures(num_fourier_freqs)
+        ff_dim = self.ff.out_dim
+        # Preserve legacy construction order for seeded runs and optimizer states.
+        self.embed_net = MLP(ff_dim, hidden_embed, embed_dim) if enforce_symmetry else None
+        self.receiver_net = None if enforce_symmetry else MLP(ff_dim, hidden_embed, embed_dim)
+        self.sender_net = None if enforce_symmetry else MLP(ff_dim, hidden_embed, embed_dim)
+        self.score_net = MLP(1 + 2 * embed_dim, hidden_score, 1)
+
+        # Move frequencies before evaluating labels: xi may already be on CUDA/MPS.
+        self.to(device=xi.device, dtype=xi.dtype)
+        if not torch.isfinite(self.ff.freqs).all():
+            raise ValueError("Fourier frequencies are not representable in xi.dtype")
+        features = self.ff(self.xi)
+        if not torch.isfinite(features).all():
+            raise ValueError("Fourier features are not representable in xi.dtype")
+        self.register_buffer("psi_xi_cached", features)
+
+    def forward(self, t: torch.Tensor | float) -> torch.Tensor:
+        """Evaluate scalar or one-dimensional finite times on the fixed grid."""
+        if isinstance(t, torch.Tensor) and t.is_complex():
+            raise ValueError("t must contain real finite times")
+        t = torch.as_tensor(t, device=self.xi.device, dtype=self.xi.dtype)
+        if t.ndim > 1 or t.numel() == 0:
+            raise ValueError("t must be a scalar or a nonempty one-dimensional vector")
+        if not torch.isfinite(t).all():
+            raise ValueError("t must contain finite times")
+        scalar = t.ndim == 0
+        t = t.reshape(-1)
+        B, N = t.numel(), self.N
+        t_in = (t / self.time_scale).reshape(B, 1, 1, 1).expand(B, N, N, 1)
+        if not torch.isfinite(t_in).all():
+            raise ValueError("scaled times must be finite in the controller dtype")
+
+        if self.enforce_symmetry:
+            h_receiver = h_sender = self.embed_net(self.psi_xi_cached)
+        else:
+            h_receiver = self.receiver_net(self.psi_xi_cached)
+            h_sender = self.sender_net(self.psi_xi_cached)
+        d = h_receiver.shape[-1]
+        h_i = h_receiver.view(1, N, 1, d).expand(B, N, N, d)
+        h_j = h_sender.view(1, 1, N, d).expand(B, N, N, d)
+        scores = self.score_net(torch.cat((t_in, h_i, h_j), dim=-1)).squeeze(-1)
+        if self.enforce_symmetry:
+            scores = 0.5 * (scores + scores.transpose(-1, -2))
+        if self.score_bound is not None:
+            scores = self.score_bound * torch.tanh(scores)
+
+        if self.enforce_symmetry:
+            if self.positive_map == "exp":
+                log_weights = scores.clamp(max=20.0)
+                weights = torch.exp(log_weights - log_weights.amax(dim=(-1, -2), keepdim=True))
+            else:
+                weights = F.softplus(scores)
+            # Rescaling first avoids overflow in row sums; tiny floors avoid 0/0
+            # when softplus underflows.  Both operations preserve symmetry.
+            weights = weights.clamp_min(torch.finfo(weights.dtype).tiny)
+            weights = weights / weights.amax(dim=(-1, -2), keepdim=True)
+            W = weights / weights.sum(dim=-1).mean(dim=-1, keepdim=True).unsqueeze(-1)
+        else:
+            W = F.softmax(scores, dim=-1)
+        return W[0] if scalar else W
diff --git a/graphon_control/diagnostics.py b/graphon_control/diagnostics.py
new file mode 100644
index 0000000..1c3320b
--- /dev/null
+++ b/graphon_control/diagnostics.py
@@ -0,0 +1,33 @@
+"""Scientific diagnostics; measured variation is not an enforced BV budget."""
+
+import math
+
+import torch
+
+
+def _check_kernel(W: torch.Tensor) -> None:
+    if W.ndim < 2 or W.shape[-2] != W.shape[-1] or W.shape[-1] == 0:
+        raise ValueError("kernel must have shape (..., N, N), N >= 1")
+
+
+def vxi_discrete(W: torch.Tensor) -> torch.Tensor:
+    """Receiver variation of graphon values wbar=N*W, preserving batch axes.
+
+    V_xi = sum_i sqrt((1/N) sum_j (wbar[i+1,j]-wbar[i,j])²).
+    The receiver interval has no periodic closing edge.
+    """
+    _check_kernel(W)
+    N = W.shape[-1]
+    differences = N * (W[..., 1:, :] - W[..., :-1, :])
+    return torch.linalg.vector_norm(differences, dim=-1).sum(dim=-1) / math.sqrt(N)
+
+
+def kernel_constraint_report(W: torch.Tensor) -> dict[str, torch.Tensor]:
+    """Report constraints and variation without asserting admissibility."""
+    _check_kernel(W)
+    return {
+        "min_weight": W.amin(dim=(-2, -1)),
+        "max_row_sum_error": (W.sum(-1) - 1).abs().amax(dim=-1),
+        "max_symmetry_error": (W - W.transpose(-1, -2)).abs().amax(dim=(-2, -1)),
+        "receiver_variation": vxi_discrete(W),
+    }
diff --git a/graphon_control/dynamics.py b/graphon_control/dynamics.py
new file mode 100644
index 0000000..7cbc9c3
--- /dev/null
+++ b/graphon_control/dynamics.py
@@ -0,0 +1,125 @@
+"""Differentiable finite-agent Euler and Euler–Maruyama dynamics.
+
+W contains quadrature weights: W_ij = wbar_ij/N. Consequently the drift is
+sum_j W_ij * phi(x_j - x_i), with no additional factor of 1/N.
+"""
+
+import math
+from typing import Callable
+
+import torch
+
+from .validation import finite_scalar, n_time_steps, positive_integer
+
+
+def mean_field_kernel(N: int, device=None, dtype=torch.float32) -> torch.Tensor:
+    """The constant graphon w=1, represented by W_ij=1/N."""
+    positive_integer(N, "N")
+    return torch.full((N, N), 1.0 / N, device=device, dtype=dtype)
+
+
+class ConstantKernel:
+    """Expose a fixed baseline through the controller's scalar/batched API."""
+
+    def __init__(self, W0: torch.Tensor):
+        if W0.ndim != 2 or W0.shape[0] != W0.shape[1] or W0.shape[0] == 0:
+            raise ValueError("W0 must have nonempty square shape (N, N)")
+        if not W0.is_floating_point() or not torch.isfinite(W0).all():
+            raise ValueError("W0 must contain finite floating point weights")
+        self.W0 = W0
+
+    def __call__(self, t: torch.Tensor) -> torch.Tensor:
+        if t.ndim == 0:
+            return self.W0
+        if t.ndim != 1:
+            raise ValueError("time must be scalar or one-dimensional")
+        return self.W0.unsqueeze(0).expand(t.shape[0], *self.W0.shape)
+
+
+def _model_dtype(model):
+    for name in ("xi", "W0"):
+        value = getattr(model, name, None)
+        if isinstance(value, torch.Tensor) and value.is_floating_point():
+            return value.dtype
+    if hasattr(model, "parameters"):
+        for parameter in model.parameters():
+            if parameter.is_floating_point():
+                return parameter.dtype
+    return torch.get_default_dtype()
+
+
+def precompute_Ws(model, NT: int, dt: float, device, chunk: int) -> torch.Tensor:
+    """Sample W(t_n) in chunks, preserving its gradient and (NT,N,N) shape."""
+    positive_integer(NT, "NT")
+    positive_integer(chunk, "chunk")
+    finite_scalar(dt, "dt", positive=True)
+    times = torch.arange(NT, device=device, dtype=_model_dtype(model)) * dt
+    outputs = []
+    for start in range(0, NT, chunk):
+        batch = times[start:start + chunk]
+        W = model(batch)
+        # Compatibility with older controllers that squeeze one-element batches.
+        if W.ndim == 2 and batch.numel() == 1:
+            W = W.unsqueeze(0)
+        if W.ndim != 3 or W.shape[0] != batch.numel() or W.shape[-2] != W.shape[-1] or W.shape[-1] == 0:
+            raise ValueError("controller must return (batch, N, N) for batched times")
+        outputs.append(W)
+    return torch.cat(outputs, dim=0)
+
+
+def rollout_with_Ws(
+    Ws: torch.Tensor,
+    x0: torch.Tensor,
+    dt: float,
+    phi: Callable[[torch.Tensor], torch.Tensor],
+    sigma: float = 0.0,
+    generator: torch.Generator | None = None,
+    brownian_increments: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Integrate a batch of states; return shape (NT+1,B,N).
+
+    Optional ``brownian_increments`` has shape (NT,B,N) and contains actual
+    Brownian increments ΔB (variance dt), rather than standardized normals.
+    It permits paired stochastic comparisons and exact replay of a rollout.
+    A generator controls automatic draws and must match the state device.
+    """
+    dt = finite_scalar(dt, "dt", positive=True)
+    sigma = finite_scalar(sigma, "sigma", nonnegative=True)
+    if x0.ndim != 2 or min(x0.shape) == 0 or not x0.is_floating_point():
+        raise ValueError("x0 must have nonempty floating point shape (B, N)")
+    B, N = x0.shape
+    if Ws.ndim != 3 or Ws.shape[0] == 0 or Ws.shape[1:] != (N, N):
+        raise ValueError("Ws must have shape (NT, N, N) with NT >= 1")
+    if Ws.device != x0.device or Ws.dtype != x0.dtype:
+        raise ValueError("Ws and x0 must have the same device and dtype")
+    if brownian_increments is not None:
+        if generator is not None:
+            raise ValueError("supply brownian_increments or generator, not both")
+        if brownian_increments.shape != (Ws.shape[0], B, N):
+            raise ValueError("brownian_increments must have shape (NT, B, N)")
+        if brownian_increments.device != x0.device or brownian_increments.dtype != x0.dtype:
+            raise ValueError("brownian_increments must match x0 device and dtype")
+    x = x0
+    xs = [x]
+    for n, W in enumerate(Ws):
+        # diff[b,i,j] = x_j - x_i: receiver is the row, sender is the column.
+        diff = x.unsqueeze(1) - x.unsqueeze(2)
+        drift = (W.unsqueeze(0) * phi(diff)).sum(dim=-1)
+        x = x + dt * drift
+        if sigma:
+            dB = (brownian_increments[n] if brownian_increments is not None else
+                  math.sqrt(dt) * torch.randn(x.shape, device=x.device, dtype=x.dtype,
+                                             generator=generator))
+            x = x + sigma * dB
+        xs.append(x)
+    return torch.stack(xs)
+
+
+def rollout_dynamics_w1(model, x0, T, dt, sigma, phi, *, generator=None,
+                        brownian_increments=None):
+    """Compatibility adapter returning lists of states and kernels."""
+    NT = n_time_steps(T, dt)
+    # Keep the generic API's scalar-time evaluations to support custom models.
+    Ws = torch.stack([model(x0.new_tensor(n * dt)) for n in range(NT)])
+    xs = rollout_with_Ws(Ws, x0, dt, phi, sigma, generator, brownian_increments)
+    return list(xs.unbind(0)), list(Ws.unbind(0))
diff --git a/graphon_control/experiments/__init__.py b/graphon_control/experiments/__init__.py
new file mode 100644
index 0000000..8ff131c
--- /dev/null
+++ b/graphon_control/experiments/__init__.py
@@ -0,0 +1 @@
+"""Experiment-specific physics and orchestration; import an individual module to use it."""
diff --git a/graphon_control/experiments/bounded_confidence.py b/graphon_control/experiments/bounded_confidence.py
new file mode 100644
index 0000000..0c58ed9
--- /dev/null
+++ b/graphon_control/experiments/bounded_confidence.py
@@ -0,0 +1,537 @@
+"""Bounded-confidence opinions experiment using the shared paper conventions.
+
+See docs/NUMERICS.md for equations and docs/DEVELOPMENT.md for reproducible runs.
+Model-specific rendering is in bounded_confidence_plots.py.
+"""
+import argparse
+import datetime
+import json
+import math
+import os
+from dataclasses import dataclass, replace
+from typing import Callable, Dict, List, Optional
+
+import numpy as np
+import torch
+
+from ..controls import TimeLabelGraphonControl
+from ..artifacts import create_run_directory, write_run_provenance, save_checkpoint
+from ..dynamics import ConstantKernel, mean_field_kernel, precompute_Ws, rollout_with_Ws
+from ..objectives import paper_cost
+from ..diagnostics import vxi_discrete
+from ..training import train_controller
+from ..validation import n_time_steps, validate_common_config
+from ..cli import add_common_arguments, configuration_from_args
+from pathlib import Path
+
+
+# ============================================================
+# Configuration (paper protocol; deviations stated in subsec:exp-bc)
+# ============================================================
+
+@dataclass
+class E2Config:
+    # interaction
+    eps: float = 1.0           # confidence radius
+    gain: float = 4.0          # gain a (Lipschitz constant of φ_ε)
+    phi_form: str = "poly"     # "poly" (compact support) | "gauss" (fallback)
+
+    # initial profile x(0, ξ) = span (ξ - 1/2) + amp sin(2π modes ξ)
+    span: float = 6.0
+    amp: float = 0.5
+    modes: int = 4
+    jitter: float = 0.1        # i.i.d. jitter on the training initial data
+
+    # label-band baseline
+    band_range: float = 0.10
+
+    # discretization
+    N: int = 64
+    T: float = 10.0
+    dt: float = 0.05
+    sigma: float = 0.0
+
+    # cost — PAPER convention (eqn:J-cost-1), with β = 1 (stated deviation from
+    # the e1 protocol's β = 10: merging clusters needs concentrated attention,
+    # and the pilot decides the final value; see --beta)
+    beta: float = 1.0
+    lambda_terminal: float = 0.0
+    gamma_target: float = 100.0
+    x_target: float = 0.0
+
+    # optimization
+    batch_size: int = 16
+    num_steps: int = 2000
+    lr: float = 1e-3
+    grad_clip: float = 1.0
+
+    # architecture (model_1.py defaults, as in e1)
+    num_fourier_freqs: int = 8
+    embed_dim: int = 64
+    hidden_embed: int = 128
+    hidden_score: int = 128
+
+    # bookkeeping
+    evaluation_seed: int = 1729
+    score_bound: Optional[float] = None
+    seed: int = 0
+    device: str = "cpu"
+    time_chunk: int = 50
+    uncontrolled_only: bool = False
+    hk_ref: bool = False       # add the sharp state-dependent HK reference arm
+
+    def __post_init__(self):
+        validate_common_config(self)
+        if not isinstance(self.uncontrolled_only, bool) or not isinstance(self.hk_ref, bool):
+            raise ValueError("uncontrolled_only and hk_ref must be booleans")
+        for name in ("eps", "gain", "span", "band_range"):
+            if not math.isfinite(getattr(self, name)) or getattr(self, name) <= 0:
+                raise ValueError(f"{name} must be finite and positive")
+        if self.phi_form not in ("poly", "gauss"):
+            raise ValueError("phi_form must be 'poly' or 'gauss'")
+        if not math.isfinite(self.amp) or isinstance(self.modes, bool) or not isinstance(self.modes, int) or self.modes < 1:
+            raise ValueError("amp must be finite and modes a positive integer")
+
+    @property
+    def NT(self) -> int:
+        return n_time_steps(self.T, self.dt)
+
+
+# ============================================================
+# E2 ingredients
+# ============================================================
+
+def phi_bc_poly(y: torch.Tensor, eps: float, gain: float) -> torch.Tensor:
+    """
+    Mollified bounded-confidence interaction
+      φ_ε(z) = a z (1 - z²/ε²)₊²,
+    C^{1,1}, odd, Lipschitz constant a, support exactly [-ε, ε].
+    """
+    u = y / eps
+    return gain * y * torch.clamp(1.0 - u * u, min=0.0) ** 2
+
+
+def phi_bc_gauss(y: torch.Tensor, eps: float, gain: float) -> torch.Tensor:
+    """Gaussian-tail fallback φ_ε(z) = a z exp(-z²/ε²) (no dead gradients)."""
+    return gain * y * torch.exp(-(y / eps) ** 2)
+
+
+def make_phi(cfg: E2Config) -> Callable[[torch.Tensor], torch.Tensor]:
+    base = {"poly": phi_bc_poly, "gauss": phi_bc_gauss}[cfg.phi_form]
+    return lambda y: base(y, cfg.eps, cfg.gain)
+
+
+def bc_profile(xi: torch.Tensor, span: float, amp: float, modes: int) -> torch.Tensor:
+    """Initial opinion profile x(0, ξ) = span (ξ - 1/2) + amp sin(2π modes ξ)."""
+    return span * (xi - 0.5) + amp * torch.sin(2.0 * math.pi * modes * xi)
+
+
+def band_kernel(xi: torch.Tensor, r: float) -> torch.Tensor:
+    """
+    Label-band kernel: smooth bump of range r in |ξ - ζ| on the INTERVAL
+    (no torus wrap — opinions are not periodic; row normalization absorbs the
+    boundary), row-normalized to a row-stochastic W. Smooth in ξ, hence of
+    finite receiver variation, unlike the sharp band indicator.
+    """
+    d = (xi.unsqueeze(0) - xi.unsqueeze(1)).abs()
+    u = (d / r).clamp(max=1.0 - 1e-6)
+    bump = torch.where(d < r, torch.exp(-1.0 / (1.0 - u ** 2)), torch.zeros_like(d))
+    return bump / bump.sum(dim=1, keepdim=True)
+
+
+def rollout_hk_reference(x0: torch.Tensor, NT: int, dt: float, eps: float) -> torch.Tensor:
+    """
+    Classical sharp Hegselmann-Krause reference: state-dependent kernel
+    W_ij(t) ∝ 1{|x_i - x_j| ≤ ε}, row-normalized, with φ(z) = z within range.
+    NOT admissible in our control class (state-dependent, infinite V_ξ);
+    reported only as a reference portrait. x0: (B, N) -> (NT+1, B, N).
+    """
+    x = x0
+    xs = [x]
+    for _ in range(NT):
+        diff = x.unsqueeze(1) - x.unsqueeze(2)            # (B, N, N): x_j - x_i
+        adj = (diff.abs() <= eps).to(x.dtype)
+        W = adj / adj.sum(dim=-1, keepdim=True)           # diagonal is always in range
+        x = x + dt * (W * diff).sum(-1)
+        xs.append(x)
+    return torch.stack(xs, dim=0)
+
+
+def e2_cost(xs, Ws, cfg: E2Config):
+    """Paper objective, equation (60); shared across both experiments."""
+    return paper_cost(xs, Ws, cfg)
+
+
+def cluster_count(x: torch.Tensor, eps: float) -> torch.Tensor:
+    """
+    Number of opinion clusters: 1 + #{gaps > ε in the sorted state}. A gap
+    larger than ε separates components for the compact-support interaction.
+    For Gaussian tails this is only a descriptive gap threshold.
+    x: (..., N) -> (...) int.
+    """
+    xs, _ = torch.sort(x, dim=-1)
+    gaps = xs[..., 1:] - xs[..., :-1]
+    return 1 + (gaps > eps).sum(dim=-1)
+
+
+def cluster_inventory(x: torch.Tensor, eps: float) -> List[Dict[str, float]]:
+    """Cluster centers and mass fractions of a single state x: (N,)."""
+    xs, _ = torch.sort(x)
+    xs = xs.cpu().numpy()
+    breaks = np.where(np.diff(xs) > eps)[0]
+    out = []
+    start = 0
+    for b in list(breaks) + [len(xs) - 1]:
+        block = xs[start:b + 1]
+        out.append({"center": float(block.mean()), "mass": len(block) / len(xs)})
+        start = b + 1
+    return out
+
+
+def control_authority(x: torch.Tensor,
+                      phi: Callable[[torch.Tensor], torch.Tensor]) -> torch.Tensor:
+    """
+    Control authority (1/N) Σ_i max_j |φ_ε(x_j - x_i)|: the strongest force any
+    admissible kernel can still exert, averaged over receivers. Decays to 0 both
+    at fragmentation (window of controllability closed) and at consensus
+    (nothing left to do) — the two are distinguished by Var. x: (..., N) -> (...).
+    """
+    diff = x.unsqueeze(-2) - x.unsqueeze(-1)   # (..., N, N): [i, j] = x_j - x_i
+    return phi(diff).abs().amax(dim=-1).mean(dim=-1)
+
+
+# ============================================================
+# Evaluation helper (shared by full and uncontrolled-only runs)
+# ============================================================
+
+def evaluate_arms(arm_specs: Dict[str, object], x0_eval: torch.Tensor,
+                  cfg: E2Config, phi, device):
+    """
+    arm_specs values: a kernel (ConstantKernel / trained model), or the string
+    "hk" for the sharp state-dependent reference. Deterministic, exact datum.
+    """
+    out = {}
+    with torch.no_grad():
+        for arm, spec in arm_specs.items():
+            if spec == "hk":
+                xs_eval = rollout_hk_reference(x0_eval, cfg.NT, cfg.dt, cfg.eps)
+                Ws_eval, parts = None, None
+            else:
+                Ws_eval = precompute_Ws(spec, cfg.NT, cfg.dt, device, cfg.time_chunk)
+                xs_eval = rollout_with_Ws(Ws_eval, x0_eval, cfg.dt, phi, sigma=0.0)
+                parts = e2_cost(xs_eval, Ws_eval, cfg)
+            x_t = xs_eval[:, 0, :]                               # (NT+1, N)
+            out[arm] = {
+                "X": x_t.cpu().numpy(),
+                "Ws": Ws_eval,
+                "var": x_t.var(dim=1, unbiased=False).cpu().numpy(),
+                "K": cluster_count(x_t, cfg.eps).cpu().numpy(),
+                "authority": control_authority(x_t, phi).cpu().numpy(),
+                "cost": None if parts is None
+                        else {kk: float(vv) for kk, vv in parts.items()},
+            }
+    return out
+
+
+def freeze_time(K: np.ndarray, dt: float) -> Optional[float]:
+    """First time after which K(t) never changes again (None if K constant)."""
+    changes = np.where(K[1:] != K[:-1])[0]
+    if len(changes) == 0:
+        return None
+    return float((changes[-1] + 1) * dt)
+
+
+# ============================================================
+# Main experiment
+# ============================================================
+
+def run_e2(cfg: E2Config, tag: str = "", out_root: str = "out_1", fig_root: str = "figures",
+           mirror_figs: bool = False, uncontrolled_only: Optional[bool] = None, make_plots: bool = True):
+    torch.manual_seed(cfg.seed)
+    np.random.seed(cfg.seed)
+    device = torch.device(cfg.device)
+    dtype = torch.float32
+
+    uncontrolled_only = cfg.uncontrolled_only if uncontrolled_only is None else uncontrolled_only
+    cfg = replace(cfg, uncontrolled_only=uncontrolled_only)
+    cfg.__post_init__()
+    run_dir = str(create_run_directory(out_root, "e2", tag or None))
+    fig_dir = os.path.join(fig_root, os.path.basename(run_dir))
+    if mirror_figs and make_plots and Path(fig_dir).exists():
+        raise FileExistsError(f"Figure export already exists: {fig_dir}")
+    write_run_provenance(run_dir, cfg, cfg.seed, cfg.evaluation_seed, device=cfg.device)
+
+    xi = (torch.arange(cfg.N, device=device, dtype=dtype) + 0.5) / cfg.N
+    base = bc_profile(xi, cfg.span, cfg.amp, cfg.modes)
+    NT = cfg.NT
+    phi = make_phi(cfg)
+
+    print("=" * 64)
+    print(f"E2 bounded confidence | eps={cfg.eps} gain={cfg.gain} "
+          f"span={cfg.span} N={cfg.N} T={cfg.T} dt={cfg.dt}")
+    print(f"steps={cfg.num_steps} batch={cfg.batch_size} lr={cfg.lr} "
+          f"beta={cfg.beta} gamma={cfg.gamma_target} phi={cfg.phi_form}")
+    print(f"device={cfg.device} tag={tag}")
+    print("=" * 64)
+
+    # ---------- validate the paper-specific fragmentation scenario ----------
+    # (the narrative claim of subsec:exp-bc is false for this config otherwise)
+    with torch.no_grad():
+        W_mf = mean_field_kernel(cfg.N, device)
+        Ws_gate = W_mf.unsqueeze(0).expand(NT, -1, -1)
+        xs_gate = rollout_with_Ws(Ws_gate, base.unsqueeze(0), cfg.dt, phi)
+        K_gate = int(cluster_count(xs_gate[-1, 0], cfg.eps))
+    print(f"sanity gate: uncontrolled mean-field arm ends with K = {K_gate} clusters")
+    if K_gate < 2 and not uncontrolled_only and cfg.phi_form == "poly":
+        raise RuntimeError(
+            f"uncontrolled arm does not fragment (K={K_gate}); "
+            f"raise --span, lower --eps, or raise --gain")
+
+    arm_specs: Dict[str, object] = {
+        "meanfield": ConstantKernel(W_mf),
+        "band": ConstantKernel(band_kernel(xi, cfg.band_range)),
+    }
+
+    # ---------- training ----------
+    model = None
+    history: Dict[str, List[float]] = {kk: [] for kk in
+                                       ["loss", "state_cost", "l2_reg_cost",
+                                        "terminal_var", "terminal_target",
+                                        "terminal_mean"]}
+    if not uncontrolled_only:
+        model = TimeLabelGraphonControl(
+            xi=xi,
+            num_fourier_freqs=cfg.num_fourier_freqs,
+            embed_dim=cfg.embed_dim,
+            hidden_embed=cfg.hidden_embed,
+            hidden_score=cfg.hidden_score,
+            time_scale=max(cfg.T, 1e-8),
+            enforce_symmetry=False,
+            score_bound=cfg.score_bound,
+        ).to(device)
+        def objective():
+            Ws = precompute_Ws(model, NT, cfg.dt, device, cfg.time_chunk)
+            x0 = base.unsqueeze(0) + cfg.jitter * torch.randn(
+                cfg.batch_size, cfg.N, device=device, dtype=dtype)
+            xs = rollout_with_Ws(Ws, x0, cfg.dt, phi, sigma=cfg.sigma)
+            return e2_cost(xs, Ws, cfg)
+
+        history, _ = train_controller(model, objective, cfg.num_steps, cfg.lr,
+                                      cfg.grad_clip, description="E2 training")
+
+        model.eval()
+        arm_specs["trained"] = model
+
+    if cfg.hk_ref:
+        arm_specs["hk"] = "hk"
+
+    torch.manual_seed(cfg.evaluation_seed)
+
+    # ---------- evaluation: deterministic, exact datum ----------
+    x0_eval = base.unsqueeze(0)
+    results = evaluate_arms(arm_specs, x0_eval, cfg, phi, device)
+    arms = list(arm_specs.keys())
+    cost_arms = [a for a in arms if results[a]["cost"] is not None]
+
+    vxi_band = float(vxi_discrete(results["band"]["Ws"][0]))
+    vxi_mf = float(vxi_discrete(results["meanfield"]["Ws"][0]))
+    vxi_trained = (vxi_discrete(results["trained"]["Ws"]).cpu().numpy()
+                   if "trained" in results else None)
+
+    t_grid = np.arange(NT + 1) * cfg.dt
+    t_grid_W = np.arange(NT) * cfg.dt
+
+    # ---------- persistence ----------
+    if model is not None:
+        save_checkpoint(os.path.join(run_dir, "model.pt"), model)
+        with open(os.path.join(run_dir, "history.json"), "w") as f:
+            json.dump(history, f, indent=2)
+
+    npz = {"t": t_grid, "t_W": t_grid_W, "xi": xi.cpu().numpy(),
+           "W_band": results["band"]["Ws"][0].cpu().numpy().astype(np.float32)}
+    for arm in arms:
+        npz[f"x_{arm}"] = results[arm]["X"]
+        npz[f"var_{arm}"] = results[arm]["var"]
+        npz[f"K_{arm}"] = results[arm]["K"]
+        npz[f"authority_{arm}"] = results[arm]["authority"]
+    if "trained" in results:
+        npz["W_trained"] = results["trained"]["Ws"].cpu().numpy().astype(np.float32)
+        npz["vxi_trained"] = vxi_trained
+    np.savez_compressed(os.path.join(run_dir, "trajectories.npz"), **npz)
+
+    var_final = {arm: float(results[arm]["var"][-1]) for arm in arms}
+    K_final = {arm: int(results[arm]["K"][-1]) for arm in arms}
+    metrics = {
+        "eps": cfg.eps, "gain": cfg.gain, "span": cfg.span, "beta": cfg.beta,
+        "costs": {arm: results[arm]["cost"] for arm in cost_arms},
+        "var_final": var_final,
+        "K_final": K_final,
+        "clusters_meanfield_T": cluster_inventory(
+            torch.from_numpy(results["meanfield"]["X"][-1]), cfg.eps),
+        "freeze_time_meanfield": freeze_time(results["meanfield"]["K"], cfg.dt),
+        "vxi": {"band": vxi_band, "meanfield": vxi_mf},
+    }
+    if "trained" in results:
+        cons = np.where(results["trained"]["var"] < 1e-2)[0]
+        metrics["consensus_time_trained"] = (float(cons[0] * cfg.dt)
+                                             if len(cons) else None)
+        metrics["terminal_mean_trained"] = results["trained"]["cost"]["terminal_mean"]
+        metrics["vxi"]["trained_max"] = float(vxi_trained.max())
+        metrics["vxi"]["trained_final"] = float(vxi_trained[-1])
+        metrics["value_gap_meanfield_minus_trained"] = (
+            results["meanfield"]["cost"]["total"] - results["trained"]["cost"]["total"])
+    with open(os.path.join(run_dir, "metrics.json"), "w") as f:
+        json.dump(metrics, f, indent=2)
+
+    if make_plots:
+        plot_run(run_dir, fig_dir=fig_dir if mirror_figs else None)
+
+    # ---------- summary ----------
+    lines = [
+        "=== E2 bounded confidence: run summary ===",
+        f"timestamp : {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
+        f"run_dir   : {run_dir}",
+        f"eps = {cfg.eps}, gain = {cfg.gain}, span = {cfg.span}, "
+        f"N = {cfg.N}, T = {cfg.T}, dt = {cfg.dt}, steps = {cfg.num_steps}, "
+        f"beta = {cfg.beta}",
+        "",
+        "arm         J_total      Var(T)     K(T)   mean(T)",
+    ]
+    for arm in arms:
+        j = (f"{results[arm]['cost']['total']:>9.4f}"
+             if results[arm]["cost"] is not None else "        —")
+        mean_T = float(results[arm]["X"][-1].mean())
+        lines.append(f"{arm:<11s} {j}  {var_final[arm]:>10.6f}   "
+                     f"{K_final[arm]:>3d}   {mean_T:>+8.4f}")
+    lines.append("")
+    lines.append(f"mean-field clusters at T: "
+                 f"{[(round(c['center'], 2), round(c['mass'], 3)) for c in metrics['clusters_meanfield_T']]}")
+    if metrics["freeze_time_meanfield"] is not None:
+        lines.append(f"mean-field freeze time: t = {metrics['freeze_time_meanfield']:.2f}")
+    if "trained" in results:
+        lines.append(
+            f"V_xi trained: max {metrics['vxi']['trained_max']:.3f}, "
+            f"final {metrics['vxi']['trained_final']:.3f} "
+            f"(band {vxi_band:.3f}, mean-field {vxi_mf:.3f})")
+        lines.append(f"value gap (mean-field - trained): "
+                     f"{metrics['value_gap_meanfield_minus_trained']:.4f}")
+        if metrics["consensus_time_trained"] is not None:
+            lines.append(f"trained consensus time (Var < 1e-2): "
+                         f"t = {metrics['consensus_time_trained']:.2f}")
+    summary = "\n".join(lines)
+    with open(os.path.join(run_dir, "run_summary.txt"), "w") as f:
+        f.write(summary + "\n")
+    print(summary)
+    print(f"\nAll outputs in: {run_dir}")
+    if mirror_figs and make_plots:
+        print(f"Paper figures mirrored to: {fig_dir}")
+    with open(os.path.join(run_dir, "complete.json"), "w") as f:
+        json.dump({"status": "complete", "plots_generated": make_plots}, f)
+    return metrics
+
+
+# ============================================================
+# CLI
+# ============================================================
+
+
+def plot_run(run_dir, fig_dir=None):
+    """Regenerate figures from archived arrays; does not train or load executable code."""
+    from .bounded_confidence_plots import (fig_state_evolution, fig_opinion_traj, fig_kernel_strip, fig_cluster_count, fig_variance, fig_authority, fig_vxi, fig_phi, fig_training_history, save_opinion_gif, ARM_LABELS, ARM_COLORS)
+    from ..plotting import save_graphon_gif
+    cfg = E2Config(**json.loads((Path(run_dir) / "config.json").read_text()))
+    history_path = Path(run_dir) / "history.json"
+    history = json.loads(history_path.read_text()) if history_path.exists() else {}
+    metrics = json.loads((Path(run_dir) / "metrics.json").read_text())
+    with np.load(Path(run_dir) / "trajectories.npz", allow_pickle=False) as archive:
+        data = dict(archive)
+    NT = cfg.NT
+    t_grid, t_grid_W = data["t"], data["t_W"]
+    xi = torch.as_tensor(data["xi"])
+    arms = list(metrics["var_final"])
+    results = {arm: {"X": data[f"x_{arm}"], "K": data[f"K_{arm}"],
+                     "var": data[f"var_{arm}"], "authority": data[f"authority_{arm}"]}
+               for arm in arms}
+    if "trained" in results:
+        results["trained"]["Ws"] = torch.as_tensor(data["W_trained"])
+    vxi_trained = data.get("vxi_trained")
+    vxi_band = metrics["vxi"]["band"]
+    if fig_dir is not None:
+        Path(fig_dir).mkdir(parents=True, exist_ok=False)
+    # ---------- figures ----------
+    def both(name: str) -> List[str]:
+        if fig_dir is None:
+            return [os.path.join(run_dir, name)]
+        return [os.path.join(run_dir, name), os.path.join(fig_dir, name)]
+
+    xs_arms_np = {arm: results[arm]["X"] for arm in arms}
+    K_series = {arm: results[arm]["K"] for arm in arms}
+    var_series = {arm: results[arm]["var"] for arm in arms}
+    auth_series = {arm: results[arm]["authority"] for arm in arms}
+
+    # paper figures show mean-field vs trained only; the band arm stays in the
+    # npz/metrics archive (uncontrolled-only runs keep whatever arms they have)
+    fig_arms = ([a for a in arms if a in ("meanfield", "trained")]
+                if "trained" in results else arms)
+
+    for p in both("state_evolution.pdf"):
+        fig_state_evolution(t_grid, xs_arms_np, fig_arms, cfg, p)
+    for p in both("opinion_traj.pdf"):
+        fig_opinion_traj(t_grid, xi.cpu().numpy(), xs_arms_np, fig_arms, cfg, p)
+    for p in both("cluster_count.pdf"):
+        fig_cluster_count(t_grid, K_series, fig_arms, p)
+    for p in both("variance.pdf"):
+        fig_variance(t_grid, var_series, fig_arms, p)
+    for p in both("authority.pdf"):
+        fig_authority(t_grid, auth_series, fig_arms, p)
+    for p in both("phi.pdf"):
+        fig_phi(cfg, p)
+    if "trained" in results:
+        strip_times = [0.0, 0.1 * cfg.T, 0.2 * cfg.T, 0.3 * cfg.T, 0.4 * cfg.T, cfg.T]
+        for p in both("kernel_strip.pdf"):
+            fig_kernel_strip(results["trained"]["Ws"], cfg.dt, p, times=strip_times)
+        for p in both("vxi.pdf"):
+            fig_vxi(t_grid_W, vxi_trained, vxi_band, p)
+        fig_training_history(history, os.path.join(run_dir, "training_history.pdf"))
+
+    # ---------- GIFs ----------
+    save_opinion_gif(t_grid, xi.cpu().numpy(), xs_arms_np, K_series, fig_arms, cfg,
+                     os.path.join(run_dir, "opinions.gif"))
+    if "trained" in results:
+        n_gif = min(60, NT)
+        gif_idx = np.linspace(0, NT - 1, n_gif).round().astype(int)
+        gif_frames = [results["trained"]["Ws"][i].cpu() * cfg.N for i in gif_idx]
+        save_graphon_gif(gif_frames, os.path.join(run_dir, "kernel_time.gif"),
+                         fps=10, t_vals=[float(i * cfg.dt) for i in gif_idx])
+
+
+    complete = Path(run_dir) / "complete.json"
+    if complete.exists():
+        payload = json.loads(complete.read_text())
+        payload["plots_generated"] = True
+        complete.write_text(json.dumps(payload) + "\n")
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(description="Bounded-confidence opinion control")
+    add_common_arguments(parser)
+    for name in ("eps", "gain", "span", "amp"):
+        parser.add_argument(f"--{name}", type=float)
+    parser.add_argument("--modes", type=int)
+    parser.add_argument("--phi-form", choices=["poly", "gauss"], dest="phi_form")
+    parser.add_argument("--uncontrolled-only", action="store_true", default=None)
+    parser.add_argument("--hk-ref", action="store_true", default=None, dest="hk_ref")
+    args = parser.parse_args(argv)
+    try:
+        if args.replot:
+            plot_run(args.replot)
+            return
+        cfg = configuration_from_args(E2Config, args)
+        run_e2(cfg, tag=args.tag, out_root=args.out_root, fig_root=args.fig_root,
+               mirror_figs=args.mirror_figures, make_plots=not args.no_plots, uncontrolled_only=args.uncontrolled_only)
+    except (ValueError, FileExistsError) as exc:
+        parser.error(str(exc))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/graphon_control/experiments/bounded_confidence_plots.py b/graphon_control/experiments/bounded_confidence_plots.py
new file mode 100644
index 0000000..2d2bfb9
--- /dev/null
+++ b/graphon_control/experiments/bounded_confidence_plots.py
@@ -0,0 +1,254 @@
+"""Model-specific figures, loaded only when rendering archived runs."""
+import math
+from typing import Dict, List
+import numpy as np
+import torch
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from ..plotting import imshow_graphon, save_graphon_gif
+from .bounded_confidence import E2Config, phi_bc_poly, phi_bc_gauss, make_phi
+
+
+ARM_LABELS = {"meanfield": r"mean-field  $w \equiv 1$",
+              "band": "label-band",
+              "trained": "trained control",
+              "hk": "sharp HK (reference)"}
+
+
+ARM_COLORS = {"meanfield": "tab:gray", "band": "tab:orange",
+              "trained": "tab:blue", "hk": "tab:green"}
+
+
+def fig_state_evolution(t_grid, xs_arms: Dict[str, np.ndarray], arms: List[str],
+                        cfg: E2Config, path: str):
+    lim = cfg.span / 2 + cfg.amp + 0.5
+    fig, axes = plt.subplots(1, len(arms), figsize=(4.3 * len(arms), 3.6),
+                             sharey=True, squeeze=False)
+    for ax, arm in zip(axes[0], arms):
+        X = xs_arms[arm]  # (NT+1, N)
+        im = ax.imshow(X.T, origin="lower", aspect="auto",
+                       extent=(t_grid[0], t_grid[-1], 0.0, 1.0),
+                       cmap="RdBu_r", vmin=-lim, vmax=lim)
+        ax.set_title(ARM_LABELS[arm])
+        ax.set_xlabel(r"$t$")
+    axes[0][0].set_ylabel(r"label $\xi$")
+    fig.colorbar(im, ax=axes, fraction=0.02, pad=0.01, label=r"$x(t,\xi)$")
+    fig.suptitle("State evolution", y=1.04)
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_opinion_traj(t_grid, xi, xs_arms: Dict[str, np.ndarray], arms: List[str],
+                     cfg: E2Config, path: str):
+    """The classical HK cluster-tree: opinion trajectories colored by label."""
+    colors = plt.cm.viridis(xi)
+    fig, axes = plt.subplots(1, len(arms), figsize=(4.3 * len(arms), 3.6),
+                             sharey=True, squeeze=False)
+    for ax, arm in zip(axes[0], arms):
+        X = xs_arms[arm]  # (NT+1, N)
+        for i in range(X.shape[1]):
+            ax.plot(t_grid, X[:, i], color=colors[i], lw=0.8, alpha=0.8)
+        ax.axhline(cfg.x_target, color="0.3", lw=1.0, ls="--")
+        ax.set_title(ARM_LABELS[arm])
+        ax.set_xlabel(r"$t$")
+        ax.grid(alpha=0.15)
+    # ε scale bar on the first panel: the interaction's dead-zone width
+    ax0 = axes[0][0]
+    t_bar = t_grid[0] + 0.05 * (t_grid[-1] - t_grid[0])
+    y0 = cfg.span / 2 + cfg.amp - cfg.eps
+    ax0.annotate("", xy=(t_bar, y0), xytext=(t_bar, y0 + cfg.eps),
+                 arrowprops=dict(arrowstyle="<->", color="0.2", lw=1.2))
+    ax0.text(t_bar + 0.02 * (t_grid[-1] - t_grid[0]), y0 + cfg.eps / 2,
+             r"$\varepsilon$", va="center", fontsize=11, color="0.2")
+    axes[0][0].set_ylabel(r"opinion $x_i(t)$")
+    fig.suptitle("Opinion trajectories", y=1.02)
+    fig.tight_layout()
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_kernel_strip(Ws: torch.Tensor, dt: float, path: str, n_panels: int = 6,
+                     times: List[float] = None):
+    """
+    Kernel heatmaps on a LOG color scale: the merge-phase attention peaks are
+    an order of magnitude above the uniform level wbar = 1, and a linear scale
+    would flatten the relax-to-uniform endpoint. Panel times default to a
+    uniform grid; pass `times` to concentrate panels in the merge window.
+    """
+    from matplotlib.colors import LogNorm
+    NT = Ws.shape[0]
+    if times is None:
+        idxs = np.linspace(0, NT - 1, n_panels).round().astype(int)
+    else:
+        idxs = np.clip(np.round(np.array(times) / dt).astype(int), 0, NT - 1)
+        n_panels = len(idxs)
+    wbars = [Ws.shape[-1] * Ws[i].cpu().numpy() for i in idxs]
+    vmax = max(w.max() for w in wbars)
+    vmin = max(min(w.min() for w in wbars), vmax * 1e-3)
+    fig, axes = plt.subplots(1, n_panels, figsize=(2.6 * n_panels, 2.9))
+    for ax, i, w in zip(axes, idxs, wbars):
+        im = ax.imshow(np.maximum(w, vmin), origin="lower", extent=(0, 1, 0, 1),
+                       aspect="equal", cmap="viridis", interpolation="bicubic",
+                       norm=LogNorm(vmin=vmin, vmax=vmax))
+        ax.set_title(rf"$t = {i * dt:.1f}$", fontsize=10)
+        ax.set_xticks([0, 1])
+        ax.set_yticks([0, 1])
+        if ax is axes[0]:
+            ax.set_ylabel(r"$\xi$")
+        ax.set_xlabel(r"$\zeta$")
+    fig.colorbar(im, ax=axes, fraction=0.02, pad=0.01, label=r"$\hat{w}(t,\xi,\zeta)$")
+    fig.suptitle("Trained kernel over time", y=1.04)
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_cluster_count(t_grid, K_series: Dict[str, np.ndarray], arms: List[str],
+                      path: str):
+    fig, ax = plt.subplots(figsize=(6.4, 4.0))
+    for arm in arms:
+        ax.step(t_grid, K_series[arm], where="post", color=ARM_COLORS[arm],
+                lw=2.0, label=ARM_LABELS[arm])
+    ax.axhline(1, color="0.85", lw=0.8, zorder=0)
+    ax.set_xlabel(r"$t$")
+    ax.set_ylabel(r"clusters $K(t)$")
+    Kmax = max(int(K_series[a].max()) for a in arms)
+    ax.set_yticks(range(1, Kmax + 1))
+    ax.set_ylim(0.75, Kmax + 0.25)
+    ax.grid(alpha=0.15)
+    ax.legend()
+    ax.set_title("Opinion cluster count")
+    fig.tight_layout()
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_variance(t_grid, var_series: Dict[str, np.ndarray], arms: List[str],
+                 path: str):
+    fig, ax = plt.subplots(figsize=(6.4, 4.0))
+    for arm in arms:
+        ax.semilogy(t_grid, np.maximum(var_series[arm], 1e-12),
+                    color=ARM_COLORS[arm], lw=2.0, label=ARM_LABELS[arm])
+    ax.set_xlabel(r"$t$")
+    ax.set_ylabel(r"$\mathrm{Var}(x(t))$")
+    ax.grid(alpha=0.25, which="both")
+    ax.legend()
+    ax.set_title("Variance: frozen plateau vs consensus")
+    fig.tight_layout()
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_authority(t_grid, auth_series: Dict[str, np.ndarray], arms: List[str],
+                  path: str):
+    fig, ax = plt.subplots(figsize=(6.4, 4.0))
+    for arm in arms:
+        ax.plot(t_grid, auth_series[arm], color=ARM_COLORS[arm], lw=2.0,
+                label=ARM_LABELS[arm])
+    ax.set_xlabel(r"$t$")
+    ax.set_ylabel(r"$\frac{1}{N}\sum_i \max_j |\varphi_\varepsilon(x_j - x_i)|$")
+    ax.grid(alpha=0.25)
+    ax.legend()
+    ax.set_title("Control authority (window of controllability)")
+    fig.tight_layout()
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_vxi(t_grid_W, vxi_trained: np.ndarray, vxi_band: float, path: str):
+    fig, ax = plt.subplots(figsize=(6.4, 4.0))
+    ax.plot(t_grid_W, vxi_trained, color=ARM_COLORS["trained"], lw=2.0,
+            label=r"trained $V_\xi(\hat{w}(t))$")
+    ax.axhline(0.0, color=ARM_COLORS["meanfield"], lw=1.5, ls="--",
+               label=r"mean-field ($V_\xi = 0$)")
+    ax.axhline(vxi_band, color=ARM_COLORS["band"], lw=1.5, ls="--",
+               label=rf"label-band ($V_\xi = {vxi_band:.1f}$)")
+    ax.set_xlabel(r"$t$")
+    ax.set_ylabel(r"$V_\xi$")
+    ax.grid(alpha=0.25)
+    ax.legend()
+    ax.set_title("Receiver variation of the kernel (measured, not enforced)")
+    fig.tight_layout()
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_phi(cfg: E2Config, path: str):
+    z = torch.linspace(-2.0 * cfg.eps, 2.0 * cfg.eps, 801)
+    phi = make_phi(cfg)
+    sharp = torch.where(z.abs() <= cfg.eps, cfg.gain * z, torch.zeros_like(z))
+    fig, ax = plt.subplots(figsize=(6.0, 3.6))
+    ax.plot(z.numpy(), sharp.numpy(), color="0.6", lw=1.5, ls="--",
+            label=r"sharp HK  $a\,z\,\mathbf{1}\{|z| \leq \varepsilon\}$")
+    ax.plot(z.numpy(), phi(z).numpy(), color="tab:blue", lw=2.2,
+            label=r"$\varphi_\varepsilon(z) = a\,z\,(1 - z^2/\varepsilon^2)_+^2$")
+    for s in (-1, 1):
+        ax.axvline(s * cfg.eps, color="0.85", lw=0.8, zorder=0)
+    ax.set_xlabel(r"$z$")
+    ax.set_ylabel(r"$\varphi(z)$")
+    ax.grid(alpha=0.2)
+    ax.legend(fontsize=9)
+    ax.set_title(rf"Bounded-confidence interaction ($\varepsilon = {cfg.eps:g}$, "
+                 rf"$a = {cfg.gain:g}$)")
+    fig.tight_layout()
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_training_history(history: Dict[str, List[float]], path: str):
+    fig, axes = plt.subplots(1, 2, figsize=(11, 3.8))
+    axes[0].plot(history["loss"], lw=1.0, color="tab:blue")
+    axes[0].set_title("total loss")
+    axes[0].set_xlabel("step")
+    axes[0].grid(alpha=0.25)
+    for key, color in [("state_cost", "tab:green"), ("l2_reg_cost", "tab:red"),
+                       ("terminal_var", "tab:purple"), ("terminal_target", "tab:brown")]:
+        axes[1].semilogy(np.maximum(np.array(history[key]), 1e-10), lw=1.0,
+                         label=key, color=color)
+    axes[1].set_title("components (log scale)")
+    axes[1].set_xlabel("step")
+    axes[1].grid(alpha=0.25, which="both")
+    axes[1].legend(fontsize=8)
+    fig.tight_layout()
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def save_opinion_gif(t_grid, xi, xs_arms: Dict[str, np.ndarray],
+                     K_series: Dict[str, np.ndarray], arms: List[str],
+                     cfg: E2Config, path: str,
+                     stride: int = 2, fps: int = 12, dpi: int = 80):
+    """Animated opinion axis: dot at (x_i(t), ξ_i), colored by label."""
+    from PIL import Image
+    from io import BytesIO
+
+    colors = plt.cm.viridis(xi)
+    lim = cfg.span / 2 + cfg.amp + 0.7
+    frames = list(range(0, len(t_grid), stride))
+    if frames[-1] != len(t_grid) - 1:
+        frames.append(len(t_grid) - 1)
+
+    pil_imgs = []
+    for n in frames:
+        fig, axes = plt.subplots(1, len(arms), figsize=(3.5 * len(arms), 3.6),
+                                 squeeze=False)
+        for ax, arm in zip(axes[0], arms):
+            x_n = xs_arms[arm][n]  # (N,)
+            ax.scatter(x_n, xi, c=colors, s=22, edgecolors="none")
+            ax.axvline(cfg.x_target, color="0.8", lw=1.0, ls="--", zorder=0)
+            ax.set_xlim(-lim, lim)
+            ax.set_ylim(-0.03, 1.03)
+            ax.set_xlabel(r"opinion $x$")
+            ax.set_title(f"{ARM_LABELS[arm]}\n$K = {int(K_series[arm][n])}$",
+                         fontsize=10)
+        axes[0][0].set_ylabel(r"label $\xi$")
+        fig.suptitle(rf"Bounded confidence:  $t = {t_grid[n]:.2f}$", y=0.99)
+        fig.tight_layout()
+        buf = BytesIO()
+        fig.savefig(buf, format="png", dpi=dpi)
+        plt.close(fig)
+        buf.seek(0)
+        pil_imgs.append(Image.open(buf).copy())
+
+    pil_imgs[0].save(path, save_all=True, append_images=pil_imgs[1:], loop=0,
+                     duration=int(1000 / fps), optimize=False)
diff --git a/graphon_control/experiments/generic.py b/graphon_control/experiments/generic.py
new file mode 100644
index 0000000..640e50c
--- /dev/null
+++ b/graphon_control/experiments/generic.py
@@ -0,0 +1,256 @@
+"""General time/label control experiment using the shared numerical core.
+
+The paper coefficient is beta, not the prototype's former beta/2. Optional
+uniform, terminal-variance and scheduled-target terms are explicit extensions.
+"""
+import argparse
+import dataclasses
+import inspect
+import json
+from dataclasses import dataclass
+from functools import partial
+from pathlib import Path
+from typing import Callable, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+
+from ..controls import FourierFeatures, MLP, TimeLabelGraphonControl
+from ..dynamics import rollout_dynamics_w1
+from ..objectives import (empirical_mean, empirical_variance, graphon_l2_penalty,
+                         graphon_uniform_penalty, uniform_graphon_target,
+                         resolve_target_schedule, sampled_cost_JNdt, trajectory_cost)
+from ..initial_conditions import (sample_x0, sample_x0_structured, sample_x0_two_clusters,
+                                  sample_x0_linear_label, sample_x0_linear_sine)
+from ..interactions import phi_default, phi_linear
+from ..training import train_controller
+from ..validation import validate_common_config
+from ..artifacts import create_run_directory, write_run_provenance, save_checkpoint
+from ..cli import add_common_arguments, configuration_from_args
+
+
+@dataclass
+class TrainConfig:
+    # discretization/noise parameters
+    N: int = 64                     # number of agents / label grid points
+    T: float = 1.0                  # total time horizon
+    dt: float = 0.02                # time step for Euler-Maruyama simulation and cost discretization
+    sigma: float = 0.15             # noise scale in dynamics
+
+    # cost parameters
+    beta: float = 1e-2              # control regularization coefficient in J^{N,Δt}
+    eta_uniform: float = 0.0        # coefficient for uniform-graphon regularizer
+    regularizer_mode: str = "l2"    # one of: "l2", "uniform", "l2_plus_uniform"
+    lambda_terminal: float = 1.0    # weight on terminal variance penalty
+    gamma_target: float = 0.0       # weight on terminal mean-tracking penalty
+    x_target: float = 0.0           # target for the terminal empirical mean
+    target_times: List[float] = dataclasses.field(default_factory=list)    # additional checkpoint times in [0, T]
+    target_values: List[float] = dataclasses.field(default_factory=list)   # target mean value at each checkpoint time
+    target_weights: List[float] = dataclasses.field(default_factory=list)  # optional weights for checkpoint targets; defaults to gamma_target if omitted
+
+    # graphon architecture
+    enforce_symmetry: bool = False   # if True, enforce W_ij = W_ji exactly
+    positive_map: str = "softplus"   # nonnegativity map for symmetric mode: "softplus" or "exp"
+
+    # Monte Carlo / optimization
+    batch_size: int = 32
+    num_steps: int = 1000
+    lr: float = 1e-3
+    grad_clip: Optional[float] = 1.0
+
+    # network architecture
+    num_fourier_freqs: int = 8    # m in the draft
+    embed_dim: int = 64           # d in the draft
+    hidden_embed: int = 128
+    hidden_score: int = 128
+    time_scale: Optional[float] = None       # rescales t before feeding to score net
+
+    # loss weighting
+    weighted_running_cost: bool = False  # if True, linearly ramp running state cost weight from ~0 to 1 over [0, T]
+
+    # interaction scaling
+    interaction_alpha: float = 1.0      # linear interaction gain; positive is attractive (consensus-seeking) under the paper convention φ(x_j - x_i)
+
+    # reproducibility / device
+    eval_batch_size: int = 128
+    evaluation_seed: int = 1729
+    score_bound: Optional[float] = None
+    seed: int = 0
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+
+    def __post_init__(self):
+        validate_common_config(self)
+        if self.regularizer_mode not in ("l2", "uniform", "l2_plus_uniform"):
+            raise ValueError("unknown regularizer_mode")
+        resolve_target_schedule(self.T, self.dt, self.gamma_target,
+                                self.target_times, self.target_values, self.target_weights)
+
+
+def train_algorithm1_w1(config: TrainConfig, phi=None,
+                        x0_sampler=sample_x0_two_clusters):
+    """Train an open-loop controller; samplers return (batch, agents)."""
+    config.__post_init__()
+    phi = partial(phi_default, alpha=config.interaction_alpha) if phi is None else phi
+    torch.manual_seed(config.seed)
+    xi = (torch.arange(config.N, device=config.device, dtype=torch.float32) + 0.5) / config.N
+    model = TimeLabelGraphonControl(
+        xi, config.num_fourier_freqs, config.embed_dim, config.hidden_embed,
+        config.hidden_score, time_scale=config.T if config.time_scale is None else config.time_scale,
+        enforce_symmetry=config.enforce_symmetry, positive_map=config.positive_map,
+        score_bound=config.score_bound)
+
+    def objective():
+        loss, aux = sampled_cost_JNdt(
+            model, x0_sampler(config.batch_size, xi), T=config.T, dt=config.dt,
+            sigma=config.sigma, beta=config.beta, lambda_terminal=config.lambda_terminal,
+            phi=phi, gamma_target=config.gamma_target, x_target=config.x_target,
+            target_times=config.target_times, target_values=config.target_values,
+            target_weights=config.target_weights, eta_uniform=config.eta_uniform,
+            regularizer_mode=config.regularizer_mode,
+            weighted_running_cost=config.weighted_running_cost)
+        return {'total': loss, **aux}
+
+    history, last = train_controller(model, objective, config.num_steps, config.lr,
+                                    config.grad_clip, description='training')
+    return model, history, xi, {key: val for key, val in last.items() if key != 'total'}
+
+
+@torch.no_grad()
+def evaluate_policy(model, xi, T, dt, sigma, batch_size,
+                    phi=phi_default, x0_sampler=sample_x0_two_clusters):
+    model.eval()
+    return rollout_dynamics_w1(model, x0_sampler(batch_size, xi), T, dt, sigma, phi)
+
+
+def save_run(out_root, config, history, final_aux, phi, x0_sampler,
+             x0_sample, run_tag='', reserved_run_dir=None):
+    """Persist configuration/provenance and training diagnostics to a new run."""
+    run_dir = (create_run_directory(out_root, 'run', run_tag or None)
+               if reserved_run_dir is None else Path(reserved_run_dir))
+    if reserved_run_dir is None:
+        write_run_provenance(run_dir, config, config.seed, config.evaluation_seed,
+                             device=config.device)
+    for filename, value in [('history.json', history), ('final_aux.json', final_aux)]:
+        (run_dir / filename).write_text(json.dumps(value, indent=2) + '\n')
+    try:
+        source = inspect.getsource(phi)
+    except (OSError, TypeError):
+        source = repr(phi)
+    summary = (f'Interaction: {getattr(phi, "__name__", repr(phi))}\n'
+               f'Sampler: {getattr(x0_sampler, "__name__", repr(x0_sampler))}\n'
+               f'Initial sample mean: {float(x0_sample.mean()):.6g}\n'
+               f'Interaction source:\n{source}\n')
+    (run_dir / 'run_summary.txt').write_text(summary)
+    return str(run_dir)
+
+
+@torch.no_grad()
+def _archive_evaluation(run_dir, model, config, xi, x0_sample, phi, sampler, batch_size):
+    xs, Ws = evaluate_policy(model, xi, config.T, config.dt, config.sigma,
+                             batch_size, phi=phi, x0_sampler=sampler)
+    schedule = resolve_target_schedule(config.T, config.dt, config.gamma_target,
+                                       config.target_times, config.target_values, config.target_weights)
+    parts = trajectory_cost(torch.stack(xs), torch.stack(Ws), dt=config.dt, beta=config.beta,
+                            lambda_terminal=config.lambda_terminal, gamma_target=config.gamma_target,
+                            x_target=config.x_target, target_schedule=schedule,
+                            eta_uniform=config.eta_uniform, regularizer_mode=config.regularizer_mode,
+                            weighted_running_cost=config.weighted_running_cost)
+    metrics = {key: float(value) for key, value in parts.items()}
+    Path(run_dir, 'metrics.json').write_text(json.dumps(metrics, indent=2) + '\n')
+    phi_grid = torch.linspace(-20, 20, 401, device=xi.device, dtype=xi.dtype)
+    np.savez_compressed(Path(run_dir) / 'trajectories.npz',
+                        xs=torch.stack(xs).cpu().numpy(), Ws=torch.stack(Ws).cpu().numpy(),
+                        xi=xi.cpu().numpy(), x0_sample=x0_sample.cpu().numpy(),
+                        phi_grid=phi_grid.cpu().numpy(), phi_values=phi(phi_grid).cpu().numpy())
+    return xs, Ws
+
+
+def run_experiment(config: TrainConfig, phi=None, x0_sampler=sample_x0_two_clusters,
+                   out_root='out_1', run_tag='', eval_batch_size=None, make_plots=True):
+    """Train, archive, evaluate and optionally render using the same shared core."""
+    eval_batch_size = config.eval_batch_size if eval_batch_size is None else eval_batch_size
+    if not isinstance(eval_batch_size, int) or isinstance(eval_batch_size, bool) or eval_batch_size < 1:
+        raise ValueError('eval_batch_size must be a positive integer')
+    config = dataclasses.replace(config, eval_batch_size=eval_batch_size)
+    reserved = create_run_directory(out_root, 'run', run_tag or None)
+    write_run_provenance(reserved, config, config.seed, config.evaluation_seed, device=config.device)
+    phi = partial(phi_default, alpha=config.interaction_alpha) if phi is None else phi
+    model, history, xi, final_aux = train_algorithm1_w1(config, phi, x0_sampler)
+    torch.manual_seed(config.evaluation_seed)
+    x0_sample = x0_sampler(config.batch_size, xi)
+    run_dir = save_run(out_root, config, history, final_aux, phi, x0_sampler, x0_sample,
+                       run_tag, reserved_run_dir=reserved)
+    save_checkpoint(Path(run_dir) / 'model.pt', model)
+    xs, Ws = _archive_evaluation(run_dir, model, config, xi, x0_sample, phi,
+                                  x0_sampler, eval_batch_size)
+    if make_plots:
+        plot_run(run_dir)
+    Path(run_dir, 'complete.json').write_text(json.dumps({'status': 'complete', 'plots_generated': make_plots}) + '\n')
+    print(f'Results saved to: {run_dir}')
+    return {'model': model, 'history': history, 'xi': xi, 'xs': xs, 'Ws': Ws, 'run_dir': run_dir}
+
+
+def plot_run(run_dir):
+    """Render saved numerical arrays without retraining or evaluating a callable."""
+    import matplotlib
+    matplotlib.use('Agg')
+    import matplotlib.pyplot as plt
+    from ..plotting import (imshow_graphon, plot_graphon_diagnostics, heterogeneity_stats,
+                            plot_state_histogram_over_time_imshow, plot_x0_samples,
+                            plot_x0_density, save_graphon_gif)
+    run_dir = Path(run_dir)
+    cfg = TrainConfig(**json.loads((run_dir / 'config.json').read_text()))
+    with np.load(run_dir / 'trajectories.npz', allow_pickle=False) as archive:
+        data = {name: torch.as_tensor(value) for name, value in archive.items()}
+    xs, Ws, xi, x0 = data['xs'], data['Ws'], data['xi'], data['x0_sample']
+    lines = []
+    for label, index in [('early', len(Ws) // 4), ('mid', len(Ws) // 2), ('late', 3 * len(Ws) // 4)]:
+        W = Ws[index]
+        lines.append(f'{label}: row sums {W.sum(-1).tolist()}; statistics {heterogeneity_stats(W)}')
+        plot_graphon_diagnostics(cfg.N * W, title_prefix=f'{label}-time ', show=False)
+        plt.gcf().savefig(run_dir / f'graphon_diag_{label}.png', dpi=150)
+        plt.close(plt.gcf())
+    (run_dir / 'eval_diagnostics.txt').write_text('\n'.join(lines) + '\n')
+    fig, ax = plt.subplots(figsize=(7, 6))
+    imshow_graphon(cfg.N * Ws[len(Ws) // 2], ax=ax, show=False)
+    fig.savefig(run_dir / 'graphon_mid.png', dpi=150)
+    plt.close(fig)
+    plot_state_histogram_over_time_imshow(list(xs), dt=cfg.dt, show=False, x_target=cfg.x_target)
+    plt.gcf().savefig(run_dir / 'state_histogram.png', dpi=150)
+    plt.close(plt.gcf())
+    for filename, plotter in [('ic_samples.png', plot_x0_samples), ('ic_density.png', plot_x0_density)]:
+        fig, _ = plotter(xi, x0)
+        fig.savefig(run_dir / filename, dpi=150)
+        plt.close(fig)
+    fig, ax = plt.subplots()
+    ax.plot(data['phi_grid'], data['phi_values'])
+    ax.set(xlabel='state difference', ylabel='interaction')
+    fig.savefig(run_dir / 'phi_plot.png', dpi=150)
+    plt.close(fig)
+    indices = np.linspace(0, len(Ws) - 1, min(60, len(Ws))).round().astype(int)
+    save_graphon_gif([cfg.N * Ws[i] for i in indices], str(run_dir / 'graphon_time.gif'),
+                     t_vals=[float(i * cfg.dt) for i in indices])
+
+    complete = run_dir / "complete.json"
+    if complete.exists():
+        payload = json.loads(complete.read_text())
+        payload["plots_generated"] = True
+        complete.write_text(json.dumps(payload) + "\n")
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(description='General time-and-label graphon control')
+    add_common_arguments(parser)
+    args = parser.parse_args(argv)
+    try:
+        if args.replot:
+            plot_run(args.replot)
+            return
+        cfg = configuration_from_args(TrainConfig, args)
+        run_experiment(cfg, out_root=args.out_root, run_tag=args.tag, make_plots=not args.no_plots)
+    except (ValueError, FileExistsError) as exc:
+        parser.error(str(exc))
+
+
+if __name__ == '__main__':
+    main()
diff --git a/graphon_control/experiments/kuramoto.py b/graphon_control/experiments/kuramoto.py
new file mode 100644
index 0000000..22a1f33
--- /dev/null
+++ b/graphon_control/experiments/kuramoto.py
@@ -0,0 +1,463 @@
+"""Kuramoto untwisting experiment using the shared paper conventions.
+
+See docs/NUMERICS.md for equations and docs/DEVELOPMENT.md for reproducible runs.
+Model-specific rendering is in kuramoto_plots.py.
+"""
+import argparse
+import datetime
+import json
+import math
+import os
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Optional
+
+import numpy as np
+import torch
+
+from ..controls import TimeLabelGraphonControl
+from ..artifacts import create_run_directory, write_run_provenance, save_checkpoint
+from ..dynamics import ConstantKernel, mean_field_kernel, precompute_Ws, rollout_with_Ws
+from ..objectives import paper_cost
+from ..diagnostics import vxi_discrete
+from ..training import train_controller
+from ..validation import n_time_steps, validate_common_config
+from ..cli import add_common_arguments, configuration_from_args
+from pathlib import Path
+
+
+# ============================================================
+# Configuration (paper protocol, §5.1)
+# ============================================================
+
+@dataclass
+class E1Config:
+    # experiment
+    k: int = 1                 # winding number of the twisted state
+    jitter: float = 0.1        # i.i.d. jitter on the training initial data (resampled each batch)
+    ring_range: float = 0.1    # range r of the smooth ring baseline
+    x0_level: int = 0          # if > 0: use the coarse datum P_M(twist) with M = x0_level
+                               # (the common datum of lem:AN-monotone; requires M | N)
+
+    # discretization
+    N: int = 64
+    T: float = 5.0
+    dt: float = 0.05
+    sigma: float = 0.0         # deterministic dynamics (theory of §3 is deterministic)
+
+    # cost — PAPER convention (main5.tex, eqn: JN definition / eqn:J-cost-1):
+    #   J = Σ_n Δt [ Var(x^n) + β ||w(t_n)||²_{L²(I²)} ] + γ (mean(x^T) - x_target)²
+    # (β, not β/2, and NO terminal-variance term; lambda_terminal kept as an
+    #  optional knob but defaults to 0 to match the paper)
+    beta: float = 10.0
+    lambda_terminal: float = 0.0
+    gamma_target: float = 100.0
+    x_target: float = 0.0
+
+    # optimization
+    batch_size: int = 16
+    num_steps: int = 2000
+    lr: float = 1e-3
+    grad_clip: float = 1.0
+
+    # architecture (model_1.py defaults)
+    num_fourier_freqs: int = 8
+    embed_dim: int = 64
+    hidden_embed: int = 128
+    hidden_score: int = 128
+
+    # bookkeeping
+    evaluation_seed: int = 1729
+    score_bound: Optional[float] = None
+    seed: int = 0
+    device: str = "cpu"
+    time_chunk: int = 50       # chunk size for the batched W(t) precomputation
+
+    def __post_init__(self):
+        validate_common_config(self)
+        if not isinstance(self.k, int) or isinstance(self.k, bool) or self.k < 1 or 2 * self.k >= self.N:
+            raise ValueError("k must be positive and satisfy 2*k < N to resolve the twist")
+        if isinstance(self.x0_level, bool) or not isinstance(self.x0_level, int):
+            raise ValueError("x0_level must be an integer")
+        if self.x0_level and (self.x0_level < 2 or self.N % self.x0_level):
+            raise ValueError("x0_level must be zero or a divisor of N >= 2")
+        if not math.isfinite(self.ring_range) or self.ring_range <= 0:
+            raise ValueError("ring_range must be finite and positive")
+
+    @property
+    def NT(self) -> int:
+        return n_time_steps(self.T, self.dt)
+
+
+# ============================================================
+# E1 ingredients
+# ============================================================
+
+def phi_sin(y: torch.Tensor) -> torch.Tensor:
+    """Kuramoto interaction φ(z) = sin z (bounded, Lipschitz, odd; φ(0)=0)."""
+    return torch.sin(y)
+
+
+def twisted_state(xi: torch.Tensor, k: int) -> torch.Tensor:
+    """k-twisted state x(0, ξ) = 2πk(ξ - 1/2)."""
+    return 2.0 * math.pi * k * (xi - 0.5)
+
+
+def coarse_twisted_state(xi: torch.Tensor, k: int, M: int) -> torch.Tensor:
+    """
+    P_M projection of the k-twisted state: block average over M cells.
+    Since the twist is linear in ξ, the block average equals the twist at the
+    M cell midpoints, repeated N/M times — piecewise constant on the level-M
+    partition, hence exactly representable at every dyadic N ≥ M. This is the
+    common initial datum x_M(0) of lem:AN-monotone / thm:main (main5.tex).
+    """
+    N = xi.numel()
+    if N % M != 0:
+        raise ValueError(f"N={N} must be a multiple of M={M}")
+    mid = (torch.arange(M, device=xi.device, dtype=xi.dtype) + 0.5) / M
+    vals = twisted_state(mid, k)                     # (M,)
+    return vals.repeat_interleave(N // M)            # (N,)
+
+
+def smooth_ring_kernel(xi: torch.Tensor, r: float) -> torch.Tensor:
+    """
+    Mollified ring kernel: smooth bump of range r in torus distance,
+    row-normalized to a row-stochastic W. Smooth in ξ, hence of finite
+    receiver variation (unlike the sharp indicator (1/2r) 1_{[-r,r]}).
+    """
+    d = (xi.unsqueeze(0) - xi.unsqueeze(1)).abs()
+    d = torch.minimum(d, 1.0 - d)  # torus distance
+    u = (d / r).clamp(max=1.0 - 1e-6)
+    bump = torch.where(d < r, torch.exp(-1.0 / (1.0 - u ** 2)), torch.zeros_like(d))
+    return bump / bump.sum(dim=1, keepdim=True)
+
+
+# ============================================================
+# Fast rollout: W depends only on t, so precompute all W(t_n)
+# in a few batched forward passes, then run the Euler loop.
+# ============================================================
+
+
+def e1_cost(xs, Ws, cfg: E1Config):
+    """Paper objective, equation (60); shared across both experiments."""
+    return paper_cost(xs, Ws, cfg)
+
+
+def order_parameter(xs: torch.Tensor) -> torch.Tensor:
+    """
+    Kuramoto order parameter r(t) = | (1/N) Σ_j e^{i x_j(t)} |.
+    xs: (..., N) -> (...)
+    """
+    return torch.sqrt(torch.cos(xs).mean(-1) ** 2 + torch.sin(xs).mean(-1) ** 2)
+
+
+def winding_number(xs: torch.Tensor) -> torch.Tensor:
+    """
+    Winding number of the state along the label circle:
+      W(t) = (1/2π) [ Σ_i wrap(x_{i+1} - x_i) + wrap(x_1 - x_N) ],
+    with wrap(d) = ((d + π) mod 2π) - π. The k-twisted state has W = k;
+    consensus has W = 0. Untwisting proceeds in integer quanta: each unit of
+    winding requires transporting the outer agents a full 2π of real phase
+    (the sine coupling is blind to exact 2π gaps), reeled in through the
+    intermediate agents. xs: (..., N) -> (...).
+    """
+    gaps = xs[..., 1:] - xs[..., :-1]
+    close = (xs[..., :1] - xs[..., -1:])
+    all_gaps = torch.cat([gaps, close], dim=-1)
+    wrapped = torch.remainder(all_gaps + math.pi, 2.0 * math.pi) - math.pi
+    return wrapped.sum(dim=-1) / (2.0 * math.pi)
+
+
+
+
+# ============================================================
+# Main experiment
+# ============================================================
+
+def run_e1(cfg: E1Config, tag: str = "", out_root: str = "out_1", fig_root: str = "figures",
+           mirror_figs: bool = False, make_plots: bool = True):
+    torch.manual_seed(cfg.seed)
+    np.random.seed(cfg.seed)
+    device = torch.device(cfg.device)
+    dtype = torch.float32
+
+    cfg.__post_init__()
+    run_dir = str(create_run_directory(out_root, "e1", tag or None))
+    fig_dir = os.path.join(fig_root, os.path.basename(run_dir))
+    if mirror_figs and make_plots and Path(fig_dir).exists():
+        raise FileExistsError(f"Figure export already exists: {fig_dir}")
+    write_run_provenance(run_dir, cfg, cfg.seed, cfg.evaluation_seed, device=cfg.device)
+
+    xi = (torch.arange(cfg.N, device=device, dtype=dtype) + 0.5) / cfg.N
+    if cfg.x0_level > 0:
+        base = coarse_twisted_state(xi, cfg.k, cfg.x0_level)
+    else:
+        base = twisted_state(xi, cfg.k)
+    NT = cfg.NT
+
+    # Sanity assert: the twisted datum (smooth or coarse) must be exactly
+    # stationary for the mean-field baseline (roots-of-unity cancellation);
+    # if this fails, the "pinned baseline" claim of §5.2 is wrong for this grid.
+    with torch.no_grad():
+        drift0 = (mean_field_kernel(cfg.N, device) * phi_sin(
+            base.unsqueeze(0) - base.unsqueeze(1))).sum(-1)
+        pin_err = float(drift0.abs().max())
+    if pin_err > 1e-5:
+        raise RuntimeError(f"baseline pinning violated: |drift|={pin_err:.2e} "
+                           f"(N={cfg.N}, k={cfg.k}, x0_level={cfg.x0_level})")
+
+    print("=" * 64)
+    print(f"E1 Kuramoto untwisting | k={cfg.k} N={cfg.N} T={cfg.T} dt={cfg.dt}")
+    print(f"steps={cfg.num_steps} batch={cfg.batch_size} lr={cfg.lr} "
+          f"beta={cfg.beta} gamma={cfg.gamma_target} lambda={cfg.lambda_terminal}")
+    print(f"device={cfg.device} tag={tag}")
+    print("=" * 64)
+
+    # ---------- model ----------
+    model = TimeLabelGraphonControl(
+        xi=xi,
+        num_fourier_freqs=cfg.num_fourier_freqs,
+        embed_dim=cfg.embed_dim,
+        hidden_embed=cfg.hidden_embed,
+        hidden_score=cfg.hidden_score,
+        time_scale=max(cfg.T, 1e-8),
+        enforce_symmetry=False,
+        score_bound=cfg.score_bound,
+    ).to(device)
+    def objective():
+        Ws = precompute_Ws(model, NT, cfg.dt, device, cfg.time_chunk)
+        x0 = base.unsqueeze(0) + cfg.jitter * torch.randn(
+            cfg.batch_size, cfg.N, device=device, dtype=dtype)
+        xs = rollout_with_Ws(Ws, x0, cfg.dt, phi_sin, sigma=cfg.sigma)
+        return e1_cost(xs, Ws, cfg)
+
+    history, _ = train_controller(model, objective, cfg.num_steps, cfg.lr,
+                                  cfg.grad_clip, description="E1 training")
+    torch.manual_seed(cfg.evaluation_seed)
+
+    # ---------- evaluation: deterministic, exact twist, all three arms ----------
+    model.eval()
+    arms: Dict[str, object] = {
+        "meanfield": ConstantKernel(mean_field_kernel(cfg.N, device)),
+        "ring": ConstantKernel(smooth_ring_kernel(xi, cfg.ring_range)),
+        "trained": model,
+    }
+    x0_eval = base.unsqueeze(0)  # (1, N), no jitter: the exact stationary point
+
+    xs_arms_np: Dict[str, np.ndarray] = {}
+    Ws_arms: Dict[str, torch.Tensor] = {}
+    r_series: Dict[str, np.ndarray] = {}
+    var_series: Dict[str, np.ndarray] = {}
+    wind_series: Dict[str, np.ndarray] = {}
+    cost_summary: Dict[str, Dict[str, float]] = {}
+
+    with torch.no_grad():
+        for arm, kernel in arms.items():
+            Ws_eval = precompute_Ws(kernel, NT, cfg.dt, device, cfg.time_chunk)
+            xs_eval = rollout_with_Ws(Ws_eval, x0_eval, cfg.dt, phi_sin, sigma=0.0)
+            parts = e1_cost(xs_eval, Ws_eval, cfg)
+            X = xs_eval[:, 0, :].cpu().numpy()          # (NT+1, N)
+            xs_arms_np[arm] = X
+            Ws_arms[arm] = Ws_eval
+            r_series[arm] = order_parameter(xs_eval[:, 0, :]).cpu().numpy()
+            var_series[arm] = xs_eval[:, 0, :].var(dim=1, unbiased=False).cpu().numpy()
+            wind_series[arm] = winding_number(xs_eval[:, 0, :]).cpu().numpy()
+            cost_summary[arm] = {kk: float(vv) for kk, vv in parts.items()}
+
+    vxi_trained = vxi_discrete(Ws_arms["trained"]).cpu().numpy()   # (NT,)
+    vxi_ring = float(vxi_discrete(Ws_arms["ring"][0]))
+    vxi_mf = float(vxi_discrete(Ws_arms["meanfield"][0]))
+
+    t_grid = np.arange(NT + 1) * cfg.dt
+    t_grid_W = np.arange(NT) * cfg.dt
+
+    # ---------- noisy ensemble evaluation (robustness runs, sigma > 0) ----------
+    # The deterministic evaluation above supplies the paper's table values (the
+    # noisy training loss carries a noise floor ~ sigma^2 T^2 / 2); here we
+    # additionally roll the trained control on the noisy dynamics it was trained
+    # for, over an ensemble of Brownian paths from the exact twist.
+    noisy_eval: Dict[str, np.ndarray] = {}
+    if cfg.sigma > 0.0:
+        n_rep = 32
+        with torch.no_grad():
+            xs_noisy = rollout_with_Ws(Ws_arms["trained"], x0_eval.expand(n_rep, -1),
+                                       cfg.dt, phi_sin, sigma=cfg.sigma)  # (NT+1, n_rep, N)
+            r_rep = order_parameter(xs_noisy).cpu().numpy()               # (NT+1, n_rep)
+            w_rep = winding_number(xs_noisy).cpu().numpy()
+        noisy_eval = {
+            "r_noisy_mean": r_rep.mean(axis=1),
+            "r_noisy_q10": np.quantile(r_rep, 0.10, axis=1),
+            "r_noisy_q90": np.quantile(r_rep, 0.90, axis=1),
+            "winding_noisy_mean": w_rep.mean(axis=1),
+            "winding_noisy_q10": np.quantile(w_rep, 0.10, axis=1),
+            "winding_noisy_q90": np.quantile(w_rep, 0.90, axis=1),
+        }
+
+    # ---------- persistence ----------
+    save_checkpoint(os.path.join(run_dir, "model.pt"), model)
+    with open(os.path.join(run_dir, "history.json"), "w") as f:
+        json.dump(history, f, indent=2)
+
+    np.savez_compressed(
+        os.path.join(run_dir, "trajectories.npz"),
+        t=t_grid, t_W=t_grid_W, xi=xi.cpu().numpy(),
+        x_meanfield=xs_arms_np["meanfield"], x_ring=xs_arms_np["ring"],
+        x_trained=xs_arms_np["trained"],
+        W_trained=Ws_arms["trained"].cpu().numpy().astype(np.float32),
+        W_ring=Ws_arms["ring"][0].cpu().numpy().astype(np.float32),
+        r_meanfield=r_series["meanfield"], r_ring=r_series["ring"],
+        r_trained=r_series["trained"],
+        var_meanfield=var_series["meanfield"], var_ring=var_series["ring"],
+        var_trained=var_series["trained"],
+        winding_meanfield=wind_series["meanfield"], winding_ring=wind_series["ring"],
+        winding_trained=wind_series["trained"],
+        vxi_trained=vxi_trained,
+        **noisy_eval,
+    )
+
+    metrics = {
+        "k": cfg.k,
+        "costs": cost_summary,
+        "r_initial": float(r_series["trained"][0]),
+        "r_final": {arm: float(r_series[arm][-1]) for arm in arms},
+        "var_final": {arm: float(var_series[arm][-1]) for arm in arms},
+        "winding_final": {arm: float(wind_series[arm][-1]) for arm in arms},
+        "terminal_mean_trained": cost_summary["trained"]["terminal_mean"],
+        "vxi": {"trained_max": float(vxi_trained.max()),
+                "trained_final": float(vxi_trained[-1]),
+                "ring": vxi_ring, "meanfield": vxi_mf},
+        "value_gap_meanfield_minus_trained":
+            cost_summary["meanfield"]["total"] - cost_summary["trained"]["total"],
+    }
+    if noisy_eval:
+        metrics["noisy_eval"] = {
+            "sigma": cfg.sigma,
+            "n_replicates": 32,
+            "r_final_mean": float(noisy_eval["r_noisy_mean"][-1]),
+            "r_final_q10": float(noisy_eval["r_noisy_q10"][-1]),
+            "r_final_q90": float(noisy_eval["r_noisy_q90"][-1]),
+            "winding_final_mean": float(noisy_eval["winding_noisy_mean"][-1]),
+        }
+    with open(os.path.join(run_dir, "metrics.json"), "w") as f:
+        json.dump(metrics, f, indent=2)
+
+    if make_plots:
+        plot_run(run_dir, fig_dir=fig_dir if mirror_figs else None)
+
+    # ---------- summary ----------
+    lines = [
+        "=== E1 Kuramoto untwisting: run summary ===",
+        f"timestamp : {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
+        f"run_dir   : {run_dir}",
+        f"k = {cfg.k}, N = {cfg.N}, T = {cfg.T}, dt = {cfg.dt}, steps = {cfg.num_steps}",
+        "",
+        "arm         J_total      r(T)      Var(T)    winding(T)",
+    ]
+    for arm in ["meanfield", "ring", "trained"]:
+        lines.append(f"{arm:<11s} {cost_summary[arm]['total']:>9.4f}  "
+                     f"{metrics['r_final'][arm]:>8.4f}  {metrics['var_final'][arm]:>10.6f}  "
+                     f"{metrics['winding_final'][arm]:>+9.3f}")
+    lines += [
+        "",
+        f"V_xi trained: max {metrics['vxi']['trained_max']:.3f}, "
+        f"final {metrics['vxi']['trained_final']:.3f} "
+        f"(ring {vxi_ring:.3f}, mean-field {vxi_mf:.3f})",
+        f"value gap (mean-field - trained): "
+        f"{metrics['value_gap_meanfield_minus_trained']:.4f}",
+    ]
+    summary = "\n".join(lines)
+    with open(os.path.join(run_dir, "run_summary.txt"), "w") as f:
+        f.write(summary + "\n")
+    print(summary)
+    print(f"\nAll outputs in: {run_dir}")
+    if mirror_figs and make_plots:
+        print(f"Paper figures mirrored to: {fig_dir}")
+    with open(os.path.join(run_dir, "complete.json"), "w") as f:
+        json.dump({"status": "complete", "plots_generated": make_plots}, f)
+    return metrics
+
+
+# ============================================================
+# CLI
+# ============================================================
+
+
+def plot_run(run_dir, fig_dir=None):
+    """Regenerate figures from archived arrays; does not train or load executable code."""
+    from .kuramoto_plots import (fig_r_panels, fig_r_combined, fig_winding, fig_variance, fig_vxi, fig_kernel_strip, fig_state_evolution, fig_training_history, save_untwist_gif, ARM_LABELS, ARM_COLORS, PAPER_ARMS)
+    from ..plotting import save_graphon_gif
+    cfg = E1Config(**json.loads((Path(run_dir) / "config.json").read_text()))
+    history = json.loads((Path(run_dir) / "history.json").read_text())
+    metrics = json.loads((Path(run_dir) / "metrics.json").read_text())
+    with np.load(Path(run_dir) / "trajectories.npz", allow_pickle=False) as archive:
+        data = dict(archive)
+    NT = cfg.NT
+    t_grid, t_grid_W = data["t"], data["t_W"]
+    xi = torch.as_tensor(data["xi"])
+    xs_arms_np = {arm: data[f"x_{arm}"] for arm in ("meanfield", "ring", "trained")}
+    r_series = {arm: data[f"r_{arm}"] for arm in xs_arms_np}
+    var_series = {arm: data[f"var_{arm}"] for arm in xs_arms_np}
+    wind_series = {arm: data[f"winding_{arm}"] for arm in xs_arms_np}
+    Ws_arms = {"trained": torch.as_tensor(data["W_trained"])}
+    vxi_trained = data["vxi_trained"]
+    vxi_ring = metrics["vxi"]["ring"]
+    if fig_dir is not None:
+        Path(fig_dir).mkdir(parents=True, exist_ok=False)
+    # ---------- figures ----------
+    def both(name: str) -> List[str]:
+        if fig_dir is None:
+            return [os.path.join(run_dir, name)]
+        return [os.path.join(run_dir, name), os.path.join(fig_dir, name)]
+
+    for p in both("r_panels.pdf"):
+        fig_r_panels(t_grid, r_series, cfg.k, p, arms=PAPER_ARMS)
+    for p in both("r_combined.pdf"):
+        fig_r_combined(t_grid, r_series, cfg.k, p, arms=PAPER_ARMS)
+    for p in both("variance.pdf"):
+        fig_variance(t_grid, var_series, p, arms=PAPER_ARMS)
+    for p in both("winding.pdf"):
+        fig_winding(t_grid, wind_series, cfg.k, p, arms=PAPER_ARMS)
+    for p in both("vxi.pdf"):
+        fig_vxi(t_grid_W, vxi_trained, vxi_ring, p)
+    for p in both("kernel_strip.pdf"):
+        fig_kernel_strip(Ws_arms["trained"], cfg.dt, p)
+    for p in both("state_evolution.pdf"):
+        fig_state_evolution(t_grid, xi.cpu().numpy(), xs_arms_np, cfg.k, p,
+                            arms=PAPER_ARMS)
+    fig_training_history(history, os.path.join(run_dir, "training_history.pdf"))
+
+    # ---------- GIFs ----------
+    save_untwist_gif(t_grid, xi.cpu().numpy(), xs_arms_np, r_series, cfg.k,
+                     os.path.join(run_dir, "untwist.gif"), arms=PAPER_ARMS)
+    n_gif = min(60, NT)
+    gif_idx = np.linspace(0, NT - 1, n_gif).round().astype(int)
+    gif_frames = [Ws_arms["trained"][i].cpu() * cfg.N for i in gif_idx]
+    save_graphon_gif(gif_frames, os.path.join(run_dir, "kernel_time.gif"),
+                     fps=10, t_vals=[float(i * cfg.dt) for i in gif_idx])
+
+
+    complete = Path(run_dir) / "complete.json"
+    if complete.exists():
+        payload = json.loads(complete.read_text())
+        payload["plots_generated"] = True
+        complete.write_text(json.dumps(payload) + "\n")
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(description="Kuramoto twisted-state control")
+    add_common_arguments(parser)
+    parser.add_argument("--k", type=int)
+    parser.add_argument("--x0-level", type=int, dest="x0_level")
+    args = parser.parse_args(argv)
+    try:
+        if args.replot:
+            plot_run(args.replot)
+            return
+        cfg = configuration_from_args(E1Config, args)
+        run_e1(cfg, tag=args.tag, out_root=args.out_root, fig_root=args.fig_root,
+               mirror_figs=args.mirror_figures, make_plots=not args.no_plots)
+    except (ValueError, FileExistsError) as exc:
+        parser.error(str(exc))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/graphon_control/experiments/kuramoto_plots.py b/graphon_control/experiments/kuramoto_plots.py
new file mode 100644
index 0000000..2eae907
--- /dev/null
+++ b/graphon_control/experiments/kuramoto_plots.py
@@ -0,0 +1,214 @@
+"""Model-specific figures, loaded only when rendering archived runs."""
+import math
+from typing import Dict, List
+import numpy as np
+import torch
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from ..plotting import imshow_graphon, save_graphon_gif
+
+
+ARM_LABELS = {"meanfield": r"mean-field  $w \equiv 1$",
+              "ring": "smooth ring",
+              "trained": "trained control"}
+
+
+ARM_COLORS = {"meanfield": "tab:gray", "ring": "tab:orange", "trained": "tab:blue"}
+
+
+PAPER_ARMS = ["meanfield", "trained"]
+
+
+def fig_r_panels(t_grid, r_series: Dict[str, np.ndarray], k: int, path: str,
+                 arms: List[str] = None):
+    arms = arms or ["meanfield", "ring", "trained"]
+    fig, axes = plt.subplots(1, len(arms), figsize=(4.4 * len(arms), 3.6),
+                             sharey=True, squeeze=False)
+    for ax, arm in zip(axes[0], arms):
+        for other in arms:
+            if other != arm:
+                ax.plot(t_grid, r_series[other], color="0.85", lw=1.0)
+        ax.plot(t_grid, r_series[arm], color=ARM_COLORS[arm], lw=2.2)
+        ax.set_title(ARM_LABELS[arm])
+        ax.set_xlabel(r"$t$")
+        ax.set_ylim(-0.03, 1.03)
+        ax.grid(alpha=0.25)
+    axes[0][0].set_ylabel(r"order parameter $r(t)$")
+    fig.suptitle(rf"Untwisting the $k={k}$ twisted state", y=1.02)
+    fig.tight_layout()
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_r_combined(t_grid, r_series, k: int, path: str, arms: List[str] = None):
+    fig, ax = plt.subplots(figsize=(6.4, 4.0))
+    for arm in (arms or ["meanfield", "ring", "trained"]):
+        ax.plot(t_grid, r_series[arm], color=ARM_COLORS[arm], lw=2.0, label=ARM_LABELS[arm])
+    ax.set_xlabel(r"$t$")
+    ax.set_ylabel(r"$r(t)$")
+    ax.set_ylim(-0.03, 1.03)
+    ax.grid(alpha=0.25)
+    ax.legend(loc="center right")
+    ax.set_title(rf"Order parameter, $k={k}$ twisted state")
+    fig.tight_layout()
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_winding(t_grid, wind_series: Dict[str, np.ndarray], k: int, path: str,
+                arms: List[str] = None):
+    fig, ax = plt.subplots(figsize=(6.4, 4.0))
+    for arm in (arms or ["meanfield", "ring", "trained"]):
+        ax.plot(t_grid, wind_series[arm], color=ARM_COLORS[arm], lw=2.0,
+                label=ARM_LABELS[arm])
+    for q in range(k + 1):
+        ax.axhline(q, color="0.85", lw=0.8, zorder=0)
+    ax.set_xlabel(r"$t$")
+    ax.set_ylabel(r"winding number $\mathcal{W}(t)$")
+    ax.set_ylim(-0.25, k + 0.25)
+    ax.grid(alpha=0.15)
+    ax.legend()
+    ax.set_title("Untwisting proceeds in integer quanta")
+    fig.tight_layout()
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_variance(t_grid, var_series: Dict[str, np.ndarray], path: str,
+                 arms: List[str] = None):
+    fig, ax = plt.subplots(figsize=(6.4, 4.0))
+    for arm in (arms or ["meanfield", "ring", "trained"]):
+        ax.semilogy(t_grid, np.maximum(var_series[arm], 1e-12),
+                    color=ARM_COLORS[arm], lw=2.0, label=ARM_LABELS[arm])
+    ax.set_xlabel(r"$t$")
+    ax.set_ylabel(r"$\mathrm{Var}(x(t))$")
+    ax.grid(alpha=0.25, which="both")
+    ax.legend()
+    ax.set_title("Variance decay")
+    fig.tight_layout()
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_vxi(t_grid_W, vxi_trained: np.ndarray, vxi_ring: float, path: str):
+    fig, ax = plt.subplots(figsize=(6.4, 4.0))
+    ax.plot(t_grid_W, vxi_trained, color=ARM_COLORS["trained"], lw=2.0,
+            label=r"trained $V_\xi(\hat{w}(t))$")
+    ax.axhline(0.0, color=ARM_COLORS["meanfield"], lw=1.5, ls="--",
+               label=r"mean-field ($V_\xi = 0$)")
+    ax.axhline(vxi_ring, color=ARM_COLORS["ring"], lw=1.5, ls="--",
+               label=rf"smooth ring ($V_\xi = {vxi_ring:.1f}$)")
+    ax.set_xlabel(r"$t$")
+    ax.set_ylabel(r"$V_\xi$")
+    ax.grid(alpha=0.25)
+    ax.legend()
+    ax.set_title("Receiver variation of the kernel (measured, not enforced)")
+    fig.tight_layout()
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_kernel_strip(Ws: torch.Tensor, dt: float, path: str, n_panels: int = 6):
+    NT = Ws.shape[0]
+    idxs = np.linspace(0, NT - 1, n_panels).round().astype(int)
+    wbars = [Ws.shape[-1] * Ws[i].cpu().numpy() for i in idxs]
+    vmin = min(w.min() for w in wbars)
+    vmax = max(w.max() for w in wbars)
+    fig, axes = plt.subplots(1, n_panels, figsize=(2.6 * n_panels, 2.9))
+    for ax, i, w in zip(axes, idxs, wbars):
+        im = ax.imshow(w, origin="lower", extent=(0, 1, 0, 1), aspect="equal",
+                       cmap="viridis", interpolation="bicubic", vmin=vmin, vmax=vmax)
+        ax.set_title(rf"$t = {i * dt:.1f}$", fontsize=10)
+        ax.set_xticks([0, 1])
+        ax.set_yticks([0, 1])
+        if ax is axes[0]:
+            ax.set_ylabel(r"$\xi$")
+        ax.set_xlabel(r"$\zeta$")
+    fig.colorbar(im, ax=axes, fraction=0.02, pad=0.01, label=r"$\hat{w}(t,\xi,\zeta)$")
+    fig.suptitle("Trained kernel over time", y=1.04)
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_state_evolution(t_grid, xi, xs_arms: Dict[str, np.ndarray], k: int, path: str,
+                        arms: List[str] = None):
+    arms = arms or ["meanfield", "ring", "trained"]
+    lim = math.pi * max(k, 1) * 1.1
+    fig, axes = plt.subplots(1, len(arms), figsize=(4.4 * len(arms), 3.6),
+                             sharey=True, squeeze=False)
+    for ax, arm in zip(axes[0], arms):
+        X = xs_arms[arm]  # (NT+1, N)
+        im = ax.imshow(X.T, origin="lower", aspect="auto",
+                       extent=(t_grid[0], t_grid[-1], 0.0, 1.0),
+                       cmap="twilight_shifted", vmin=-lim, vmax=lim)
+        ax.set_title(ARM_LABELS[arm])
+        ax.set_xlabel(r"$t$")
+    axes = axes[0]
+    axes[0].set_ylabel(r"label $\xi$")
+    fig.colorbar(im, ax=axes, fraction=0.02, pad=0.01, label=r"$x(t,\xi)$")
+    fig.suptitle("State evolution", y=1.04)
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def fig_training_history(history: Dict[str, List[float]], path: str):
+    fig, axes = plt.subplots(1, 2, figsize=(11, 3.8))
+    axes[0].plot(history["loss"], lw=1.0, color="tab:blue")
+    axes[0].set_title("total loss")
+    axes[0].set_xlabel("step")
+    axes[0].grid(alpha=0.25)
+    for key, color in [("state_cost", "tab:green"), ("l2_reg_cost", "tab:red"),
+                       ("terminal_var", "tab:purple"), ("terminal_target", "tab:brown")]:
+        axes[1].semilogy(np.maximum(np.array(history[key]), 1e-10), lw=1.0, label=key, color=color)
+    axes[1].set_title("components (log scale)")
+    axes[1].set_xlabel("step")
+    axes[1].grid(alpha=0.25, which="both")
+    axes[1].legend(fontsize=8)
+    fig.tight_layout()
+    fig.savefig(path, bbox_inches="tight")
+    plt.close(fig)
+
+
+def save_untwist_gif(t_grid, xi, xs_arms: Dict[str, np.ndarray],
+                     r_series: Dict[str, np.ndarray], k: int, path: str,
+                     stride: int = 2, fps: int = 12, dpi: int = 80,
+                     arms: List[str] = None):
+    """
+    Animated circle view: each agent is a dot at angle x_i(t) on the unit circle,
+    colored by its label ξ. One panel per arm.
+    """
+    from PIL import Image
+    from io import BytesIO
+
+    arms = arms or ["meanfield", "ring", "trained"]
+    colors = plt.cm.hsv(xi)
+    frames = list(range(0, len(t_grid), stride))
+    if frames[-1] != len(t_grid) - 1:
+        frames.append(len(t_grid) - 1)
+
+    pil_imgs = []
+    theta = np.linspace(0, 2 * np.pi, 200)
+    for n in frames:
+        fig, axes = plt.subplots(1, len(arms), figsize=(3.5 * len(arms), 3.8),
+                                 squeeze=False)
+        for ax, arm in zip(axes[0], arms):
+            x_n = xs_arms[arm][n]  # (N,)
+            ax.plot(np.cos(theta), np.sin(theta), color="0.85", lw=1.0, zorder=1)
+            ax.scatter(np.cos(x_n), np.sin(x_n), c=colors, s=28,
+                       edgecolors="none", zorder=2)
+            ax.set_xlim(-1.25, 1.25)
+            ax.set_ylim(-1.25, 1.25)
+            ax.set_aspect("equal")
+            ax.axis("off")
+            ax.set_title(f"{ARM_LABELS[arm]}\n$r = {r_series[arm][n]:.3f}$", fontsize=10)
+        fig.suptitle(rf"$k={k}$ twist:  $t = {t_grid[n]:.2f}$", y=0.99)
+        fig.tight_layout()
+        buf = BytesIO()
+        fig.savefig(buf, format="png", dpi=dpi)
+        plt.close(fig)
+        buf.seek(0)
+        pil_imgs.append(Image.open(buf).copy())
+
+    pil_imgs[0].save(path, save_all=True, append_images=pil_imgs[1:], loop=0,
+                     duration=int(1000 / fps), optimize=False)
diff --git a/graphon_control/initial_conditions.py b/graphon_control/initial_conditions.py
new file mode 100644
index 0000000..23b6306
--- /dev/null
+++ b/graphon_control/initial_conditions.py
@@ -0,0 +1,72 @@
+"""Reusable initial-state samplers."""
+import math
+from typing import Union
+import torch
+
+
+def sample_x0(
+    batch_size: int,
+    N: int,
+    mean: float,
+    std: float,
+    device: Union[str, torch.device],
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    return mean + std * torch.randn(batch_size, N, device=device, dtype=dtype)
+
+
+def sample_x0_structured(
+    batch_size: int,
+    xi: torch.Tensor,
+    noise_std: float = 0.2,
+) -> torch.Tensor:
+    """
+    xi: (N,)
+    Sinusoidal base profile with additive noise.
+    """
+    base = 1.5 * torch.sin(2 * torch.pi * xi) + 0.8 * torch.cos(4 * torch.pi * xi)
+    return base.unsqueeze(0) + noise_std * torch.randn(batch_size, xi.numel(), device=xi.device, dtype=xi.dtype)
+
+
+def sample_x0_two_clusters(
+    batch_size: int,
+    xi: torch.Tensor,
+    noise_std: float = 0.15,
+) -> torch.Tensor:
+    """
+    xi: (N,)
+    Agents with xi < 0.5 start near -1, the rest near +1.
+    """
+    base = torch.where(xi < 0.5, torch.tensor(-1.0, device=xi.device, dtype=xi.dtype), torch.tensor(1.0, device=xi.device, dtype=xi.dtype))
+    return base.unsqueeze(0) + noise_std * torch.randn(batch_size, xi.numel(), device=xi.device, dtype=xi.dtype)
+
+
+def sample_x0_linear_label(
+    batch_size: int,
+    xi: torch.Tensor,
+    noise_std: float = 0.15,
+) -> torch.Tensor:
+    """
+    xi: (N,)
+    Base profile 20*(ξ - 1/2) with additive Gaussian noise.
+    Agents with higher labels start at higher states.
+    """
+    base = 20.0 * (xi - 0.5)
+    return base.unsqueeze(0) + noise_std * torch.randn(batch_size, xi.numel(), device=xi.device, dtype=xi.dtype)
+
+
+def sample_x0_linear_sine(
+    batch_size: int,
+    xi: torch.Tensor,
+    amplitude: float = 7.0,
+    n_periods: float = 5.0,
+    noise_std: float = 1.0,
+) -> torch.Tensor:
+    """
+    xi: (N,)
+    Base profile: 20*(ξ - 1/2) + amplitude * sin(2π * n_periods * ξ)
+    with additive Gaussian noise. Combines a linear label-dependent slope
+    with multi-oscillation sinusoidal structure.
+    """
+    base = 20.0 * (xi - 0.5) + amplitude * torch.sin(2 * math.pi * n_periods * xi)
+    return base.unsqueeze(0) + noise_std * torch.randn(batch_size, xi.numel(), device=xi.device, dtype=xi.dtype)
diff --git a/graphon_control/interactions.py b/graphon_control/interactions.py
new file mode 100644
index 0000000..e4a614b
--- /dev/null
+++ b/graphon_control/interactions.py
@@ -0,0 +1,24 @@
+"""Reusable interaction laws."""
+import math
+from typing import Union
+import torch
+
+
+def phi_default(y: torch.Tensor, alpha: float = 1.0) -> torch.Tensor:
+    """
+    Default saturating interaction φ(y) = alpha * tanh(y).
+
+    The rollout evaluates φ at x_j - x_i (paper convention, eqn FN definition
+    in main5.tex), so alpha > 0 yields attractive, consensus-seeking drift.
+    """
+    return alpha * torch.tanh(y)
+
+
+def phi_linear(y: torch.Tensor, alpha: float = 1.0) -> torch.Tensor:
+    """
+    Linear interaction φ(y) = alpha * y.
+
+    The rollout evaluates φ at x_j - x_i (paper convention, eqn FN definition
+    in main5.tex), so alpha > 0 yields attractive, consensus-seeking drift.
+    """
+    return alpha * y
diff --git a/graphon_control/objectives.py b/graphon_control/objectives.py
new file mode 100644
index 0000000..e9db29d
--- /dev/null
+++ b/graphon_control/objectives.py
@@ -0,0 +1,179 @@
+"""One discretized objective shared by generic and paper experiment APIs.
+
+The paper coefficient is beta * ||w||², without an additional 1/2. Here
+W=wbar/N, so ||w||² is represented by sum_ij W_ij². The optional uniform
+regularizer retains its separate legacy convention eta_uniform/2.
+"""
+
+from collections.abc import Sequence
+
+import torch
+
+from .dynamics import mean_field_kernel, rollout_dynamics_w1
+from .validation import finite_scalar, n_time_steps
+
+
+def empirical_mean(x: torch.Tensor) -> torch.Tensor:
+    return x.mean(dim=-1)
+
+
+def empirical_variance(x: torch.Tensor) -> torch.Tensor:
+    """Population variance along the agent axis, using a stable centered sum."""
+    return (x - x.mean(dim=-1, keepdim=True)).square().mean(dim=-1)
+
+
+def resolve_target_schedule(T, dt, gamma_target, target_times, target_values,
+                            target_weights) -> list[dict[str, float]]:
+    """Snap checkpoint targets to the nearest rollout time, including endpoints."""
+    NT = n_time_steps(T, dt)
+    finite_scalar(gamma_target, "gamma_target", nonnegative=True)
+    if len(target_times) != len(target_values):
+        raise ValueError("target_times and target_values must have the same length")
+    if len(target_weights) not in (0, len(target_times)):
+        raise ValueError("target_weights must be empty or match target_times")
+    if len(target_times) == 0:
+        return []
+    if len(target_weights) == 0 and gamma_target == 0:
+        raise ValueError("scheduled targets require target_weights or nonzero gamma_target")
+    weights = target_weights if len(target_weights) else [gamma_target] * len(target_times)
+    schedule = []
+    for time, value, weight in zip(target_times, target_values, weights):
+        time = finite_scalar(time, "target time")
+        value = finite_scalar(value, "target value")
+        weight = finite_scalar(weight, "target weight", nonnegative=True)
+        if not 0 <= time <= T:
+            raise ValueError(f"target time {time} must lie in [0, T={T}]")
+        index = min(NT, max(0, round(time / dt)))
+        schedule.append({"index": float(index), "time": index * dt,
+                         "value": value, "weight": weight})
+    return schedule
+
+
+def graphon_l2_penalty(W: torch.Tensor) -> torch.Tensor:
+    """Approximate ||w||²; batched inputs retain their leading axes."""
+    if W.ndim < 2 or W.shape[-2] != W.shape[-1] or W.shape[-1] == 0:
+        raise ValueError("W must have shape (..., N, N), N >= 1")
+    return W.square().sum(dim=(-2, -1))
+
+
+def uniform_graphon_target(N, device=None, dtype=torch.float32):
+    return mean_field_kernel(N, device, dtype)
+
+
+def graphon_uniform_penalty(W: torch.Tensor) -> torch.Tensor:
+    """Approximate ||w-1||² in the quadrature-weight representation."""
+    # Validate shape before constructing the uniform target.
+    graphon_l2_penalty(W)
+    return graphon_l2_penalty(W - uniform_graphon_target(W.shape[-1], W.device, W.dtype))
+
+
+def trajectory_cost(
+    xs: torch.Tensor,
+    Ws: torch.Tensor,
+    *,
+    dt: float,
+    beta: float,
+    lambda_terminal: float = 0.0,
+    gamma_target: float = 0.0,
+    x_target: float = 0.0,
+    target_schedule: Sequence[dict] | None = None,
+    eta_uniform: float = 0.0,
+    regularizer_mode: str = "l2",
+    weighted_running_cost: bool = False,
+) -> dict[str, torch.Tensor]:
+    """Compute tensor-valued components for xs=(NT+1,B,N), Ws=(NT,N,N).
+
+    Running terms use the left endpoint at each Euler step. The optional
+    running-state ramp is (n+1)/NT, matching the former generic trainer.
+    ``terminal_var`` is the unweighted diagnostic; its weighted contribution
+    is ``terminal_var_cost``.
+    """
+    finite_scalar(dt, "dt", positive=True)
+    for name, value in (("beta", beta), ("lambda_terminal", lambda_terminal),
+                        ("gamma_target", gamma_target), ("eta_uniform", eta_uniform)):
+        finite_scalar(value, name, nonnegative=True)
+    finite_scalar(x_target, "x_target")
+    if regularizer_mode not in ("l2", "uniform", "l2_plus_uniform"):
+        raise ValueError("regularizer_mode must be 'l2', 'uniform', or 'l2_plus_uniform'")
+    if xs.ndim != 3 or min(xs.shape) == 0 or xs.shape[0] < 2:
+        raise ValueError("xs must have shape (NT+1, B, N), with NT >= 1")
+    NT, N = xs.shape[0] - 1, xs.shape[-1]
+    if Ws.shape != (NT, N, N):
+        raise ValueError("Ws must have shape (NT, N, N) matching xs")
+    if Ws.device != xs.device or Ws.dtype != xs.dtype:
+        raise ValueError("Ws and xs must have the same device and dtype")
+    zero = xs.new_zeros(())
+    variance = empirical_variance(xs).mean(dim=-1)
+    running_variance = variance[:-1]
+    if weighted_running_cost:
+        running_variance = running_variance * torch.arange(1, NT + 1, device=xs.device,
+                                                           dtype=xs.dtype) / NT
+    state_cost = dt * running_variance.sum()
+    l2_cost = (dt * beta * graphon_l2_penalty(Ws).sum()
+               if regularizer_mode in ("l2", "l2_plus_uniform") else zero)
+    uniform_cost = (dt * eta_uniform / 2 * graphon_uniform_penalty(Ws).sum()
+                   if regularizer_mode in ("uniform", "l2_plus_uniform") else zero)
+    means_T = empirical_mean(xs[-1])
+    target_tracking = (means_T - x_target).square().mean()
+    terminal_var_cost = lambda_terminal * variance[-1]
+    terminal_target_cost = gamma_target * target_tracking
+    scheduled_cost = zero
+    for target in target_schedule or ():
+        index = target["index"]
+        if int(index) != index or not 0 <= index <= NT:
+            raise ValueError("target index must be an integer in [0, NT]")
+        weight = finite_scalar(target["weight"], "target weight", nonnegative=True)
+        value = finite_scalar(target["value"], "target value")
+        tracking = (empirical_mean(xs[int(index)]) - value).square().mean()
+        scheduled_cost = scheduled_cost + weight * tracking
+    terminal_cost = terminal_var_cost + terminal_target_cost
+    total = state_cost + l2_cost + uniform_cost + terminal_cost + scheduled_cost
+    return {
+        "total": total,
+        "state_cost": state_cost,
+        "l2_reg_cost": l2_cost,
+        "uniform_cost": uniform_cost,
+        "scheduled_target_cost": scheduled_cost,
+        "terminal_var": variance[-1],
+        "terminal_mean": means_T.mean(),
+        "target_tracking": target_tracking,
+        "terminal_var_cost": terminal_var_cost,
+        "terminal_target_cost": terminal_target_cost,
+        "terminal_target": terminal_target_cost,
+        "terminal_cost": terminal_cost,
+    }
+
+
+def paper_cost(xs, Ws, cfg) -> dict[str, torch.Tensor]:
+    """Shared paper objective with the experiment scripts' established keys."""
+    parts = trajectory_cost(xs, Ws, dt=cfg.dt, beta=cfg.beta,
+                            lambda_terminal=cfg.lambda_terminal,
+                            gamma_target=cfg.gamma_target, x_target=cfg.x_target)
+    return {key: parts[key] for key in ("total", "state_cost", "l2_reg_cost",
+                                       "terminal_var", "terminal_target", "terminal_mean")}
+
+
+def sampled_cost_JNdt(
+    model, x0, T, dt, sigma, beta, lambda_terminal, phi,
+    gamma_target=0.0, x_target=0.0, target_times=None, target_values=None,
+    target_weights=None, eta_uniform=0.0, regularizer_mode="l2",
+    weighted_running_cost=False,
+):
+    """Legacy generic adapter: return a differentiable loss and float metrics.
+
+    The beta convention now agrees with the manuscript and both experiments.
+    To retain the old generic regularization strength, divide beta by two.
+    RNG streams can differ because zero-noise rollouts no longer draw noise.
+    """
+    schedule = resolve_target_schedule(T, dt, gamma_target, target_times or [],
+                                       target_values or [], target_weights or [])
+    xs, Ws = rollout_dynamics_w1(model, x0, T, dt, sigma, phi)
+    parts = trajectory_cost(
+        torch.stack(xs), torch.stack(Ws), dt=dt, beta=beta,
+        lambda_terminal=lambda_terminal, gamma_target=gamma_target, x_target=x_target,
+        target_schedule=schedule, eta_uniform=eta_uniform, regularizer_mode=regularizer_mode,
+        weighted_running_cost=weighted_running_cost,
+    )
+    aux = {("total_cost" if key == "total" else key): float(value.detach())
+           for key, value in parts.items() if key != "terminal_target"}
+    return parts["total"], aux
diff --git a/graphon_control/plotting.py b/graphon_control/plotting.py
new file mode 100644
index 0000000..a461e3a
--- /dev/null
+++ b/graphon_control/plotting.py
@@ -0,0 +1,508 @@
+"""Plot graphon densities and trajectories; callers convert W to N*W for density displays."""
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+
+
+def xi_grid(N: int, device=None, dtype=torch.float32):
+    """
+    Cell centers in [0,1]:
+        xi_i = (i + 0.5)/N
+    Returns shape (N,).
+    """
+    dev = device if device is not None else "cpu"
+    return (torch.arange(N, device=dev, dtype=dtype) + 0.5) / N
+
+
+def _as_numpy(A):
+    if torch.is_tensor(A):
+        return A.detach().cpu().numpy()
+    return np.asarray(A)
+
+
+def _default_vmin_vmax(Wnp, symmetric=False, q=None):
+    """
+    Choose plotting range.
+    - symmetric=True: use [-m,m]
+    - q: if not None, use percentile clipping, e.g. q=0.01
+    """
+    if q is not None:
+        lo = np.quantile(Wnp, q)
+        hi = np.quantile(Wnp, 1 - q)
+    else:
+        lo = np.min(Wnp)
+        hi = np.max(Wnp)
+
+    if symmetric:
+        m = max(abs(lo), abs(hi))
+        return -m, m
+    return lo, hi
+
+
+def imshow_graphon(
+    W,
+    title="",
+    vmin=None,
+    vmax=None,
+    ax=None,
+    show=True,
+    cmap="viridis",
+    interpolation="bicubic",
+    colorbar=True,
+    symmetric=False,
+    percentile_clip=None,
+):
+    """
+    Nicely display W_{ij} as a graphon approximation w(xi_i, zeta_j).
+
+    x-axis = zeta in [0,1]
+    y-axis = xi   in [0,1]
+
+    Parameters
+    ----------
+    interpolation:
+        'nearest', 'bilinear', 'bicubic', etc.
+        For a smoother graphon-like appearance, 'bicubic' is nice.
+    symmetric:
+        Use symmetric color scale around 0 (useful if entries can be signed).
+    percentile_clip:
+        Example: 0.01 clips 1% low/high tails for better contrast.
+    """
+    Wnp = _as_numpy(W)
+
+    if Wnp.ndim != 2 or Wnp.shape[0] != Wnp.shape[1]:
+        raise ValueError(f"W must be square 2D array, got shape {Wnp.shape}")
+
+    if vmin is None or vmax is None:
+        auto_vmin, auto_vmax = _default_vmin_vmax(
+            Wnp, symmetric=symmetric, q=percentile_clip
+        )
+        if vmin is None:
+            vmin = auto_vmin
+        if vmax is None:
+            vmax = auto_vmax
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(6.5, 5.5))
+
+    im = ax.imshow(
+        Wnp,
+        origin="lower",
+        extent=(0.0, 1.0, 0.0, 1.0),
+        aspect="equal",
+        cmap=cmap,
+        interpolation=interpolation,
+        vmin=vmin,
+        vmax=vmax,
+    )
+
+    ax.set_title(title if title else r"Graphon heatmap")
+    ax.set_xlabel(r"$\zeta$")
+    ax.set_ylabel(r"$\xi$")
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+
+    if colorbar:
+        plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+
+    if show:
+        plt.tight_layout()
+        plt.show()
+
+    return ax
+
+
+
+def plot_slices_over_zeta(
+    W,
+    xis=(0.1, 0.3, 0.5, 0.7, 0.9),
+    title="",
+    ax=None,
+    show=True,
+    linewidth=2.0,
+):
+    """
+    Fix xi and plot zeta -> w(xi, zeta) for a few xi values.
+    """
+    Wnp = _as_numpy(W)
+    N = Wnp.shape[0]
+    zeta = (np.arange(N) + 0.5) / N
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(7, 4.5))
+
+    for xi_val in xis:
+        i = int(np.clip(np.floor(xi_val * N), 0, N - 1))
+        xi_actual = (i + 0.5) / N
+        ax.plot(zeta, Wnp[i, :], linewidth=linewidth, label=rf"$\xi\approx {xi_actual:.2f}$")
+
+    ax.set_title(title if title else r"Slices $\zeta \mapsto w(\xi,\zeta)$")
+    ax.set_xlabel(r"$\zeta$")
+    ax.set_ylabel(r"$w(\xi,\zeta)$")
+    ax.set_xlim(0, 1)
+    ax.grid(alpha=0.25)
+    ax.legend()
+
+    if show:
+        plt.tight_layout()
+        plt.show()
+
+    return ax
+
+
+def plot_slices_over_xi(
+    W,
+    zetas=(0.1, 0.3, 0.5, 0.7, 0.9),
+    title="",
+    ax=None,
+    show=True,
+    linewidth=2.0,
+):
+    """
+    Fix zeta and plot xi -> w(xi, zeta) for a few zeta values.
+    """
+    Wnp = _as_numpy(W)
+    N = Wnp.shape[0]
+    xi = (np.arange(N) + 0.5) / N
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(7, 4.5))
+
+    for zeta_val in zetas:
+        j = int(np.clip(np.floor(zeta_val * N), 0, N - 1))
+        zeta_actual = (j + 0.5) / N
+        ax.plot(xi, Wnp[:, j], linewidth=linewidth, label=rf"$\zeta\approx {zeta_actual:.2f}$")
+
+    ax.set_title(title if title else r"Slices $\xi \mapsto w(\xi,\zeta)$")
+    ax.set_xlabel(r"$\xi$")
+    ax.set_ylabel(r"$w(\xi,\zeta)$")
+    ax.set_xlim(0, 1)
+    ax.grid(alpha=0.25)
+    ax.legend()
+
+    if show:
+        plt.tight_layout()
+        plt.show()
+
+    return ax
+
+
+def plot_graphon_diagnostics(
+    W,
+    title_prefix="",
+    cmap="viridis",
+    interpolation="bicubic",
+    percentile_clip=None,
+    symmetric=False,
+    show=True,
+):
+    """
+    Three-panel diagnostic:
+      1) heatmap
+      2) slices in zeta for several xi
+      3) row/column masses
+
+    Very useful for seeing whether W is:
+      - uniform,
+      - banded,
+      - blocky,
+      - strongly heterogeneous.
+    """
+    Wnp = _as_numpy(W)
+    N = Wnp.shape[0]
+    grid = (np.arange(N) + 0.5) / N
+
+    # W contains graphon density values; integration uses normalized measure.
+    row_mass = Wnp.mean(axis=1)
+    col_mass = Wnp.mean(axis=0)
+
+    if percentile_clip is not None:
+        vmin, vmax = _default_vmin_vmax(Wnp, symmetric=symmetric, q=percentile_clip)
+    else:
+        vmin, vmax = _default_vmin_vmax(Wnp, symmetric=symmetric, q=None)
+
+    fig, axes = plt.subplots(1, 3, figsize=(16, 4.6))
+
+    # Panel 1: heatmap
+    im = axes[0].imshow(
+        Wnp,
+        origin="lower",
+        extent=(0.0, 1.0, 0.0, 1.0),
+        aspect="equal",
+        cmap=cmap,
+        interpolation=interpolation,
+        vmin=vmin,
+        vmax=vmax,
+    )
+    axes[0].set_title(f"{title_prefix}heatmap")
+    axes[0].set_xlabel(r"$\zeta$")
+    axes[0].set_ylabel(r"$\xi$")
+    plt.colorbar(im, ax=axes[0], fraction=0.046, pad=0.04)
+
+    # Panel 2: slices
+    xis = (0.1, 0.3, 0.5, 0.7, 0.9)
+    for xi_val in xis:
+        i = int(np.clip(np.floor(xi_val * N), 0, N - 1))
+        axes[1].plot(grid, Wnp[i, :], linewidth=2, label=rf"$\xi\approx {(i+0.5)/N:.2f}$")
+    axes[1].set_title(f"{title_prefix}slices over $\\zeta$")
+    axes[1].set_xlabel(r"$\zeta$")
+    axes[1].set_ylabel(r"$w(\xi,\zeta)$")
+    axes[1].grid(alpha=0.25)
+    axes[1].legend(fontsize=9)
+
+    # Panel 3: row/column masses
+    axes[2].plot(grid, row_mass, linewidth=2, label="row mass")
+    axes[2].plot(grid, col_mass, linewidth=2, linestyle="--", label="column mass")
+    axes[2].set_title(f"{title_prefix}row/column mass")
+    axes[2].set_xlabel(r"$\xi$ or $\zeta$")
+    axes[2].set_ylabel("mass")
+    axes[2].grid(alpha=0.25)
+    axes[2].legend()
+
+    plt.tight_layout()
+    if show:
+        plt.show()
+
+    return axes
+
+
+def heterogeneity_stats(W):
+    """
+    Scalar summaries for how structured W is.
+    """
+    Wt = W if torch.is_tensor(W) else torch.tensor(W, dtype=torch.float32)
+    eps = 1e-12
+
+    P = (Wt / (Wt.sum(dim=1, keepdim=True).clamp_min(eps))).clamp_min(eps)
+    row_entropy = float((-P * P.log()).sum(dim=1).mean().item())
+
+    N = Wt.shape[0]
+    k = max(1, int(0.05 * N))
+    topk_mass = float(torch.topk(Wt, k=k, dim=1).values.sum(dim=1).mean().item())
+
+    uniform = torch.full_like(Wt, 1.0 / N)
+    row_l2_dev = float(((Wt - uniform).pow(2).sum(dim=1).mean()).sqrt().item())
+
+    return {
+        "mean_row_entropy": row_entropy,
+        "mean_topk_mass_(k=5%N)": topk_mass,
+        "mean_row_L2_dev_from_uniform": row_l2_dev,
+        "min": float(Wt.min().item()),
+        "max": float(Wt.max().item()),
+        "mean": float(Wt.mean().item()),
+    }
+
+
+def frob_norm(W):
+    Wnp = _as_numpy(W)
+    return float(np.linalg.norm(Wnp.reshape(-1), ord=2))
+
+
+def dist_frob(Wa, Wb):
+    A = _as_numpy(Wa)
+    B = _as_numpy(Wb)
+    return float(np.linalg.norm((A - B).reshape(-1), ord=2))
+
+
+@torch.no_grad()
+def sample_graphon_on_fine_grid(model_or_callable, t=None, x=None, M=None, device="cpu"):
+    """Return density values for a controller's own grid or a label callable.
+
+    A TimeLabelGraphonControl stores a fixed grid: this helper evaluates that
+    grid only (M must equal N). It does not claim neural-grid refinement.
+    Other callables must have signature (XI,ZETA), or (t,x,XI,ZETA) when
+    t or x is supplied. Exceptions inside a callable are never swallowed.
+    """
+    from .controls import TimeLabelGraphonControl
+    if isinstance(model_or_callable, TimeLabelGraphonControl):
+        model = model_or_callable
+        if M is not None and M != model.N:
+            raise ValueError("controller evaluation requires M=N; arbitrary grids are unsupported")
+        weights = model(0.0 if t is None else t)
+        if weights.ndim != 2:
+            raise ValueError("plotting requires one scalar time")
+        return (model.N * weights).detach().cpu()
+    M = 200 if M is None else M
+    if isinstance(M, bool) or not isinstance(M, int) or M < 1:
+        raise ValueError("M must be a positive integer")
+    grid = (torch.arange(M, device=device, dtype=torch.float32) + 0.5) / M
+    XI, ZETA = torch.meshgrid(grid, grid, indexing="ij")
+    values = (model_or_callable(XI, ZETA) if t is None and x is None
+              else model_or_callable(t, x, XI, ZETA))
+    if values.shape != (M, M):
+        raise ValueError("kernel callable must return (M,M) density values")
+    return values.detach().cpu()
+
+
+@torch.no_grad()
+def plot_state_histogram_over_time_imshow(
+    xs,
+    dt: float,
+    bins=80,
+    value_range=None,
+    density=True,
+    cmap="viridis",
+    title="Empirical state histogram over time",
+    x_target=None,
+    show=True,
+):
+    X = torch.stack(xs, dim=0).detach().cpu().numpy()   # (Tsteps, B, N)
+    Tsteps = X.shape[0]
+
+    flat_all = X.reshape(-1)
+    if value_range is None:
+        xmin = float(flat_all.min())
+        xmax = float(flat_all.max())
+        pad = 0.05 * max(1e-8, xmax - xmin)
+        xmin -= pad
+        xmax += pad
+        value_range = (xmin, xmax)
+
+    x_edges = np.linspace(value_range[0], value_range[1], bins + 1)
+    x_centers = 0.5 * (x_edges[:-1] + x_edges[1:])
+
+    H = np.zeros((Tsteps, bins), dtype=float)
+    for n in range(Tsteps):
+        vals = X[n].reshape(-1)
+        hist, _ = np.histogram(vals, bins=x_edges, density=density)
+        H[n, :] = hist
+
+    _, ax = plt.subplots(1, 1, figsize=(8, 5))
+    im = ax.imshow(
+        H,
+        origin="lower",
+        aspect="auto",
+        extent=(x_centers[0], x_centers[-1], 0.0, dt * (Tsteps - 1)),
+        cmap=cmap,
+    )
+    ax.set_xlabel("state x")
+    ax.set_ylabel("time t")
+    ax.set_title(title)
+    plt.colorbar(im, ax=ax, label="density" if density else "count")
+
+    # Overlay: mean state trajectory
+    t_axis = np.linspace(0.0, dt * (Tsteps - 1), Tsteps)
+    mean_vals = X.reshape(Tsteps, -1).mean(axis=1)
+    ax.plot(mean_vals, t_axis, color="white", lw=1.5, label=r"$\bar{x}(t)$")
+
+    # Overlay: target state
+    if x_target is not None:
+        ax.axvline(x_target, color="white", lw=1.0, ls="--",
+                   label=rf"$x^* = {x_target}$")
+
+    ax.legend(loc="upper right", fontsize=8, framealpha=0.4)
+
+    if show:
+        plt.show()
+
+    return H, x_edges, ax
+
+
+def plot_x0_samples(xi, x0_batch, n_show=12, title="Sample initial-condition profiles"):
+    """
+    Line plot of n_show samples from x0_batch (B, N) vs xi (N,).
+    Returns fig, ax.
+    """
+    fig, ax = plt.subplots(figsize=(10, 4))
+    for i in range(min(n_show, x0_batch.shape[0])):
+        ax.plot(xi.cpu().numpy(), x0_batch[i].detach().cpu().numpy(), lw=0.8)
+    ax.set_xlabel(r"$\xi$")
+    ax.set_ylabel(r"$x_0(\xi)$")
+    ax.set_title(title)
+    fig.tight_layout()
+    return fig, ax
+
+
+def plot_x0_density(xi, x0_batch, gridsize=40, title="Initial-state density across labels and batch"):
+    """
+    Hexbin density of (xi, x0) pooled over all batch samples.
+    x0_batch shape: (B, N). xi shape: (N,).
+    Returns fig, ax.
+    """
+    xi_np = xi.cpu().numpy()
+    x0_np = x0_batch.detach().cpu().numpy()
+    B, N = x0_np.shape
+    xi_rep = np.tile(xi_np, B)
+    x0_rep = x0_np.reshape(-1)
+
+    fig, ax = plt.subplots(figsize=(10, 4))
+    hb = ax.hexbin(xi_rep, x0_rep, gridsize=gridsize, cmap="Blues", mincnt=1, linewidths=0.2)
+    fig.colorbar(hb, ax=ax, label="count")
+    ax.set_xlabel(r"$\xi$")
+    ax.set_ylabel(r"$x_0$")
+    ax.set_title(title)
+    fig.tight_layout()
+    return fig, ax
+
+
+def plot_phi(phi_fn, z_range=(-20.0, 20.0), n_points=400, title=None):
+    """
+    Plot phi_fn(z) vs z over z_range. Returns fig, ax.
+    """
+    import torch as _torch
+    zz = np.linspace(z_range[0], z_range[1], n_points)
+    z_t = _torch.tensor(zz, dtype=_torch.float32)
+    with _torch.no_grad():
+        yy = phi_fn(z_t).numpy()
+    fig, ax = plt.subplots(figsize=(5, 4))
+    ax.plot(zz, yy, lw=1.5)
+    ax.axhline(0, color="k", lw=0.5, ls="--")
+    ax.axvline(0, color="k", lw=0.5, ls="--")
+    ax.set_xlabel(r"$z$")
+    ax.set_ylabel(r"$\phi(z)$")
+    ax.set_title(title or r"Interaction kernel $\phi(z)$")
+    fig.tight_layout()
+    return fig, ax
+
+
+def save_graphon_gif(frames, path, fps=10, cmap="viridis", percentile_clip=0.01,
+                     t_vals=None, figsize=(5, 4), dpi=80):
+    """
+    frames : list of (N, N) tensors, one per time point.
+    t_vals : optional list of floats (same length as frames) used for per-frame title.
+    Renders each frame as a proper matplotlib figure (axes, colorbar, labels)
+    and saves as an animated GIF via Pillow.
+    """
+    from PIL import Image
+    from io import BytesIO
+
+    # Global color scale computed once for visual consistency across frames
+    all_vals = np.concatenate([_as_numpy(f).ravel() for f in frames])
+    q = (percentile_clip * 100) if percentile_clip else 0.0
+    vmin = float(np.percentile(all_vals, q))
+    vmax = float(np.percentile(all_vals, 100.0 - q))
+
+    pil_imgs = []
+    for i, W in enumerate(frames):
+        t_str = f"$t = {t_vals[i]:.2f}$" if t_vals is not None else f"frame {i}"
+        title = r"Graphon $w(\xi,\zeta)$,  " + t_str
+
+        fig, ax = plt.subplots(figsize=figsize)
+        imshow_graphon(
+            W,
+            title=title,
+            cmap=cmap,
+            interpolation="bicubic",
+            percentile_clip=None,   # use pre-computed vmin/vmax
+            vmin=vmin,
+            vmax=vmax,
+            ax=ax,
+            show=False,
+        )
+        fig.tight_layout()
+
+        buf = BytesIO()
+        fig.savefig(buf, format="png", dpi=dpi)
+        plt.close(fig)
+        buf.seek(0)
+        pil_imgs.append(Image.open(buf).copy())
+
+    duration_ms = int(1000 / fps)
+    pil_imgs[0].save(
+        path,
+        save_all=True,
+        append_images=pil_imgs[1:],
+        loop=0,
+        duration=duration_ms,
+        optimize=False,
+    )
\ No newline at end of file
diff --git a/graphon_control/training.py b/graphon_control/training.py
new file mode 100644
index 0000000..6a1d4c9
--- /dev/null
+++ b/graphon_control/training.py
@@ -0,0 +1,76 @@
+"""Shared optimization loop with finite loss and gradient checks."""
+
+import math
+from collections.abc import Callable, Mapping
+from numbers import Real
+
+import torch
+from tqdm.auto import tqdm
+
+from .validation import finite_scalar, positive_integer
+
+
+def _metric_value(value, name: str) -> float:
+    if isinstance(value, torch.Tensor):
+        if value.numel() != 1:
+            raise ValueError(f"objective component {name!r} must be scalar")
+        value = float(value.detach())
+    if not isinstance(value, Real) or not math.isfinite(value):
+        raise FloatingPointError(f"objective component {name!r} is not finite: {value!r}")
+    return float(value)
+
+
+def train_controller(model, objective: Callable[[], Mapping], num_steps: int, lr: float,
+                     grad_clip: float | None = 1.0, description: str = "training"):
+    """Optimize a fresh stochastic objective each iteration with Adam.
+
+    ``objective`` returns a mapping whose ``total`` is a differentiable scalar
+    tensor. Other values may be scalar tensors or numeric metrics. History
+    records ``total`` under ``loss`` and preserves the remaining component keys.
+    Nonfinite values raise before an optimizer update can corrupt parameters.
+    """
+    positive_integer(num_steps, "num_steps")
+    finite_scalar(lr, "lr", positive=True)
+    if grad_clip is not None:
+        finite_scalar(grad_clip, "grad_clip", positive=True)
+    parameters = [parameter for parameter in model.parameters() if parameter.requires_grad]
+    if not parameters:
+        raise ValueError("model must have trainable parameters")
+    optimizer = torch.optim.Adam(parameters, lr=lr)
+    history = {"loss": []}
+    expected_keys = None
+    last_parts = {}
+    for step in tqdm(range(num_steps), desc=description, unit="step"):
+        model.train()
+        optimizer.zero_grad(set_to_none=True)
+        parts = objective()
+        if not isinstance(parts, Mapping) or "total" not in parts:
+            raise ValueError("objective must return a mapping containing 'total'")
+        total = parts["total"]
+        if not isinstance(total, torch.Tensor) or total.ndim != 0 or not total.requires_grad:
+            raise ValueError("objective 'total' must be a differentiable scalar tensor")
+        if "loss" in parts:
+            raise ValueError("objective component 'loss' is reserved for total history")
+        if expected_keys is None:
+            expected_keys = set(parts)
+            history.update({key: [] for key in parts if key != "total"})
+        elif set(parts) != expected_keys:
+            raise ValueError("objective component keys changed during training")
+        last_parts = {key: _metric_value(value, key) for key, value in parts.items()}
+        total.backward()
+        gradients = [p.grad for p in parameters if p.grad is not None]
+        if not gradients:
+            raise ValueError("objective is disconnected from the controller parameters")
+        if any(not torch.isfinite(gradient).all() for gradient in gradients):
+            optimizer.zero_grad(set_to_none=True)
+            raise FloatingPointError(f"nonfinite controller gradient at optimization step {step}")
+        if grad_clip is not None:
+            torch.nn.utils.clip_grad_norm_(parameters, grad_clip, error_if_nonfinite=True)
+        optimizer.step()
+        if any(not torch.isfinite(parameter).all() for parameter in parameters):
+            raise FloatingPointError(f"nonfinite controller parameter at optimization step {step}")
+        history["loss"].append(last_parts["total"])
+        for key in parts:
+            if key != "total":
+                history[key].append(last_parts[key])
+    return history, last_parts
diff --git a/graphon_control/validation.py b/graphon_control/validation.py
new file mode 100644
index 0000000..31094b5
--- /dev/null
+++ b/graphon_control/validation.py
@@ -0,0 +1,83 @@
+"""Validation shared by the generic trainer and paper experiments."""
+
+import math
+from numbers import Integral, Real
+
+
+def positive_integer(value, name: str) -> int:
+    """Require an integer count, rejecting booleans and silent truncation."""
+    if isinstance(value, bool) or not isinstance(value, Integral) or value < 1:
+        raise ValueError(f"{name} must be a positive integer; got {value!r}")
+    return int(value)
+
+
+def finite_scalar(value, name: str, *, positive=False, nonnegative=False) -> float:
+    if isinstance(value, bool) or not isinstance(value, Real) or not math.isfinite(value):
+        raise ValueError(f"{name} must be finite; got {value!r}")
+    if positive and value <= 0:
+        raise ValueError(f"{name} must be positive; got {value!r}")
+    if nonnegative and value < 0:
+        raise ValueError(f"{name} must be nonnegative; got {value!r}")
+    return float(value)
+
+
+def n_time_steps(T: float, dt: float) -> int:
+    """Return T/dt only when it is a positive integer (up to roundoff).
+
+    A partial final step is not implemented. Rejecting it prevents a rollout
+    from silently stopping before or after its advertised horizon.
+    """
+    T = finite_scalar(T, "T", positive=True)
+    dt = finite_scalar(dt, "dt", positive=True)
+    ratio = T / dt
+    if not math.isfinite(ratio):
+        raise ValueError("T/dt must be a finite positive integer")
+    NT = round(ratio)
+    if NT < 1 or not math.isclose(ratio, NT, rel_tol=0.0, abs_tol=1e-8):
+        raise ValueError(f"T/dt must be a positive integer; got T={T}, dt={dt}")
+    return NT
+
+
+def validate_common_config(cfg) -> None:
+    """Validate common dataclass fields without coupling to one experiment."""
+    n_time_steps(cfg.T, cfg.dt)
+    for name in ("N", "batch_size", "num_steps", "num_fourier_freqs", "embed_dim",
+                 "hidden_embed", "hidden_score", "time_chunk"):
+        if hasattr(cfg, name):
+            positive_integer(getattr(cfg, name), name)
+    for name in ("lr", "time_scale"):
+        if hasattr(cfg, name):
+            if name == "time_scale" and cfg.time_scale is None:
+                continue  # None explicitly selects the experiment horizon T.
+            finite_scalar(getattr(cfg, name), name, positive=True)
+    for name in ("sigma", "beta", "lambda_terminal", "gamma_target", "eta_uniform", "jitter"):
+        if hasattr(cfg, name):
+            finite_scalar(getattr(cfg, name), name, nonnegative=True)
+    if getattr(cfg, "score_bound", None) is not None:
+        finite_scalar(cfg.score_bound, "score_bound", nonnegative=True)
+    if hasattr(cfg, "interaction_alpha"):
+        finite_scalar(cfg.interaction_alpha, "interaction_alpha")
+    if getattr(cfg, "grad_clip", None) is not None:
+        finite_scalar(cfg.grad_clip, "grad_clip", positive=True)
+    if hasattr(cfg, "x_target"):
+        finite_scalar(cfg.x_target, "x_target")
+    for name in ("seed", "evaluation_seed"):
+        if hasattr(cfg, name):
+            seed = getattr(cfg, name)
+            if isinstance(seed, bool) or not isinstance(seed, Integral) or seed < 0:
+                raise ValueError(f"{name} must be a nonnegative integer")
+    if hasattr(cfg, "regularizer_mode") and cfg.regularizer_mode not in (
+        "l2", "uniform", "l2_plus_uniform"
+    ):
+        raise ValueError("regularizer_mode must be 'l2', 'uniform', or 'l2_plus_uniform'")
+    if hasattr(cfg, "positive_map") and cfg.positive_map not in ("softplus", "exp"):
+        raise ValueError("positive_map must be 'softplus' or 'exp'")
+    if any(hasattr(cfg, key) for key in ("target_times", "target_values", "target_weights")):
+        # Local import avoids a module import cycle.
+        from .objectives import resolve_target_schedule
+
+        resolve_target_schedule(
+            cfg.T, cfg.dt, getattr(cfg, "gamma_target", 0.0),
+            getattr(cfg, "target_times", []), getattr(cfg, "target_values", []),
+            getattr(cfg, "target_weights", []),
+        )
diff --git a/model_1.py b/model_1.py
index 25ff639..b5cf34a 100644
--- a/model_1.py
+++ b/model_1.py
@@ -1,1078 +1,7 @@
-# model_1.py
-
-import dataclasses
-import datetime
-import inspect
-import json
-import math
-import os
-from dataclasses import dataclass
-from typing import Callable, Dict, List, Optional, Union
-
-# PyTorch for modeling and training
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-# import vizualization utilities from model_1_viz.py
-from tqdm import tqdm
-from model_1_viz import imshow_graphon, plot_graphon_diagnostics, heterogeneity_stats, plot_state_histogram_over_time_imshow, plot_x0_samples, plot_x0_density, plot_phi, save_graphon_gif
-
-
-# ============================================================
-# Random run-tag generator
-# ============================================================
-
-_TAG_ADJECTIVES = [
-    "amber", "arctic", "azure", "bold", "bright", "calm", "cardinal",
-    "cobalt", "coral", "crimson", "crystal", "dark", "dawn", "dense", "dim",
-    "dusty", "dusk", "ember", "faint", "firm", "flat", "fluid", "foggy",
-    "frozen", "gilded", "golden", "green", "grey", "hollow", "hushed", "icy",
-    "iron", "jade", "keen", "light", "lunar", "muted", "narrow", "noble",
-    "ochre", "pale", "quiet", "rapid", "rigid", "rough", "ruby", "russet",
-    "sandy", "scarlet", "sharp", "silent", "silver", "sleek", "slow", "solar",
-    "sparse", "stark", "stern", "still", "stone", "stormy", "subtle", "swift",
-    "tawny", "teal", "thin", "tidal", "tiny", "topaz", "turbid", "twilight",
-    "vast", "velvet", "vivid", "warm", "white", "wild", "windy", "wiry",
-]
-
-_TAG_NOUNS = [
-    "basin", "beacon", "bloom", "boulder", "brook", "canopy", "canyon",
-    "cliff", "cloud", "coast", "comet", "coral", "crater", "creek", "crest",
-    "delta", "dune", "eddy", "fern", "field", "fjord", "flare", "flint",
-    "fog", "forest", "frost", "gale", "gap", "geyser", "glacier", "glen",
-    "gorge", "grove", "gulf", "heath", "inlet", "island", "kelp", "knoll",
-    "lagoon", "larch", "ledge", "marsh", "meadow", "mesa", "moor", "moss",
-    "nebula", "peak", "pine", "plain", "pond", "pool", "prism", "ravine",
-    "reef", "ridge", "rift", "river", "rock", "shore", "slope", "spring",
-    "storm", "stream", "summit", "swamp", "tide", "timber", "torrent", "trail",
-    "vale", "valley", "veil", "void", "wave", "wisp",
-]
-
-
-def _random_run_tag() -> str:
-    import random
-    return f"{random.choice(_TAG_ADJECTIVES)}_{random.choice(_TAG_NOUNS)}"
-
-
-# ============================================================
-# Configuration
-# ============================================================
-
-@dataclass
-class TrainConfig:
-    # discretization/noise parameters
-    N: int = 64                     # number of agents / label grid points
-    T: float = 1.0                  # total time horizon 
-    dt: float = 0.02                # time step for Euler-Maruyama simulation and cost discretization
-    sigma: float = 0.15             # noise scale in dynamics
-
-    # cost parameters
-    beta: float = 1e-2              # control regularization coefficient in J^{N,Δt}
-    eta_uniform: float = 0.0        # coefficient for uniform-graphon regularizer
-    regularizer_mode: str = "l2"    # one of: "l2", "uniform", "l2_plus_uniform"
-    lambda_terminal: float = 1.0    # weight on terminal variance penalty
-    gamma_target: float = 0.0       # weight on terminal mean-tracking penalty
-    x_target: float = 0.0           # target for the terminal empirical mean
-    target_times: List[float] = dataclasses.field(default_factory=list)    # additional checkpoint times in [0, T]
-    target_values: List[float] = dataclasses.field(default_factory=list)   # target mean value at each checkpoint time
-    target_weights: List[float] = dataclasses.field(default_factory=list)  # optional weights for checkpoint targets; defaults to gamma_target if omitted
-
-    # graphon architecture
-    enforce_symmetry: bool = False   # if True, enforce W_ij = W_ji exactly
-    positive_map: str = "softplus"   # nonnegativity map for symmetric mode: "softplus" or "exp"
-
-    # Monte Carlo / optimization
-    batch_size: int = 32
-    num_steps: int = 1000
-    lr: float = 1e-3
-    grad_clip: Optional[float] = 1.0
-
-    # network architecture
-    num_fourier_freqs: int = 8    # m in the draft
-    embed_dim: int = 64           # d in the draft
-    hidden_embed: int = 128
-    hidden_score: int = 128
-    time_scale: float = 1.0       # rescales t before feeding to score net
-
-    # loss weighting
-    weighted_running_cost: bool = False  # if True, linearly ramp running state cost weight from ~0 to 1 over [0, T]
-
-    # interaction scaling
-    interaction_alpha: float = 1.0      # linear interaction gain; positive is attractive (consensus-seeking) under the paper convention φ(x_j - x_i)
-
-    # reproducibility / device
-    seed: int = 0
-    device: str = "cuda" if torch.cuda.is_available() else "cpu"
-
-
-# ============================================================
-# Utility: interaction nonlinearity φ
-# ============================================================
-
-def phi_default(y: torch.Tensor, alpha: float = 1.0) -> torch.Tensor:
-    """
-    Default saturating interaction φ(y) = alpha * tanh(y).
-
-    The rollout evaluates φ at x_j - x_i (paper convention, eqn FN definition
-    in main5.tex), so alpha > 0 yields attractive, consensus-seeking drift.
-    """
-    return alpha * torch.tanh(y)
-
-
-def phi_linear(y: torch.Tensor, alpha: float = 1.0) -> torch.Tensor:
-    """
-    Linear interaction φ(y) = alpha * y.
-
-    The rollout evaluates φ at x_j - x_i (paper convention, eqn FN definition
-    in main5.tex), so alpha > 0 yields attractive, consensus-seeking drift.
-    """
-    return alpha * y
-
-
-# ============================================================
-# Fourier features for labels ξ in [0,1]
-# ψ(ξ) = (sin(2π 2^k ξ), cos(2π 2^k ξ))_{k=0}^{m-1}
-# ============================================================
-
-class FourierFeatures(nn.Module):
-    def __init__(self, num_freqs: int):
-        super().__init__()
-        self.num_freqs = num_freqs
-        self.freqs: torch.Tensor
-        freqs = 2.0 ** torch.arange(num_freqs, dtype=torch.float32)
-        self.register_buffer("freqs", freqs)
-
-    @property
-    def out_dim(self) -> int:
-        return 2 * self.num_freqs
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        x: (...,) in [0,1]
-        returns: (..., 2*num_freqs)
-        """
-        x = x.unsqueeze(-1)  # (..., 1)
-        angles = 2.0 * math.pi * x * self.freqs  # (..., m)
-        return torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1)
-
-
-# ============================================================
-# Simple MLP
-# ============================================================
-
-class MLP(nn.Module):
-    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, depth: int = 2):
-        super().__init__()
-        layers = []
-        d = in_dim
-        for _ in range(depth):
-            layers.append(nn.Linear(d, hidden_dim))
-            layers.append(nn.ReLU())
-            d = hidden_dim
-        layers.append(nn.Linear(d, out_dim))
-        self.net = nn.Sequential(*layers)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.net(x)
-
-
-# ============================================================
-# Graphon control architecture (Model 1: time + label)
-#
-# enforce_symmetry=False (directed):
-#   receiver h_i = φθ(ψ(ξ_i)),  sender h̃_j = φ̃θ(ψ(ξ_j))
-#   score s_ij(t) = gθ(t, h_i, h̃_j),  W = row-softmax (self-loops permitted,
-#   matching the class U_N without the zero-diagonal constraint in main5.tex)
-#
-# enforce_symmetry=True (undirected):
-#   shared h_i = φθ(ψ(ξ_i))
-#   A_ij = gθ(t, h_i, h_j),  S = 0.5*(A + Aᵀ),  W = positive_map(S) / avg_row_sum
-# ============================================================
-
-class TimeLabelGraphonControl(nn.Module):
-    def __init__(
-        self,
-        xi: torch.Tensor,
-        num_fourier_freqs: int,
-        embed_dim: int,
-        hidden_embed: int,
-        hidden_score: int,
-        time_scale: float = 1.0,
-        enforce_symmetry: bool = False,
-        positive_map: str = "softplus",
-    ):
-        super().__init__()
-        self.xi: torch.Tensor
-        self.psi_xi_cached: torch.Tensor
-        self.register_buffer("xi", xi)          # (N,)
-        self.N = xi.numel()
-        self.time_scale = time_scale
-        self.enforce_symmetry = enforce_symmetry
-        self.positive_map = positive_map
-
-        self.ff = FourierFeatures(num_fourier_freqs)
-        ff_dim = self.ff.out_dim
-
-        # Symmetric mode:  single shared embedding  (h_i = h_j basis, symmetry by construction)
-        # Asymmetric mode: separate receiver/sender embeddings (directed attention, row-softmax)
-        self.embed_net:    Optional[MLP] = MLP(ff_dim, hidden_embed, embed_dim, depth=2) if enforce_symmetry else None
-        self.receiver_net: Optional[MLP] = None if enforce_symmetry else MLP(ff_dim, hidden_embed, embed_dim, depth=2)
-        self.sender_net:   Optional[MLP] = None if enforce_symmetry else MLP(ff_dim, hidden_embed, embed_dim, depth=2)
-
-        # Score net gθ : R^{1 + 2d} -> R  (shared across both modes)
-        self.score_net = MLP(1 + 2 * embed_dim, hidden_score, 1, depth=2)
-
-        # Precompute label features / embeddings at the fixed label grid
-        psi_xi = self.ff(self.xi)                          # (N, 2m)
-        self.register_buffer("psi_xi_cached", psi_xi)
-
-    def forward(self, t: torch.Tensor) -> torch.Tensor:
-        """
-        t: scalar tensor or shape (B,)
-        returns: (N, N) if scalar t, (B, N, N) if batched t
-
-        enforce_symmetry=False  →  directed, row-stochastic  (each row sums to 1)
-        enforce_symmetry=True   →  undirected, symmetric,    average row sum = 1
-        """
-        N = self.N
-
-        if t.ndim == 0:
-            t = t.unsqueeze(0)
-
-        B = t.shape[0]
-        t_in = (t / self.time_scale).reshape(B, 1, 1, 1).expand(B, N, N, 1)
-
-        if self.enforce_symmetry:
-            assert self.embed_net is not None
-            h = self.embed_net(self.psi_xi_cached)              # (N, d)
-            d = h.shape[-1]
-            h_i = h.view(1, N, 1, d).expand(B, N, N, d)
-            h_j = h.view(1, 1, N, d).expand(B, N, N, d)
-
-            A = self.score_net(torch.cat([t_in, h_i, h_j], dim=-1)).squeeze(-1)  # (B, N, N)
-            A = 0.5 * (A + A.transpose(-1, -2))                 # exact symmetry
-
-            W = torch.exp(A.clamp(max=20.0)) if self.positive_map == "exp" else F.softplus(A)
-
-            # Normalize by avg row sum — scalar per batch element, preserves symmetry
-            avg_row_sum = W.sum(dim=-1).mean(dim=-1, keepdim=True).unsqueeze(-1).clamp_min(1e-8)
-            W = W / avg_row_sum
-
-        else:
-            assert self.receiver_net is not None and self.sender_net is not None
-            receiver_h = self.receiver_net(self.psi_xi_cached)  # (N, d)
-            sender_h   = self.sender_net(self.psi_xi_cached)    # (N, d)
-            d = receiver_h.shape[-1]
-            h_i = receiver_h.view(1, N, 1, d).expand(B, N, N, d)
-            h_j = sender_h.view(1, 1, N, d).expand(B, N, N, d)
-
-            scores = self.score_net(torch.cat([t_in, h_i, h_j], dim=-1)).squeeze(-1)  # (B, N, N)
-            W = F.softmax(scores, dim=-1)
-
-        if W.shape[0] == 1:
-            return W[0]
-        return W
-
-
-# ============================================================
-# Simulation of dynamics (5)
-# x_i^{n+1} = x_i^n + Δt (1/N) Σ_j W_ij(t_n) φ(x_j^n - x_i^n) + σ ΔB_i^n
-# ============================================================
-
-def rollout_dynamics_w1(
-    model: TimeLabelGraphonControl,
-    x0: torch.Tensor,
-    T: float,
-    dt: float,
-    sigma: float,
-    phi: Callable[[torch.Tensor], torch.Tensor],
-):
-    """
-    x0: (B, N)
-    returns:
-        xs: list of length NT+1, each entry (B, N)
-        Ws: list of length NT, each entry (N, N)   # W1 depends only on time
-    """
-    device = x0.device
-    B, N = x0.shape
-    NT = int(round(T / dt))
-
-    x = x0
-    xs = [x]
-    Ws = []
-
-    sqrt_dt = math.sqrt(dt)
-
-    for n in range(NT):
-        t_n = torch.tensor(n * dt, device=device, dtype=x0.dtype)
-        W = model(t_n)  # (N, N)
-        Ws.append(W)
-
-        # pairwise differences diff[b, i, j] = x_j - x_i, shape (B, N, N)
-        # (paper convention: drift_i = Σ_j W_ij φ(x_j - x_i))
-        diff = x.unsqueeze(1).expand(B, N, N) - x.unsqueeze(2).expand(B, N, N)
-
-        interaction = phi(diff)                      # (B, N, N)
-        drift = (W.unsqueeze(0) * interaction).sum(dim=-1)          # (B, N)
-
-        noise = sigma * sqrt_dt * torch.randn(B, N, device=device, dtype=x0.dtype)
-        x = x + dt * drift + noise
-        xs.append(x)
-
-    return xs, Ws
-
-
-# ============================================================
-# Cost J^{N,Δt}(θ)
-# Σ_n Δt [ Var(x^n) + (β/2)||W(t_n)||^2 ]
-# with Var(x) = (1/N)Σ_k x_k^2 - ( (1/N)Σ_k x_k )^2
-# ============================================================
-
-def empirical_variance(x: torch.Tensor) -> torch.Tensor:
-    """
-    x: (B, N)
-    returns: (B,)
-    """
-    mean = x.mean(dim=-1)
-    second_moment = (x ** 2).mean(dim=-1)
-    return second_moment - mean ** 2
-
-
-def empirical_mean(x: torch.Tensor) -> torch.Tensor:
-    """
-    x: (B, N)
-    returns: (B,)
-    """
-    return x.mean(dim=-1)
-
-
-def resolve_target_schedule(
-    T: float,
-    dt: float,
-    gamma_target: float,
-    target_times: List[float],
-    target_values: List[float],
-    target_weights: List[float],
-) -> List[Dict[str, float]]:
-    """
-    Resolve user-specified checkpoint targets onto the discrete simulation grid.
-
-    Each target time is mapped to the nearest rollout index via round(t / dt).
-    """
-    if len(target_times) != len(target_values):
-        raise ValueError(
-            "target_times and target_values must have the same length; "
-            f"got {len(target_times)} and {len(target_values)}"
-        )
-
-    if target_weights and len(target_weights) != len(target_times):
-        raise ValueError(
-            "target_weights must be empty or have the same length as target_times; "
-            f"got {len(target_weights)} and {len(target_times)}"
-        )
-
-    if not target_times:
-        return []
-
-    resolved_weights = target_weights
-    if not resolved_weights:
-        if gamma_target == 0.0:
-            raise ValueError(
-                "target_weights is empty and gamma_target is 0.0, so scheduled targets would carry no weight. "
-                "Set target_weights explicitly or use a nonzero gamma_target."
-            )
-        resolved_weights = [gamma_target] * len(target_times)
-
-    NT = int(round(T / dt))
-    schedule: List[Dict[str, float]] = []
-    for target_time, target_value, target_weight in zip(target_times, target_values, resolved_weights):
-        if target_time < 0.0 or target_time > T:
-            raise ValueError(f"target time {target_time} must lie in [0, T={T}]")
-
-        target_index = int(round(target_time / dt))
-        target_index = max(0, min(target_index, NT))
-        snapped_time = target_index * dt
-        schedule.append(
-            {
-                "index": float(target_index),
-                "time": snapped_time,
-                "value": float(target_value),
-                "weight": float(target_weight),
-            }
-        )
-
-    return schedule
-
-
-def graphon_l2_penalty(W: torch.Tensor) -> torch.Tensor:
-    """
-    Discrete approximation of ||w||^2_{L^2(I^2)} when W is a row-stochastic
-    weight matrix representing quadrature weights rather than raw graphon values.
-    Since w_ij ~ N * W_ij, we use sum(W^2), not mean(W^2).
-    """
-    return (W ** 2).sum()
-
-
-def uniform_graphon_target(
-    N: int,
-    device: torch.device,
-    dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    """
-    Return the discrete uniform graphon compatible with the row-stochastic
-    architecture used here (self-loops permitted; there is no zero-diagonal
-    constraint, matching the class U_N in main5.tex).
-
-    The learned W is row-stochastic, so the uniform target is U_ij = 1/N for
-    all entries — the row-stochastic representation of the constant graphon
-    w == 1 — NOT the all-ones matrix, which would violate row normalization.
-    """
-    return torch.full((N, N), 1.0 / N, device=device, dtype=dtype)
-
-
-def graphon_uniform_penalty(W: torch.Tensor) -> torch.Tensor:
-    """
-    Discrete approximation of ||w - u||^2_{L^2(I^2)} in the row-stochastic
-    representation. The correct scale is sum((W - U)^2), not mean((W - U)^2).
-    """
-    N = W.shape[0]
-    U = uniform_graphon_target(N, device=W.device, dtype=W.dtype)
-    return ((W - U) ** 2).sum()
-
-
-def sampled_cost_JNdt(
-    model: TimeLabelGraphonControl,
-    x0: torch.Tensor,
-    T: float,
-    dt: float,
-    sigma: float,
-    beta: float,
-    lambda_terminal: float,
-    phi: Callable[[torch.Tensor], torch.Tensor],
-    gamma_target: float = 0.0,
-    x_target: float = 0.0,
-    target_times: Optional[List[float]] = None,
-    target_values: Optional[List[float]] = None,
-    target_weights: Optional[List[float]] = None,
-    eta_uniform: float = 0.0,
-    regularizer_mode: str = "l2",
-    weighted_running_cost: bool = False,
-):
-    """
-    Monte Carlo sampled objective corresponding to Algorithm 1.
-    x0: (B, N)
-
-    regularizer_mode controls which graphon penalties are active:
-      "l2"             — only (beta/2)||W||^2 dt
-      "uniform"        — only (eta_uniform/2)||W - U||^2 dt
-      "l2_plus_uniform"— both penalties together
-
-    returns:
-        loss: scalar
-        aux: dict with diagnostics
-    """
-    if regularizer_mode not in ("l2", "uniform", "l2_plus_uniform"):
-        raise ValueError(
-            f"regularizer_mode must be 'l2', 'uniform', or 'l2_plus_uniform'; "
-            f"got '{regularizer_mode}'"
-        )
-
-    xs, Ws = rollout_dynamics_w1(model, x0, T=T, dt=dt, sigma=sigma, phi=phi)
-
-    device = x0.device
-    dtype = x0.dtype
-    running_state_cost:   torch.Tensor = torch.zeros((), device=device, dtype=dtype)
-    running_l2_reg_cost: torch.Tensor = torch.zeros((), device=device, dtype=dtype)
-    running_uniform_cost: torch.Tensor = torch.zeros((), device=device, dtype=dtype)
-    scheduled_target_cost: torch.Tensor = torch.zeros((), device=device, dtype=dtype)
-
-    resolved_schedule = resolve_target_schedule(
-        T=T,
-        dt=dt,
-        gamma_target=gamma_target,
-        target_times=target_times or [],
-        target_values=target_values or [],
-        target_weights=target_weights or [],
-    )
-
-    NT = len(Ws)
-    for n in range(NT):
-        x_n = xs[n]           # (B, N)
-        W_n = Ws[n]           # (N, N)
-
-        var_n         = empirical_variance(x_n).mean()
-        l2_reg_n      = graphon_l2_penalty(W_n)
-        uniform_reg_n = graphon_uniform_penalty(W_n)
-
-        w_n = (n + 1) * dt / T if weighted_running_cost else 1.0
-        running_state_cost = running_state_cost + dt * w_n * var_n
-
-        if regularizer_mode == "l2":
-            running_l2_reg_cost = running_l2_reg_cost + dt * (beta / 2.0) * l2_reg_n
-        elif regularizer_mode == "uniform":
-            running_uniform_cost = running_uniform_cost + dt * (eta_uniform / 2.0) * uniform_reg_n
-        else:  # "l2_plus_uniform"
-            running_l2_reg_cost = running_l2_reg_cost + dt * (beta / 2.0) * l2_reg_n
-            running_uniform_cost = running_uniform_cost + dt * (eta_uniform / 2.0) * uniform_reg_n
-
-    x_T = xs[-1]                                   # (B, N)
-    terminal_var = empirical_variance(x_T).mean()
-
-    mean_T = empirical_mean(x_T)                   # (B,)
-    target_tracking = ((mean_T - x_target) ** 2).mean()
-
-    for target_spec in resolved_schedule:
-        target_index = int(target_spec["index"])
-        mean_at_target = empirical_mean(xs[target_index])
-        scheduled_target_cost = scheduled_target_cost + target_spec["weight"] * ((mean_at_target - target_spec["value"]) ** 2).mean()
-
-    terminal_var_cost = lambda_terminal * terminal_var
-    terminal_target_cost = gamma_target * target_tracking
-    terminal_cost = terminal_var_cost + terminal_target_cost
-
-    total = running_state_cost + running_l2_reg_cost + running_uniform_cost + terminal_cost + scheduled_target_cost
-
-    aux = {
-        "total_cost": float(total.detach().cpu()),
-        "state_cost": float(running_state_cost.detach().cpu()),
-        "l2_reg_cost": float(running_l2_reg_cost.detach().cpu()),
-        "uniform_cost": float(running_uniform_cost.detach().cpu()),
-        "scheduled_target_cost": float(scheduled_target_cost.detach().cpu()),
-        "terminal_var": float(terminal_var.detach().cpu()),
-        "terminal_mean": float(mean_T.mean().detach().cpu()),
-        "target_tracking": float(target_tracking.detach().cpu()),
-        "terminal_var_cost": float(terminal_var_cost.detach().cpu()),
-        "terminal_target_cost": float(terminal_target_cost.detach().cpu()),
-        "terminal_cost": float(terminal_cost.detach().cpu()),
-    }
-    return total, aux
-
-
-# ============================================================
-# Sampling initial conditions
-# ============================================================
-
-def sample_x0(
-    batch_size: int,
-    N: int,
-    mean: float,
-    std: float,
-    device: Union[str, torch.device],
-    dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    return mean + std * torch.randn(batch_size, N, device=device, dtype=dtype)
-
-
-def sample_x0_structured(
-    batch_size: int,
-    xi: torch.Tensor,
-    noise_std: float = 0.2,
-) -> torch.Tensor:
-    """
-    xi: (N,)
-    Sinusoidal base profile with additive noise.
-    """
-    base = 1.5 * torch.sin(2 * torch.pi * xi) + 0.8 * torch.cos(4 * torch.pi * xi)
-    return base.unsqueeze(0) + noise_std * torch.randn(batch_size, xi.numel(), device=xi.device)
-
-
-def sample_x0_two_clusters(
-    batch_size: int,
-    xi: torch.Tensor,
-    noise_std: float = 0.15,
-) -> torch.Tensor:
-    """
-    xi: (N,)
-    Agents with xi < 0.5 start near -1, the rest near +1.
-    """
-    base = torch.where(xi < 0.5, torch.tensor(-1.0, device=xi.device), torch.tensor(1.0, device=xi.device))
-    return base.unsqueeze(0) + noise_std * torch.randn(batch_size, xi.numel(), device=xi.device)
-
-
-def sample_x0_linear_label(
-    batch_size: int,
-    xi: torch.Tensor,
-    noise_std: float = 0.15,
-) -> torch.Tensor:
-    """
-    xi: (N,)
-    Base profile 20*(ξ - 1/2) with additive Gaussian noise.
-    Agents with higher labels start at higher states.
-    """
-    base = 20.0 * (xi - 0.5)
-    return base.unsqueeze(0) + noise_std * torch.randn(batch_size, xi.numel(), device=xi.device)
-
-
-def sample_x0_linear_sine(
-    batch_size: int,
-    xi: torch.Tensor,
-    amplitude: float = 7.0,
-    n_periods: float = 5.0,
-    noise_std: float = 1.0,
-) -> torch.Tensor:
-    """
-    xi: (N,)
-    Base profile: 20*(ξ - 1/2) + amplitude * sin(2π * n_periods * ξ)
-    with additive Gaussian noise. Combines a linear label-dependent slope
-    with multi-oscillation sinusoidal structure.
-    """
-    base = 20.0 * (xi - 0.5) + amplitude * torch.sin(2 * math.pi * n_periods * xi)
-    return base.unsqueeze(0) + noise_std * torch.randn(batch_size, xi.numel(), device=xi.device, dtype=xi.dtype)
-
-
-# ============================================================
-# Training loop = Algorithm 1
-# Pick sample S = (X0, ΔW), compute sampled cost, backprop, update θ
-# ============================================================
-
-def train_algorithm1_w1(
-    config: TrainConfig,
-    phi: Callable[[torch.Tensor], torch.Tensor] = phi_default,
-    x0_sampler: Callable[[int, torch.Tensor], torch.Tensor] = sample_x0_two_clusters,
-):
-    torch.manual_seed(config.seed)
-
-    device = torch.device(config.device)
-    dtype = torch.float32
-
-    N = config.N
-
-    # label discretization ξ_i = (i + 1/2)/N
-    xi = (torch.arange(N, device=device, dtype=dtype) + 0.5) / N
-
-    model = TimeLabelGraphonControl(
-        xi=xi,
-        num_fourier_freqs=config.num_fourier_freqs,
-        embed_dim=config.embed_dim,
-        hidden_embed=config.hidden_embed,
-        hidden_score=config.hidden_score,
-        time_scale=max(config.T, 1e-8),
-        enforce_symmetry=config.enforce_symmetry,
-        positive_map=config.positive_map,
-    ).to(device)
-
-    optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
-
-    history: Dict[str, List[float]] = {
-        "loss": [],
-        "state_cost": [],
-        "l2_reg_cost": [],
-        "uniform_cost": [],
-        "scheduled_target_cost": [],
-        "terminal_var": [],
-        "terminal_cost": [],
-        "terminal_mean": [],
-        "target_tracking": [],
-        "terminal_var_cost": [],
-        "terminal_target_cost": [],
-    }
-
-    pbar = tqdm(range(config.num_steps), desc="training", unit="step")
-    for _ in pbar:
-        model.train()
-        optimizer.zero_grad()
-
-        # New Monte Carlo sample S each iteration:
-        # initial states and Brownian increments are generated inside sampled_cost_JNdt
-        x0 = x0_sampler(config.batch_size, xi)
-
-        loss, aux = sampled_cost_JNdt(
-            model=model,
-            x0=x0,
-            T=config.T,
-            dt=config.dt,
-            sigma=config.sigma,
-            beta=config.beta,
-            lambda_terminal=config.lambda_terminal,
-            phi=phi,
-            gamma_target=config.gamma_target,
-            x_target=config.x_target,
-            target_times=config.target_times,
-            target_values=config.target_values,
-            target_weights=config.target_weights,
-            eta_uniform=config.eta_uniform,
-            regularizer_mode=config.regularizer_mode,
-            weighted_running_cost=config.weighted_running_cost,
-        )
-
-        loss.backward()
-
-        if config.grad_clip is not None:
-            torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
-
-        optimizer.step()
-
-        history["loss"].append(aux["total_cost"])
-        history["state_cost"].append(aux["state_cost"])
-        history["l2_reg_cost"].append(aux["l2_reg_cost"])
-        history["uniform_cost"].append(aux["uniform_cost"])
-        history["scheduled_target_cost"].append(aux["scheduled_target_cost"])
-        history["terminal_var"].append(aux["terminal_var"])
-        history["terminal_cost"].append(aux["terminal_cost"])
-        history["terminal_mean"].append(aux["terminal_mean"])
-        history["target_tracking"].append(aux["target_tracking"])
-        history["terminal_var_cost"].append(aux["terminal_var_cost"])
-        history["terminal_target_cost"].append(aux["terminal_target_cost"])
-        last_aux = aux
-
-        pbar.set_postfix(
-            loss=f"{aux['total_cost']:.4f}",
-            run=f"{aux['state_cost']:.4f}",
-            term_var=f"{aux['terminal_var']:.4f}",
-            term_mean=f"{aux['terminal_mean']:.4f}",
-            target=f"{aux['terminal_target_cost'] + aux['scheduled_target_cost']:.4f}",
-            reg=f"{aux['l2_reg_cost'] + aux['uniform_cost']:.4f}",
-        )
-
-    return model, history, xi, last_aux
-
-
-# ============================================================
-# Run persistence: save all artifacts to out_root/run_<timestamp>/
-# ============================================================
-
-def save_run(
-    out_root: str,
-    config: TrainConfig,
-    history: Dict[str, List[float]],
-    final_aux: dict,
-    phi: Callable,
-    x0_sampler: Callable,
-    x0_sample: torch.Tensor,
-    run_tag: str = "",
-) -> str:
-    """
-    Save config, history, final diagnostics, and a human-readable summary
-    to a timestamped directory under out_root.
-
-    x0_sample: (B, N) — one representative initial batch for statistics.
-    Returns the path of the created run directory.
-    """
-    ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    dir_name = f"run_{ts}" + (f"_{run_tag}" if run_tag else "")
-    run_dir = os.path.join(out_root, dir_name)
-    os.makedirs(run_dir, exist_ok=True)
-
-    # config.json
-    with open(os.path.join(run_dir, "config.json"), "w") as f:
-        json.dump(dataclasses.asdict(config), f, indent=2)
-
-    # history.json
-    with open(os.path.join(run_dir, "history.json"), "w") as f:
-        json.dump(history, f, indent=2)
-
-    # final_aux.json
-    with open(os.path.join(run_dir, "final_aux.json"), "w") as f:
-        json.dump(final_aux, f, indent=2)
-
-    # x0 statistics
-    x0_np = x0_sample.detach().cpu()
-    x0_stats = {
-        "mean": float(x0_np.mean()),
-        "std":  float(x0_np.std()),
-        "min":  float(x0_np.min()),
-        "max":  float(x0_np.max()),
-    }
-
-    # phi source (best-effort)
-    try:
-        phi_source = inspect.getsource(phi)
-    except (OSError, TypeError):
-        phi_source = "<source unavailable>"
-
-    # run_summary.txt
-    cfg_d = dataclasses.asdict(config)
-    lines = [
-        "=== Run Summary ===",
-        f"timestamp : {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
-        f"run_dir   : {run_dir}",
-        "",
-        "--- Config ---",
-    ]
-    for k, v in cfg_d.items():
-        lines.append(f"  {k} = {v}")
-    lines += [
-        "",
-        "--- Interaction function ---",
-        f"  phi : {getattr(phi, '__name__', repr(phi))}",
-        "  source:",
-    ]
-    for src_line in phi_source.splitlines():
-        lines.append(f"    {src_line}")
-    lines += [
-        "",
-        "--- Initial sampler ---",
-        f"  x0_sampler : {getattr(x0_sampler, '__name__', repr(x0_sampler))}",
-        "",
-        "--- x0 statistics (first batch) ---",
-        f"  mean  : {x0_stats['mean']:+.6f}",
-        f"  std   :  {x0_stats['std']:.6f}",
-        f"  min   : {x0_stats['min']:+.6f}",
-        f"  max   : {x0_stats['max']:+.6f}",
-        "",
-        "--- Final diagnostics ---",
-    ]
-    for k, v in final_aux.items():
-        lines.append(f"  {k:<24s}: {v:.6f}")
-
-    with open(os.path.join(run_dir, "run_summary.txt"), "w") as f:
-        f.write("\n".join(lines) + "\n")
-
-    return run_dir
-
-
-# ============================================================
-# run_experiment: train + save + evaluate in one call
-# ============================================================
-
-def run_experiment(
-    config: TrainConfig,
-    phi: Callable[[torch.Tensor], torch.Tensor] = phi_default,
-    x0_sampler: Callable[[int, torch.Tensor], torch.Tensor] = sample_x0_two_clusters,
-    out_root: str = "out_1",
-    run_tag: str = "",
-    eval_batch_size: int = 128,
-) -> dict:
-    """
-    Train the model, save all artifacts, evaluate the policy, and return
-    a results dict with keys: model, history, xi, xs, Ws, run_dir.
-    """
-    import matplotlib
-    matplotlib.use("Agg")   # non-interactive backend for file saving
-    import matplotlib.pyplot as plt
-
-    print("=" * 60)
-    print(f"  N={config.N}  T={config.T}  dt={config.dt}  sigma={config.sigma}")
-    print(f"  beta={config.beta}  eta_uniform={config.eta_uniform}  reg={config.regularizer_mode}")
-    print(f"  batch={config.batch_size}  steps={config.num_steps}  lr={config.lr}")
-    print(f"  lambda_terminal={config.lambda_terminal}  gamma_target={config.gamma_target}  x_target={config.x_target}")
-    print(f"  scheduled_targets={len(config.target_times)}")
-    print(f"  phi={getattr(phi, '__name__', repr(phi))}")
-    print(f"  x0_sampler={getattr(x0_sampler, '__name__', repr(x0_sampler))}")
-    print(f"  tag={run_tag!r}")
-    print("=" * 60)
-
-    model, history, xi, final_aux = train_algorithm1_w1(config, phi=phi, x0_sampler=x0_sampler)
-
-    # Representative x0 batch for statistics
-    device = next(model.parameters()).device
-    xi_cpu = xi.to(device)
-    x0_sample = x0_sampler(config.batch_size, xi_cpu)
-
-    run_dir = save_run(
-        out_root=out_root,
-        config=config,
-        history=history,
-        final_aux=final_aux,
-        phi=phi,
-        x0_sampler=x0_sampler,
-        x0_sample=x0_sample,
-        run_tag=run_tag,
-    )
-
-    xs, Ws = evaluate_policy(
-        model=model,
-        xi=xi,
-        T=config.T,
-        dt=config.dt,
-        sigma=config.sigma,
-        batch_size=eval_batch_size,
-        phi=phi,
-        x0_sampler=x0_sampler,
-    )
-
-    NT_eval = len(Ws)
-    time_indices = {
-        "early": NT_eval // 4,
-        "mid":   NT_eval // 2,
-        "late":  3 * NT_eval // 4,
-    }
-
-    # Save 3-panel graphon diagnostic plots and collect W stats
-    eval_lines = [
-        "=== Evaluation Diagnostics ===",
-        f"eval_batch_size : {eval_batch_size}",
-        f"NT_eval         : {NT_eval}",
-        "",
-        "--- Initial / terminal variance ---",
-        f"  init_var : {float(empirical_variance(xs[0]).mean()):.6f}",
-        f"  term_var : {float(empirical_variance(xs[-1]).mean()):.6f}",
-        "",
-    ]
-    resolved_schedule = resolve_target_schedule(
-        T=config.T,
-        dt=config.dt,
-        gamma_target=config.gamma_target,
-        target_times=config.target_times,
-        target_values=config.target_values,
-        target_weights=config.target_weights,
-    )
-    if resolved_schedule:
-        eval_lines.append("--- Scheduled mean targets ---")
-        for target_spec in resolved_schedule:
-            target_index = int(target_spec["index"])
-            realized_mean = float(empirical_mean(xs[target_index]).mean())
-            eval_lines.append(
-                f"  t={target_spec['time']:.4f}  idx={target_index:4d}  target={target_spec['value']:+.6f}  realized_mean={realized_mean:+.6f}  weight={target_spec['weight']:.6f}"
-            )
-        eval_lines.append("")
-    for label, idx in time_indices.items():
-        W = Ws[idx].cpu()
-        stats = heterogeneity_stats(W)
-        row_sums = W.sum(dim=-1)
-        sym_err = float((W - W.T).abs().max())
-
-        eval_lines += [
-            f"--- W at {label} time (step {idx}/{NT_eval}) ---",
-            f"  max_symmetry_error : {sym_err:.2e}",
-            f"  mean_row_sum       : {float(row_sums.mean()):.6f}",
-            f"  min_row_sum        : {float(row_sums.min()):.6f}",
-            f"  max_row_sum        : {float(row_sums.max()):.6f}",
-        ]
-        for sk, sv in stats.items():
-            eval_lines.append(f"  {sk:<30s}: {sv:.6f}")
-        eval_lines.append("")
-
-        # 3-panel diagnostic figure
-        plot_graphon_diagnostics(
-            W,
-            title_prefix=f"{label}-time ",
-            interpolation="bicubic",
-            percentile_clip=0.01,
-            show=False,
-        )
-        fig_diag = plt.gcf()
-        fig_diag.savefig(os.path.join(run_dir, f"graphon_diag_{label}.png"), dpi=150)
-        plt.close(fig_diag)
-
-    with open(os.path.join(run_dir, "eval_diagnostics.txt"), "w") as f:
-        f.write("\n".join(eval_lines) + "\n")
-
-    # Save mid-time graphon heatmap
-    W_mid = Ws[time_indices["mid"]].cpu()
-    fig, ax = plt.subplots(figsize=(7, 6))
-    imshow_graphon(
-        W_mid,
-        title=r"Learned graphon $w(\xi,\zeta)$ at mid-time",
-        interpolation="bicubic",
-        cmap="viridis",
-        percentile_clip=0.01,
-        ax=ax,
-        show=False,
-    )
-    fig.tight_layout()
-    fig.savefig(os.path.join(run_dir, "graphon_mid.png"), dpi=150)
-    plt.close(fig)
-
-    # Save state histogram over time
-    plot_state_histogram_over_time_imshow(xs, dt=config.dt, show=False,
-                                           x_target=config.x_target)
-    fig_hist = plt.gcf()
-    fig_hist.savefig(os.path.join(run_dir, "state_histogram.png"), dpi=150)
-    plt.close(fig_hist)
-
-    # Save IC sample-profile and density plots
-    fig_ic_s, _ = plot_x0_samples(xi_cpu, x0_sample)
-    fig_ic_s.savefig(os.path.join(run_dir, "ic_samples.png"), dpi=150)
-    plt.close(fig_ic_s)
-
-    fig_ic_d, _ = plot_x0_density(xi_cpu, x0_sample)
-    fig_ic_d.savefig(os.path.join(run_dir, "ic_density.png"), dpi=150)
-    plt.close(fig_ic_d)
-
-    # Save phi plot
-    fig_phi, _ = plot_phi(phi)
-    fig_phi.savefig(os.path.join(run_dir, "phi_plot.png"), dpi=150)
-    plt.close(fig_phi)
-
-    # Save graphon temporal GIF: sweep t from 0 to T using the final trained model
-    dtype = torch.float32
-    with torch.no_grad():
-        t_vals_tensor = torch.linspace(0.0, config.T, steps=60, device=device, dtype=dtype)
-        gif_frames = [model(t).cpu() for t in t_vals_tensor]
-    save_graphon_gif(gif_frames, os.path.join(run_dir, "graphon_time.gif"),
-                     fps=10, t_vals=t_vals_tensor.cpu().tolist())
-
-    print(f"Results saved to: {run_dir}")
-    return {"model": model, "history": history, "xi": xi, "xs": xs, "Ws": Ws, "run_dir": run_dir}
-
-
-# ============================================================
-# Evaluation helper
-# ============================================================
-
-@torch.no_grad()
-def evaluate_policy(
-    model: TimeLabelGraphonControl,
-    xi: torch.Tensor,
-    T: float,
-    dt: float,
-    sigma: float,
-    batch_size: int,
-    phi: Callable[[torch.Tensor], torch.Tensor] = phi_default,
-    x0_sampler: Callable[[int, torch.Tensor], torch.Tensor] = sample_x0_two_clusters,
-):
-    x0 = x0_sampler(batch_size, xi)
-    xs, Ws = rollout_dynamics_w1(model, x0, T=T, dt=dt, sigma=sigma, phi=phi)
-    return xs, Ws
-
-
-# ============================================================
-# Example usage
-# ============================================================
-
-if __name__ == "__main__":
-    from functools import partial
-
-    # Define the training configuration for the experiment.
-    cfg = TrainConfig(
-        # discretization / noise
-        N=32,                   # number of agents (discretization size for graphon)
-        T=10.0,                  # total time horizon
-        dt=0.02,                # time step for Euler-Maruyama simulation and cost discretization
-        sigma=0.0,             # noise scale in dynamics
-
-        # target 
-        x_target=5.0,           # target state for terminal mean-tracking cost
-        target_times=[5.0],   # user-specified target times for intermediate checkpoints (must lie in [0, T])
-        target_values=[-5.0], # corresponding target mean values at the specified target times
-        target_weights=[], # weights for the scheduled target costs; if empty, defaults to gamma_target for all targets
-        # cost parameters
-        lambda_terminal=1.0,   # terminal variance cost weight
-        gamma_target=100.0,     # target state cost weight
-        regularizer_mode="l2",  # regularization mode ("l2", "uniform", or "l2_plus_uniform")
-        eta_uniform=10.0,       # uniform cost weight
-        beta=10.0,              # control regularization coefficient in J^{N,Δt}
-        
-        # loss weighting
-        weighted_running_cost = False,  # if True, linearly ramp running state cost weight from ~0 to 1 over [0, T]
-
-        # interaction scaling
-        interaction_alpha=1.0,    # linear interaction gain; positive is attractive (consensus-seeking) under the paper convention φ(x_j - x_i)
-
-        # NN parameters
-        num_steps=500,          # training steps
-        batch_size=64,          # Monte Carlo samples per step
-        lr=1e-3,                # learning rate
-        num_fourier_freqs=8,    # number of Fourier features for label embedding
-        embed_dim=64,           # embedding dimension for graphon control architecture
-        hidden_embed=128,       # hidden dimension for embedding MLPs
-        hidden_score=128,       # hidden dimension for score MLP
-        enforce_symmetry=False, # whether to enforce W_ij = W_ji (undirected graphon) or allow asymmetry (directed graphon)
-    )
-
-    # Run the experiment with the specified configuration, interaction function, and initial state sampler.
-    results = run_experiment(
-        config=cfg,
-        phi=partial(phi_linear, alpha=cfg.interaction_alpha),
-        x0_sampler=sample_x0_linear_sine,
-        out_root="out_1",
-        run_tag=_random_run_tag(),
-    )
-
-    # Example of running a second experiment with a different phi and sampler:
-    # from functools import partial
-    # run_experiment(cfg, phi=lambda y: y, x0_sampler=sample_x0_structured,
-    #                out_root="out_1", run_tag="linear_structured")
+"""Compatibility entry point for the reusable graphon_control package."""
+from graphon_control.experiments.generic import *  # noqa: F401,F403
+from graphon_control.experiments.generic import main
+from graphon_control.artifacts import generated_run_tag as _random_run_tag
 
+if __name__ == '__main__':
+    main()
diff --git a/model_1_viz.py b/model_1_viz.py
index 6a6b1c3..0572aee 100644
--- a/model_1_viz.py
+++ b/model_1_viz.py
@@ -1,498 +1,2 @@
-import numpy as np
-import torch
-import matplotlib.pyplot as plt
-
-
-def xi_grid(N: int, device=None, dtype=torch.float32):
-    """
-    Cell centers in [0,1]:
-        xi_i = (i + 0.5)/N
-    Returns shape (N,).
-    """
-    dev = device if device is not None else "cpu"
-    return (torch.arange(N, device=dev, dtype=dtype) + 0.5) / N
-
-
-def _as_numpy(A):
-    if torch.is_tensor(A):
-        return A.detach().cpu().numpy()
-    return np.asarray(A)
-
-
-def _default_vmin_vmax(Wnp, symmetric=False, q=None):
-    """
-    Choose plotting range.
-    - symmetric=True: use [-m,m]
-    - q: if not None, use percentile clipping, e.g. q=0.01
-    """
-    if q is not None:
-        lo = np.quantile(Wnp, q)
-        hi = np.quantile(Wnp, 1 - q)
-    else:
-        lo = np.min(Wnp)
-        hi = np.max(Wnp)
-
-    if symmetric:
-        m = max(abs(lo), abs(hi))
-        return -m, m
-    return lo, hi
-
-
-def imshow_graphon(
-    W,
-    title="",
-    vmin=None,
-    vmax=None,
-    ax=None,
-    show=True,
-    cmap="viridis",
-    interpolation="bicubic",
-    colorbar=True,
-    symmetric=False,
-    percentile_clip=None,
-):
-    """
-    Nicely display W_{ij} as a graphon approximation w(xi_i, zeta_j).
-
-    x-axis = zeta in [0,1]
-    y-axis = xi   in [0,1]
-
-    Parameters
-    ----------
-    interpolation:
-        'nearest', 'bilinear', 'bicubic', etc.
-        For a smoother graphon-like appearance, 'bicubic' is nice.
-    symmetric:
-        Use symmetric color scale around 0 (useful if entries can be signed).
-    percentile_clip:
-        Example: 0.01 clips 1% low/high tails for better contrast.
-    """
-    Wnp = _as_numpy(W)
-
-    if Wnp.ndim != 2 or Wnp.shape[0] != Wnp.shape[1]:
-        raise ValueError(f"W must be square 2D array, got shape {Wnp.shape}")
-
-    if vmin is None or vmax is None:
-        auto_vmin, auto_vmax = _default_vmin_vmax(
-            Wnp, symmetric=symmetric, q=percentile_clip
-        )
-        if vmin is None:
-            vmin = auto_vmin
-        if vmax is None:
-            vmax = auto_vmax
-
-    if ax is None:
-        fig, ax = plt.subplots(figsize=(6.5, 5.5))
-
-    im = ax.imshow(
-        Wnp,
-        origin="lower",
-        extent=(0.0, 1.0, 0.0, 1.0),
-        aspect="equal",
-        cmap=cmap,
-        interpolation=interpolation,
-        vmin=vmin,
-        vmax=vmax,
-    )
-
-    ax.set_title(title if title else r"Graphon heatmap")
-    ax.set_xlabel(r"$\zeta$")
-    ax.set_ylabel(r"$\xi$")
-    ax.set_xlim(0, 1)
-    ax.set_ylim(0, 1)
-
-    if colorbar:
-        plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
-
-    if show:
-        plt.tight_layout()
-        plt.show()
-
-    return ax
-
-
-
-def plot_slices_over_zeta(
-    W,
-    xis=(0.1, 0.3, 0.5, 0.7, 0.9),
-    title="",
-    ax=None,
-    show=True,
-    linewidth=2.0,
-):
-    """
-    Fix xi and plot zeta -> w(xi, zeta) for a few xi values.
-    """
-    Wnp = _as_numpy(W)
-    N = Wnp.shape[0]
-    zeta = (np.arange(N) + 0.5) / N
-
-    if ax is None:
-        fig, ax = plt.subplots(figsize=(7, 4.5))
-
-    for xi_val in xis:
-        i = int(np.clip(np.floor(xi_val * N), 0, N - 1))
-        xi_actual = (i + 0.5) / N
-        ax.plot(zeta, Wnp[i, :], linewidth=linewidth, label=rf"$\xi\approx {xi_actual:.2f}$")
-
-    ax.set_title(title if title else r"Slices $\zeta \mapsto w(\xi,\zeta)$")
-    ax.set_xlabel(r"$\zeta$")
-    ax.set_ylabel(r"$w(\xi,\zeta)$")
-    ax.set_xlim(0, 1)
-    ax.grid(alpha=0.25)
-    ax.legend()
-
-    if show:
-        plt.tight_layout()
-        plt.show()
-
-    return ax
-
-
-def plot_slices_over_xi(
-    W,
-    zetas=(0.1, 0.3, 0.5, 0.7, 0.9),
-    title="",
-    ax=None,
-    show=True,
-    linewidth=2.0,
-):
-    """
-    Fix zeta and plot xi -> w(xi, zeta) for a few zeta values.
-    """
-    Wnp = _as_numpy(W)
-    N = Wnp.shape[0]
-    xi = (np.arange(N) + 0.5) / N
-
-    if ax is None:
-        fig, ax = plt.subplots(figsize=(7, 4.5))
-
-    for zeta_val in zetas:
-        j = int(np.clip(np.floor(zeta_val * N), 0, N - 1))
-        zeta_actual = (j + 0.5) / N
-        ax.plot(xi, Wnp[:, j], linewidth=linewidth, label=rf"$\zeta\approx {zeta_actual:.2f}$")
-
-    ax.set_title(title if title else r"Slices $\xi \mapsto w(\xi,\zeta)$")
-    ax.set_xlabel(r"$\xi$")
-    ax.set_ylabel(r"$w(\xi,\zeta)$")
-    ax.set_xlim(0, 1)
-    ax.grid(alpha=0.25)
-    ax.legend()
-
-    if show:
-        plt.tight_layout()
-        plt.show()
-
-    return ax
-
-
-def plot_graphon_diagnostics(
-    W,
-    title_prefix="",
-    cmap="viridis",
-    interpolation="bicubic",
-    percentile_clip=None,
-    symmetric=False,
-    show=True,
-):
-    """
-    Three-panel diagnostic:
-      1) heatmap
-      2) slices in zeta for several xi
-      3) row/column masses
-
-    Very useful for seeing whether W is:
-      - uniform,
-      - banded,
-      - blocky,
-      - strongly heterogeneous.
-    """
-    Wnp = _as_numpy(W)
-    N = Wnp.shape[0]
-    grid = (np.arange(N) + 0.5) / N
-
-    row_mass = Wnp.sum(axis=1)
-    col_mass = Wnp.sum(axis=0)
-
-    if percentile_clip is not None:
-        vmin, vmax = _default_vmin_vmax(Wnp, symmetric=symmetric, q=percentile_clip)
-    else:
-        vmin, vmax = _default_vmin_vmax(Wnp, symmetric=symmetric, q=None)
-
-    fig, axes = plt.subplots(1, 3, figsize=(16, 4.6))
-
-    # Panel 1: heatmap
-    im = axes[0].imshow(
-        Wnp,
-        origin="lower",
-        extent=(0.0, 1.0, 0.0, 1.0),
-        aspect="equal",
-        cmap=cmap,
-        interpolation=interpolation,
-        vmin=vmin,
-        vmax=vmax,
-    )
-    axes[0].set_title(f"{title_prefix}heatmap")
-    axes[0].set_xlabel(r"$\zeta$")
-    axes[0].set_ylabel(r"$\xi$")
-    plt.colorbar(im, ax=axes[0], fraction=0.046, pad=0.04)
-
-    # Panel 2: slices
-    xis = (0.1, 0.3, 0.5, 0.7, 0.9)
-    for xi_val in xis:
-        i = int(np.clip(np.floor(xi_val * N), 0, N - 1))
-        axes[1].plot(grid, Wnp[i, :], linewidth=2, label=rf"$\xi\approx {(i+0.5)/N:.2f}$")
-    axes[1].set_title(f"{title_prefix}slices over $\\zeta$")
-    axes[1].set_xlabel(r"$\zeta$")
-    axes[1].set_ylabel(r"$w(\xi,\zeta)$")
-    axes[1].grid(alpha=0.25)
-    axes[1].legend(fontsize=9)
-
-    # Panel 3: row/column masses
-    axes[2].plot(grid, row_mass, linewidth=2, label="row mass")
-    axes[2].plot(grid, col_mass, linewidth=2, linestyle="--", label="column mass")
-    axes[2].set_title(f"{title_prefix}row/column mass")
-    axes[2].set_xlabel(r"$\xi$ or $\zeta$")
-    axes[2].set_ylabel("mass")
-    axes[2].grid(alpha=0.25)
-    axes[2].legend()
-
-    plt.tight_layout()
-    if show:
-        plt.show()
-
-    return axes
-
-
-def heterogeneity_stats(W):
-    """
-    Scalar summaries for how structured W is.
-    """
-    Wt = W if torch.is_tensor(W) else torch.tensor(W, dtype=torch.float32)
-    eps = 1e-12
-
-    P = (Wt / (Wt.sum(dim=1, keepdim=True).clamp_min(eps))).clamp_min(eps)
-    row_entropy = float((-P * P.log()).sum(dim=1).mean().item())
-
-    N = Wt.shape[0]
-    k = max(1, int(0.05 * N))
-    topk_mass = float(torch.topk(Wt, k=k, dim=1).values.sum(dim=1).mean().item())
-
-    uniform = torch.full_like(Wt, 1.0 / N)
-    row_l2_dev = float(((Wt - uniform).pow(2).sum(dim=1).mean()).sqrt().item())
-
-    return {
-        "mean_row_entropy": row_entropy,
-        "mean_topk_mass_(k=5%N)": topk_mass,
-        "mean_row_L2_dev_from_uniform": row_l2_dev,
-        "min": float(Wt.min().item()),
-        "max": float(Wt.max().item()),
-        "mean": float(Wt.mean().item()),
-    }
-
-
-def frob_norm(W):
-    Wnp = _as_numpy(W)
-    return float(np.linalg.norm(Wnp.reshape(-1), ord=2))
-
-
-def dist_frob(Wa, Wb):
-    A = _as_numpy(Wa)
-    B = _as_numpy(Wb)
-    return float(np.linalg.norm((A - B).reshape(-1), ord=2))
-
-
-@torch.no_grad()
-def sample_graphon_on_fine_grid(model_or_callable, t=None, x=None, M=200, device="cpu"):
-    """
-    Sample either:
-      - a callable w(t,x,xi,zeta), or
-      - a callable w(xi,zeta),
-    on a fine M x M grid, for a smoother display.
-
-    This is useful if your model can evaluate directly at arbitrary (xi,zeta),
-    rather than only returning an N x N matrix.
-    """
-    grid = (torch.arange(M, device=device, dtype=torch.float32) + 0.5) / M
-    XI, ZETA = torch.meshgrid(grid, grid, indexing="ij")
-
-    try:
-        W = model_or_callable(t, x, XI, ZETA)
-    except TypeError:
-        W = model_or_callable(XI, ZETA)
-
-    return W.detach().cpu()
-
-
-@torch.no_grad()
-def plot_state_histogram_over_time_imshow(
-    xs,
-    dt: float,
-    bins=80,
-    value_range=None,
-    density=True,
-    cmap="viridis",
-    title="Empirical state histogram over time",
-    x_target=None,
-    show=True,
-):
-    X = torch.stack(xs, dim=0).detach().cpu().numpy()   # (Tsteps, B, N)
-    Tsteps = X.shape[0]
-
-    flat_all = X.reshape(-1)
-    if value_range is None:
-        xmin = float(flat_all.min())
-        xmax = float(flat_all.max())
-        pad = 0.05 * max(1e-8, xmax - xmin)
-        xmin -= pad
-        xmax += pad
-        value_range = (xmin, xmax)
-
-    x_edges = np.linspace(value_range[0], value_range[1], bins + 1)
-    x_centers = 0.5 * (x_edges[:-1] + x_edges[1:])
-
-    H = np.zeros((Tsteps, bins), dtype=float)
-    for n in range(Tsteps):
-        vals = X[n].reshape(-1)
-        hist, _ = np.histogram(vals, bins=x_edges, density=density)
-        H[n, :] = hist
-
-    _, ax = plt.subplots(1, 1, figsize=(8, 5))
-    im = ax.imshow(
-        H,
-        origin="lower",
-        aspect="auto",
-        extent=(x_centers[0], x_centers[-1], 0.0, dt * (Tsteps - 1)),
-        cmap=cmap,
-    )
-    ax.set_xlabel("state x")
-    ax.set_ylabel("time t")
-    ax.set_title(title)
-    plt.colorbar(im, ax=ax, label="density" if density else "count")
-
-    # Overlay: mean state trajectory
-    t_axis = np.linspace(0.0, dt * (Tsteps - 1), Tsteps)
-    mean_vals = X.reshape(Tsteps, -1).mean(axis=1)
-    ax.plot(mean_vals, t_axis, color="white", lw=1.5, label=r"$\bar{x}(t)$")
-
-    # Overlay: target state
-    if x_target is not None:
-        ax.axvline(x_target, color="white", lw=1.0, ls="--",
-                   label=rf"$x^* = {x_target}$")
-
-    ax.legend(loc="upper right", fontsize=8, framealpha=0.4)
-
-    if show:
-        plt.show()
-
-    return H, x_edges, ax
-
-
-def plot_x0_samples(xi, x0_batch, n_show=12, title="Sample initial-condition profiles"):
-    """
-    Line plot of n_show samples from x0_batch (B, N) vs xi (N,).
-    Returns fig, ax.
-    """
-    fig, ax = plt.subplots(figsize=(10, 4))
-    for i in range(min(n_show, x0_batch.shape[0])):
-        ax.plot(xi.cpu().numpy(), x0_batch[i].detach().cpu().numpy(), lw=0.8)
-    ax.set_xlabel(r"$\xi$")
-    ax.set_ylabel(r"$x_0(\xi)$")
-    ax.set_title(title)
-    fig.tight_layout()
-    return fig, ax
-
-
-def plot_x0_density(xi, x0_batch, gridsize=40, title="Initial-state density across labels and batch"):
-    """
-    Hexbin density of (xi, x0) pooled over all batch samples.
-    x0_batch shape: (B, N). xi shape: (N,).
-    Returns fig, ax.
-    """
-    xi_np = xi.cpu().numpy()
-    x0_np = x0_batch.detach().cpu().numpy()
-    B, N = x0_np.shape
-    xi_rep = np.tile(xi_np, B)
-    x0_rep = x0_np.reshape(-1)
-
-    fig, ax = plt.subplots(figsize=(10, 4))
-    hb = ax.hexbin(xi_rep, x0_rep, gridsize=gridsize, cmap="Blues", mincnt=1, linewidths=0.2)
-    fig.colorbar(hb, ax=ax, label="count")
-    ax.set_xlabel(r"$\xi$")
-    ax.set_ylabel(r"$x_0$")
-    ax.set_title(title)
-    fig.tight_layout()
-    return fig, ax
-
-
-def plot_phi(phi_fn, z_range=(-20.0, 20.0), n_points=400, title=None):
-    """
-    Plot phi_fn(z) vs z over z_range. Returns fig, ax.
-    """
-    import torch as _torch
-    zz = np.linspace(z_range[0], z_range[1], n_points)
-    z_t = _torch.tensor(zz, dtype=_torch.float32)
-    with _torch.no_grad():
-        yy = phi_fn(z_t).numpy()
-    fig, ax = plt.subplots(figsize=(5, 4))
-    ax.plot(zz, yy, lw=1.5)
-    ax.axhline(0, color="k", lw=0.5, ls="--")
-    ax.axvline(0, color="k", lw=0.5, ls="--")
-    ax.set_xlabel(r"$z$")
-    ax.set_ylabel(r"$\phi(z)$")
-    ax.set_title(title or r"Interaction kernel $\phi(z)$")
-    fig.tight_layout()
-    return fig, ax
-
-
-def save_graphon_gif(frames, path, fps=10, cmap="viridis", percentile_clip=0.01,
-                     t_vals=None, figsize=(5, 4), dpi=80):
-    """
-    frames : list of (N, N) tensors, one per time point.
-    t_vals : optional list of floats (same length as frames) used for per-frame title.
-    Renders each frame as a proper matplotlib figure (axes, colorbar, labels)
-    and saves as an animated GIF via Pillow.
-    """
-    from PIL import Image
-    from io import BytesIO
-
-    # Global color scale computed once for visual consistency across frames
-    all_vals = np.concatenate([_as_numpy(f).ravel() for f in frames])
-    q = (percentile_clip * 100) if percentile_clip else 0.0
-    vmin = float(np.percentile(all_vals, q))
-    vmax = float(np.percentile(all_vals, 100.0 - q))
-
-    pil_imgs = []
-    for i, W in enumerate(frames):
-        t_str = f"$t = {t_vals[i]:.2f}$" if t_vals is not None else f"frame {i}"
-        title = r"Graphon $w(\xi,\zeta)$,  " + t_str
-
-        fig, ax = plt.subplots(figsize=figsize)
-        imshow_graphon(
-            W,
-            title=title,
-            cmap=cmap,
-            interpolation="bicubic",
-            percentile_clip=None,   # use pre-computed vmin/vmax
-            vmin=vmin,
-            vmax=vmax,
-            ax=ax,
-            show=False,
-        )
-        fig.tight_layout()
-
-        buf = BytesIO()
-        fig.savefig(buf, format="png", dpi=dpi)
-        plt.close(fig)
-        buf.seek(0)
-        pil_imgs.append(Image.open(buf).copy())
-
-    duration_ms = int(1000 / fps)
-    pil_imgs[0].save(
-        path,
-        save_all=True,
-        append_images=pil_imgs[1:],
-        loop=0,
-        duration=duration_ms,
-        optimize=False,
-    )
\ No newline at end of file
+"""Compatibility imports for graphon_control.plotting."""
+from graphon_control.plotting import *  # noqa: F401,F403
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..4104658
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,31 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "graphon-control"
+version = "0.1.0"
+description = "Differentiable time-and-label control of graphon collective dynamics"
+readme = "README.md"
+requires-python = ">=3.10"
+license = {text = "MIT"}
+authors = [{name = "Joseph K. Miller"}, {name = "Brandon Su"}]
+dependencies = ["torch>=2.2", "numpy>=1.24", "matplotlib>=3.7", "tqdm>=4.66", "pillow>=10"]
+
+[project.optional-dependencies]
+dev = ["pytest>=8"]
+
+[project.scripts]
+graphon-kuramoto = "graphon_control.experiments.kuramoto:main"
+graphon-opinions = "graphon_control.experiments.bounded_confidence:main"
+graphon-train = "graphon_control.experiments.generic:main"
+
+[tool.setuptools.packages.find]
+include = ["graphon_control*"]
+
+[tool.setuptools]
+py-modules = ["model_1", "model_1_viz", "e1_kuramoto", "e2_bc", "e1_msweep", "e1_noise_figure"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-ra"
diff --git a/requirements-tested.txt b/requirements-tested.txt
new file mode 100644
index 0000000..f987b11
--- /dev/null
+++ b/requirements-tested.txt
@@ -0,0 +1,7 @@
+# Direct dependency versions used for the verified CPU test run (Python 3.12).
+torch==2.14.1
+numpy==2.4.3
+matplotlib==3.10.8
+tqdm==4.70.1
+pillow==12.1.1
+pytest==9.1.1
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..fe4647c
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,7 @@
+"""Keep numerical tests small and deterministic on shared CI hosts."""
+import os
+os.environ.setdefault('MPLBACKEND', 'Agg')
+os.environ.setdefault('OMP_NUM_THREADS', '1')
+os.environ.setdefault('MKL_NUM_THREADS', '1')
+import torch
+torch.set_num_threads(1)
diff --git a/tests/test_artifacts.py b/tests/test_artifacts.py
new file mode 100644
index 0000000..a504b7b
--- /dev/null
+++ b/tests/test_artifacts.py
@@ -0,0 +1,228 @@
+"""Artifact safety and postprocessing checks; no optimization is run here."""
+
+import json
+from pathlib import Path
+import subprocess
+
+import numpy as np
+import pytest
+import torch
+
+from graphon_control.artifacts import (
+    SCHEMA_VERSION, atomic_write_json, collect_paper_figures,
+    create_run_directory, load_controller, save_checkpoint, write_run_provenance,
+)
+import e1_msweep
+from e1_noise_figure import load_comparison, main as noise_main, validate_run_configs
+
+
+def test_run_directories_are_unique_and_refuse_collisions(tmp_path):
+    first = create_run_directory(tmp_path, "e1")
+    second = create_run_directory(tmp_path, "e1")
+    assert first != second
+    assert len(first.name.removeprefix("e1_")) == 32
+    create_run_directory(tmp_path, "e1", "paper")
+    with pytest.raises(FileExistsError, match="Choose a new tag"):
+        create_run_directory(tmp_path, "e1", "paper")
+    for invalid in ("", "../elsewhere", "/absolute", "two/parts", "a\\b", "a b"):
+        with pytest.raises(ValueError):
+            create_run_directory(tmp_path, "e1", invalid)
+
+
+def test_atomic_json_does_not_replace_or_publish_invalid_data(tmp_path):
+    target = atomic_write_json(tmp_path / "record.json", {"value": 1})
+    with pytest.raises(FileExistsError):
+        atomic_write_json(target, {"value": 2})
+    assert json.loads(target.read_text()) == {"value": 1}
+    with pytest.raises(ValueError):
+        atomic_write_json(tmp_path / "invalid.json", {"value": float("nan")})
+    assert not (tmp_path / "invalid.json").exists()
+    assert not list(tmp_path.glob(".record.json.*"))
+
+
+def test_provenance_records_configuration_seeds_and_environment(tmp_path):
+    directory = create_run_directory(tmp_path, "e1", "run")
+    config = {"N": 4, "device": "cpu", "seed": 7}
+    target = write_run_provenance(directory, config, 7, 1007)
+    saved = json.loads(target.read_text())
+    assert json.loads((directory / "config.json").read_text()) == config
+    assert saved["schema_version"] == SCHEMA_VERSION
+    assert (saved["training_seed"], saved["evaluation_seed"]) == (7, 1007)
+    assert saved["libraries"]["torch"]
+    assert saved["device"]["requested"] == "cpu"
+    assert saved["code"]["commit"]
+    assert len(saved["config_sha256"]) == 64
+    with pytest.raises(FileExistsError):
+        write_run_provenance(directory, config, 7, 1007)
+
+
+def test_checkpoint_roundtrip_and_raw_state_dict_compatibility(tmp_path):
+    source, restored = torch.nn.Linear(2, 3), torch.nn.Linear(2, 3)
+    path = save_checkpoint(tmp_path / "model.pt", source, metadata={"N": 2})
+    state = torch.load(path, weights_only=True)
+    assert set(state) == {"weight", "bias"}
+    assert load_controller(path, restored)["N"] == 2
+    for left, right in zip(source.parameters(), restored.parameters()):
+        torch.testing.assert_close(left, right)
+    with pytest.raises(FileExistsError):
+        save_checkpoint(path, source)
+    legacy = tmp_path / "legacy.pt"
+    torch.save(source.state_dict(), legacy)
+    assert load_controller(legacy, restored) == {}
+
+
+def _manifest(tmp_path, figures):
+    path = tmp_path / "selected.json"
+    path.write_text(json.dumps({"figures": figures, "output_dir": "exports"}))
+    return path
+
+
+def test_collector_preflights_all_sources_without_partial_writes(tmp_path):
+    source = tmp_path / "one.pdf"
+    source.write_bytes(b"first figure")
+    manifest = _manifest(tmp_path, {"first.pdf": "one.pdf", "second.pdf": "missing.pdf"})
+    with pytest.raises(ValueError, match="Missing figure"):
+        collect_paper_figures(manifest)
+    assert not (tmp_path / "exports").exists()
+    # Even an existing destination is untouched when another source is missing.
+    (tmp_path / "exports").mkdir()
+    retained = tmp_path / "exports" / "first.pdf"
+    retained.write_bytes(b"existing")
+    with pytest.raises(ValueError):
+        collect_paper_figures(manifest, overwrite=True)
+    assert retained.read_bytes() == b"existing"
+
+
+def test_collector_uses_manifest_paths_and_requires_explicit_overwrite(tmp_path, monkeypatch):
+    (tmp_path / "one.pdf").write_bytes(b"selected run")
+    manifest = _manifest(tmp_path, {"first.pdf": "one.pdf"})
+    monkeypatch.chdir(tmp_path.parent)
+    [destination] = collect_paper_figures(manifest)
+    assert destination.read_bytes() == b"selected run"
+    with pytest.raises(ValueError, match="Export exists"):
+        collect_paper_figures(manifest)
+    collect_paper_figures(manifest, overwrite=True)
+    manifest.write_text(json.dumps({"figures": {"../bad.pdf": "one.pdf"}}))
+    with pytest.raises(ValueError, match="safe path component"):
+        collect_paper_figures(manifest)
+
+
+def test_shell_collector_can_run_outside_repository(tmp_path):
+    (tmp_path / "one.pdf").write_bytes(b"figure")
+    manifest = _manifest(tmp_path, {"first.pdf": "one.pdf"})
+    script = Path(__file__).resolve().parents[1] / "collect_paper_figures.sh"
+    process = subprocess.run([str(script), str(manifest)], cwd=tmp_path,
+                             capture_output=True, text=True)
+    assert process.returncode == 0, process.stderr
+    assert (tmp_path / "exports" / "first.pdf").read_bytes() == b"figure"
+
+
+def _comparison_configs():
+    clean = {"N": 4, "k": 1, "T": 1.0, "dt": 0.5, "sigma": 0.0}
+    return {**clean, "sigma": 0.3}, clean
+
+
+@pytest.mark.parametrize("field,value", [("N", 8), ("k", 2), ("T", 2), ("dt", 0.25),
+                                       ("sigma", 0.2), ("score_bound", 1.0)])
+def test_noise_comparison_rejects_incompatible_reference(field, value):
+    noisy, clean = _comparison_configs()
+    clean[field] = value
+    with pytest.raises(ValueError):
+        validate_run_configs(noisy, clean)
+
+
+def _write_comparison(directory, config, noisy):
+    directory.mkdir()
+    atomic_write_json(directory / "config.json", config)
+    data = {"t": np.array([0.0, 0.5, 1.0])}
+    for key in ("r", "winding"):
+        suffixes = ("noisy_mean", "noisy_q10", "noisy_q90") if noisy else ("trained",)
+        for suffix in suffixes:
+            data[f"{key}_{suffix}"] = np.zeros(3)
+    np.savez(directory / "trajectories.npz", **data)
+
+
+def test_noise_comparison_checks_saved_grid_and_output_collision(tmp_path):
+    noisy_config, clean_config = _comparison_configs()
+    noisy, clean = tmp_path / "noisy", tmp_path / "clean"
+    _write_comparison(noisy, noisy_config, True)
+    _write_comparison(clean, clean_config, False)
+    loaded_config, _, _ = load_comparison(noisy, clean)
+    assert loaded_config == noisy_config
+    output = tmp_path / "panel.pdf"
+    output.write_bytes(b"keep")
+    with pytest.raises(SystemExit) as exc:
+        noise_main(["--noisy", str(noisy), "--noiseless", str(clean), "--out", str(output)])
+    assert exc.value.code == 1
+    assert output.read_bytes() == b"keep"
+    with np.load(clean / "trajectories.npz") as archive:
+        data = dict(archive)
+    data["t"] = np.array([0.0, 0.4, 1.0])
+    np.savez(clean / "trajectories.npz", **data)
+    with pytest.raises(ValueError, match="time grids differ"):
+        load_comparison(noisy, clean)
+
+
+def _write_sweep_cell(root, M, seed, steps=1):
+    directory = e1_msweep.cell_dir(M, seed, root)
+    directory.mkdir(parents=True)
+    config = e1_msweep._expected_config(M, seed, steps)
+    write_run_provenance(directory, config, seed, seed + 1000)
+    (directory / "model.pt").write_bytes(b"checkpoint fixture")
+    atomic_write_json(directory / "metrics.json", {
+        "costs": {"trained": {"total": M}, "meanfield": {"total": 10}},
+        "r_final": {"trained": 1}, "var_final": {"trained": 0},
+        "vxi": {"trained_max": 2},
+    })
+    atomic_write_json(directory / "complete.json", {"status": "complete"})
+    return directory
+
+
+def test_sweep_refuses_missing_or_partial_aggregation(tmp_path, monkeypatch):
+    monkeypatch.setattr(e1_msweep, "MS", [2, 4])
+    monkeypatch.setattr(e1_msweep, "SEEDS", [0])
+    with pytest.raises(ValueError, match="No complete sweep"):
+        e1_msweep.aggregate(steps=1, out_root=tmp_path)
+    assert list(tmp_path.iterdir()) == []
+    _write_sweep_cell(tmp_path, 2, 0)
+    with pytest.raises(ValueError, match="incomplete"):
+        e1_msweep.aggregate(steps=1, out_root=tmp_path)
+    assert len(list(tmp_path.iterdir())) == 1
+    summary_dir = e1_msweep.aggregate(steps=1, out_root=tmp_path, allow_partial=True)
+    summary = json.loads((summary_dir / "msweep_summary.json").read_text())
+    assert summary["complete"] is False
+    assert summary["missing"] == [{"M": 4, "seed": 0}]
+    assert (summary_dir / "convergence_M.pdf").is_file()
+
+
+def test_sweep_reuse_checks_provenance_and_config(tmp_path):
+    directory = _write_sweep_cell(tmp_path, 2, 0)
+    assert e1_msweep.done(2, 0, steps=1, out_root=tmp_path)
+    config_path = directory / "config.json"
+    config = json.loads(config_path.read_text())
+    config["beta"] = 999
+    atomic_write_json(config_path, config, overwrite=True)
+    with pytest.raises(ValueError, match="config beta"):
+        e1_msweep.done(2, 0, steps=1, out_root=tmp_path)
+
+
+def test_sweep_does_not_reuse_runs_that_failed_after_saving_metrics(tmp_path):
+    directory = _write_sweep_cell(tmp_path, 2, 0)
+    (directory / "complete.json").unlink()
+    with pytest.raises(ValueError, match="Incomplete or incompatible"):
+        e1_msweep.done(2, 0, steps=1, out_root=tmp_path)
+
+
+def test_sweep_creates_log_parent_and_reports_child_failure(tmp_path, monkeypatch):
+    monkeypatch.setattr(e1_msweep, "MS", [2])
+    monkeypatch.setattr(e1_msweep, "SEEDS", [0])
+    class FailedProcess:
+        pid = 123
+        returncode = 3
+        def poll(self):
+            return self.returncode
+    monkeypatch.setattr(e1_msweep.subprocess, "Popen", lambda *args, **kwargs: FailedProcess())
+    out_root = tmp_path / "new" / "out"
+    with pytest.raises(RuntimeError, match=r"failed \(3\)"):
+        e1_msweep.launch_all(1, out_root=out_root, lanes=1)
+    assert (out_root / "e1_msweep_M2_s0.log").is_file()
diff --git a/tests/test_controls.py b/tests/test_controls.py
new file mode 100644
index 0000000..eaee019
--- /dev/null
+++ b/tests/test_controls.py
@@ -0,0 +1,194 @@
+import io
+import math
+
+import pytest
+import torch
+
+from graphon_control.controls import FourierFeatures, MLP, TimeLabelGraphonControl
+
+
+def control(xi=None, **kwargs):
+    if xi is None:
+        xi = torch.linspace(0, 1, 5)
+    return TimeLabelGraphonControl(
+        xi, num_fourier_freqs=2, embed_dim=4, hidden_embed=7, hidden_score=9, **kwargs
+    )
+
+
+def test_directed_row_stochastic_with_self_interactions():
+    model = control()
+    W = model(torch.tensor([0.0, 0.5, 1.0]))
+    assert W.shape == (3, 5, 5)
+    assert torch.all(W > 0)
+    assert torch.all(W.diagonal(dim1=-2, dim2=-1) > 0)
+    torch.testing.assert_close(W.sum(-1), torch.ones(3, 5))
+
+
+def test_time_shapes_distinguish_scalar_and_one_element_batch():
+    model = control()
+    scalar = model(torch.tensor(0.2))
+    singleton = model(torch.tensor([0.2]))
+    batch = model(torch.tensor([0.2, 0.4]))
+    assert scalar.shape == (5, 5)
+    assert singleton.shape == (1, 5, 5)
+    assert batch.shape == (2, 5, 5)
+    torch.testing.assert_close(scalar, singleton[0])
+    torch.testing.assert_close(scalar, batch[0])
+    torch.testing.assert_close(scalar, model(0.2))
+
+
+@pytest.mark.parametrize("positive_map", ["softplus", "exp"])
+def test_symmetric_alternative_has_mean_row_normalization(positive_map):
+    # Fix a nonconstant score landscape so row-stochasticity cannot pass by chance.
+    model = control(enforce_symmetry=True, positive_map=positive_map)
+    with torch.no_grad():
+        for parameter in model.parameters():
+            parameter.fill_(0.1)
+    W = model(torch.tensor([0.0, 1.0]))
+    torch.testing.assert_close(W, W.transpose(-1, -2))
+    torch.testing.assert_close(W.sum(-1).mean(-1), torch.ones(2))
+    assert not torch.allclose(W.sum(-1), torch.ones(2, 5))
+
+
+@pytest.mark.parametrize("positive_map", ["softplus", "exp"])
+def test_symmetric_normalization_handles_underflow(positive_map):
+    model = control(enforce_symmetry=True, positive_map=positive_map)
+    with torch.no_grad():
+        model.score_net.net[-1].weight.zero_()
+        model.score_net.net[-1].bias.fill_(-1000)
+    W = model(0.0)
+    assert torch.isfinite(W).all()
+    torch.testing.assert_close(W.sum(-1).mean(), torch.tensor(1.0))
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_parameters_buffers_and_outputs_inherit_label_dtype(dtype):
+    model = control(torch.linspace(0, 1, 4, dtype=dtype))
+    for value in list(model.parameters()) + list(model.buffers()):
+        assert value.dtype == dtype
+        assert value.device.type == "cpu"
+    assert model(0.1).dtype == dtype
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps"])
+def test_accelerator_construction_and_time_conversion(device):
+    available = torch.cuda.is_available() if device == "cuda" else torch.backends.mps.is_available()
+    if not available:
+        pytest.skip(f"{device} unavailable")
+    model = control(torch.linspace(0, 1, 4, device=device))
+    for value in list(model.parameters()) + list(model.buffers()):
+        assert value.device.type == device
+    # CPU time inputs may be used with a controller constructed on an accelerator.
+    W = model(torch.tensor([0.0, 0.2]))
+    assert W.device.type == device
+    torch.testing.assert_close(W.sum(-1).cpu(), torch.ones(2, 4))
+
+
+@pytest.mark.parametrize("bound", [0.0, 0.2, 1.0])
+def test_bounded_scores_imply_graphon_density_cap(bound):
+    model = control(torch.linspace(0, 1, 9), score_bound=bound)
+    with torch.no_grad():
+        for parameter in model.parameters():
+            parameter.mul_(100)
+    W = model(torch.tensor([0.0, 0.5]))
+    assert torch.all(9 * W <= math.exp(2 * bound) + 1e-6)
+    torch.testing.assert_close(W.sum(-1), torch.ones(2, 9))
+    if bound == 0:
+        torch.testing.assert_close(W, torch.full_like(W, 1 / 9))
+
+
+def test_checkpoint_state_names_and_roundtrip():
+    model = control(score_bound=0.7)
+    expected = {"xi", "psi_xi_cached", "ff.freqs"}
+    for module in ("receiver_net", "sender_net", "score_net"):
+        for index in (0, 2, 4):
+            for kind in ("weight", "bias"):
+                expected.add(f"{module}.net.{index}.{kind}")
+    assert set(model.state_dict()) == expected
+    expected_parameter_order = [
+        f"{module}.net.{index}.{kind}"
+        for module in ("receiver_net", "sender_net", "score_net")
+        for index in (0, 2, 4)
+        for kind in ("weight", "bias")
+    ]
+    assert [name for name, _ in model.named_parameters()] == expected_parameter_order
+    checkpoint = io.BytesIO()
+    torch.save(model.state_dict(), checkpoint)
+    checkpoint.seek(0)
+    restored = control(score_bound=0.7)
+    restored.load_state_dict(torch.load(checkpoint, weights_only=True))
+    torch.testing.assert_close(model(torch.tensor([0.2, 0.8])), restored(torch.tensor([0.2, 0.8])))
+
+
+@pytest.mark.parametrize("symmetric", [False, True])
+def test_finite_nonzero_gradients_through_control(symmetric):
+    torch.manual_seed(3)
+    model = control(torch.linspace(0, 1, 5, dtype=torch.float64), enforce_symmetry=symmetric)
+    times = torch.tensor([0.1, 0.7], dtype=torch.float64, requires_grad=True)
+    W = model(times)
+    coefficients = torch.arange(W.numel(), dtype=W.dtype).reshape_as(W)
+    (W * coefficients.square()).sum().backward()
+    assert times.grad is not None and torch.isfinite(times.grad).all()
+    gradients = [p.grad for p in model.parameters()]
+    assert all(g is not None and torch.isfinite(g).all() for g in gradients)
+    assert any(torch.any(g != 0) for g in gradients)
+
+
+def test_label_grid_is_fixed_against_external_mutation():
+    xi = torch.linspace(0, 1, 5)
+    model = control(xi)
+    before = model(0.0)
+    xi.zero_()
+    torch.testing.assert_close(model(0.0), before)
+    torch.testing.assert_close(model.xi, torch.linspace(0, 1, 5))
+
+
+@pytest.mark.parametrize("xi", [
+    torch.tensor([], dtype=torch.float32), torch.tensor([[0.0, 1.0]]),
+    torch.tensor([0, 1]), torch.tensor([0.0, float("nan")]),
+    torch.tensor([0.0, float("inf")]), torch.tensor([-0.1, 0.5]),
+    torch.tensor([0.0, 1.1]), torch.tensor([0.7, 0.2]), [0.0, 1.0],
+])
+def test_invalid_label_grids(xi):
+    with pytest.raises(ValueError, match="xi"):
+        control(xi)
+
+
+@pytest.mark.parametrize("name,value", [
+    ("num_fourier_freqs", 0), ("num_fourier_freqs", True),
+    ("embed_dim", -1), ("hidden_embed", 1.5), ("hidden_score", 0),
+    ("time_scale", 0), ("time_scale", -1), ("time_scale", float("inf")),
+    ("time_scale", float("nan")), ("positive_map", "relu"),
+    ("enforce_symmetry", "yes"), ("score_bound", -0.1),
+    ("score_bound", float("inf")), ("score_bound", float("nan")),
+])
+def test_invalid_configuration(name, value):
+    kwargs = dict(num_fourier_freqs=2, embed_dim=4, hidden_embed=7, hidden_score=9)
+    kwargs[name] = value
+    with pytest.raises(ValueError, match=name):
+        TimeLabelGraphonControl(torch.linspace(0, 1, 5), **kwargs)
+
+
+@pytest.mark.parametrize("t", [torch.tensor([]), torch.zeros(2, 1),
+                                    torch.tensor(float("nan")), torch.tensor([float("inf")]),
+                                    torch.tensor(1 + 2j)])
+def test_invalid_times(t):
+    with pytest.raises(ValueError, match="t must"):
+        control()(t)
+
+
+def test_scaled_times_must_be_representable():
+    with pytest.raises(ValueError, match="scaled times"):
+        control(time_scale=1e-300)(1.0)
+
+
+def test_fourier_feature_convention():
+    features = FourierFeatures(2)(torch.tensor([0.0, 0.25], dtype=torch.float64))
+    expected = torch.tensor([[0.0, 0.0, 1.0, 1.0], [1.0, 0.0, 0.0, -1.0]], dtype=torch.float64)
+    torch.testing.assert_close(features, expected, atol=1e-14, rtol=1e-14)
+
+
+@pytest.mark.parametrize("args", [(0, 2, 1), (2, -1, 1), (2, 2, 0), (2, 2, 1, -1)])
+def test_mlp_rejects_invalid_sizes(args):
+    with pytest.raises(ValueError):
+        MLP(*args)
diff --git a/tests/test_experiments.py b/tests/test_experiments.py
new file mode 100644
index 0000000..ec427c3
--- /dev/null
+++ b/tests/test_experiments.py
@@ -0,0 +1,148 @@
+"""Experiment entry points preserve the common scientific and CLI contracts."""
+
+import dataclasses
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+
+import numpy as np
+import pytest
+import torch
+
+from graphon_control.dynamics import mean_field_kernel
+from graphon_control.experiments import bounded_confidence, generic, kuramoto
+
+
+torch.set_num_threads(1)
+
+
+@pytest.mark.parametrize("config_type", [kuramoto.E1Config, bounded_confidence.E2Config, generic.TrainConfig])
+def test_default_configuration_json_roundtrip(config_type):
+    config = config_type()
+    restored = config_type(**json.loads(json.dumps(dataclasses.asdict(config))))
+    assert dataclasses.asdict(restored) == dataclasses.asdict(config)
+
+
+@pytest.mark.parametrize("invalid", [{"T": 1.0, "dt": 0.3}, {"sigma": -1}, {"N": 0},
+    {"evaluation_seed": -2}, {"score_bound": float("nan")}])
+@pytest.mark.parametrize("config_type", [kuramoto.E1Config, bounded_confidence.E2Config, generic.TrainConfig])
+def test_invalid_experiment_configuration(config_type, invalid):
+    with pytest.raises(ValueError):
+        config_type(**invalid)
+
+
+@pytest.mark.parametrize("config_type", [kuramoto.E1Config, bounded_confidence.E2Config])
+def test_experiment_jitter_must_be_nonnegative(config_type):
+    with pytest.raises(ValueError, match="jitter"):
+        config_type(jitter=-1)
+
+
+@pytest.mark.parametrize("level", [4.0, True, -2, 3])
+def test_kuramoto_coarse_level_requires_integer_partition(level):
+    with pytest.raises(ValueError, match="x0_level"):
+        kuramoto.E1Config(N=8, x0_level=level)
+
+
+def test_bounded_confidence_modes_rejects_bool():
+    with pytest.raises(ValueError, match="modes"):
+        bounded_confidence.E2Config(modes=True)
+
+
+def test_e1_and_e2_use_identical_objective_for_identical_trajectories():
+    xs = torch.tensor([[[0.0, 2.0, 4.0]], [[0.5, 2.0, 3.5]], [[1.0, 2.0, 3.0]]])
+    Ws = mean_field_kernel(3).expand(2, 3, 3)
+    parameters = dict(N=8, T=0.2, dt=0.1, beta=0.7, gamma_target=1.4,
+                      lambda_terminal=0.3, x_target=-1.0)
+    # The objective depends on provided trajectories rather than cfg.N.
+    one = kuramoto.e1_cost(xs, Ws, kuramoto.E1Config(**parameters))
+    two = bounded_confidence.e2_cost(xs, Ws, bounded_confidence.E2Config(**parameters))
+    assert one.keys() == two.keys()
+    for key in one:
+        assert torch.equal(one[key], two[key])
+
+
+@pytest.mark.parametrize("module,config_type,runner", [
+    (kuramoto, kuramoto.E1Config, "run_e1"),
+    (bounded_confidence, bounded_confidence.E2Config, "run_e2"),
+    (generic, generic.TrainConfig, "run_experiment"),
+])
+def test_cli_loads_json_flags_override_and_smoke_caps_model(tmp_path, monkeypatch,
+                                                          module, config_type, runner):
+    config = config_type(N=32, T=0.2, dt=0.1, beta=2.0, batch_size=3,
+                         num_steps=123, embed_dim=24)
+    path = tmp_path / "config.json"
+    path.write_text(json.dumps(dataclasses.asdict(config)))
+    seen = {}
+    def capture(cfg, **kwargs):
+        seen.update(cfg=cfg, kwargs=kwargs)
+    monkeypatch.setattr(module, runner, capture)
+    module.main(["--config", str(path), "--beta", "3.5", "--smoke", "--no-plots", "--tag", "check"])
+    restored = seen["cfg"]
+    assert restored.beta == 3.5
+    assert restored.T == 0.2 and restored.dt == 0.1 and restored.batch_size == 3
+    assert restored.num_steps == 20
+    assert restored.N == 16 and restored.embed_dim == 8
+    assert restored.num_fourier_freqs == 2
+    assert not seen["kwargs"]["make_plots"]
+
+
+@pytest.mark.parametrize("module", [kuramoto, bounded_confidence, generic])
+def test_cli_reports_invalid_time_grid_before_running(module, capsys):
+    with pytest.raises(SystemExit) as error:
+        module.main(["--T", "1", "--dt", "0.3", "--no-plots"])
+    assert error.value.code == 2
+    assert "T/dt must be" in capsys.readouterr().err
+
+
+@pytest.mark.parametrize("script", ["e1_kuramoto.py", "e2_bc.py", "model_1.py"])
+def test_compatibility_cli_help_works_without_display(tmp_path, script):
+    environment = os.environ.copy()
+    environment.pop("DISPLAY", None)
+    environment.update(MPLBACKEND="Agg", MPLCONFIGDIR=str(tmp_path / "mpl"),
+                       XDG_CACHE_HOME=str(tmp_path / "cache"))
+    root = Path(__file__).resolve().parents[1]
+    result = subprocess.run([sys.executable, str(root / script), "--help"],
+                            cwd=tmp_path, env=environment, capture_output=True,
+                            text=True, timeout=45)
+    assert result.returncode == 0, result.stderr
+    assert "--config" in result.stdout and "--no-plots" in result.stdout
+
+
+@pytest.mark.parametrize("module,config_type,runner,prefix", [
+    (kuramoto, kuramoto.E1Config, "run_e1", "e1"),
+    (bounded_confidence, bounded_confidence.E2Config, "run_e2", "e2"),
+])
+def test_short_paper_experiment_archives_reproducible_numerical_run(tmp_path,
+                                                               module, config_type, runner, prefix):
+    cfg = config_type(N=8, T=0.1, dt=0.05, batch_size=2, num_steps=2,
+                      num_fourier_freqs=2, embed_dim=4, hidden_embed=8,
+                      hidden_score=8, seed=0)
+    getattr(module, runner)(cfg, tag="numerical", out_root=str(tmp_path), make_plots=False)
+    run = tmp_path / f"{prefix}_numerical"
+    assert config_type(**json.loads((run / "config.json").read_text())) == cfg
+    history = json.loads((run / "history.json").read_text())
+    assert len(history["loss"]) == 2
+    assert all(np.isfinite(history["loss"]))
+    assert (run / "model.pt").is_file()
+    assert (run / "metrics.json").is_file()
+    assert json.loads((run / "complete.json").read_text())["status"] == "complete"
+    with np.load(run / "trajectories.npz", allow_pickle=False) as archive:
+        assert archive["x_trained"].shape == (3, 8)
+        assert archive["W_trained"].shape == (2, 8, 8)
+        assert np.isfinite(archive["x_trained"]).all()
+    assert not list(run.glob("*.pdf"))
+    assert not list(run.glob("*.gif"))
+
+
+def test_generic_default_time_scale_uses_horizon_and_runs(tmp_path):
+    cfg = generic.TrainConfig(N=4, T=0.1, dt=0.05, sigma=0, batch_size=2,
+                              num_steps=2, num_fourier_freqs=2, embed_dim=4,
+                              hidden_embed=8, hidden_score=8)
+    result = generic.run_experiment(cfg, out_root=tmp_path, run_tag="generic",
+                                    eval_batch_size=2, make_plots=False)
+    assert result["model"].time_scale == cfg.T
+    assert len(result["history"]["loss"]) == 2
+    assert len(result["xs"]) == 3
+    assert Path(result["run_dir"], "complete.json").is_file()
diff --git a/tests/test_numerics.py b/tests/test_numerics.py
new file mode 100644
index 0000000..edf2047
--- /dev/null
+++ b/tests/test_numerics.py
@@ -0,0 +1,211 @@
+"""Analytic and scaling checks for the manuscript's numerical conventions."""
+
+import math
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from graphon_control.diagnostics import kernel_constraint_report, vxi_discrete
+from graphon_control.dynamics import (
+    ConstantKernel, mean_field_kernel, precompute_Ws, rollout_dynamics_w1, rollout_with_Ws,
+)
+from graphon_control.objectives import (
+    empirical_variance, graphon_l2_penalty, graphon_uniform_penalty, paper_cost,
+    resolve_target_schedule, sampled_cost_JNdt, trajectory_cost,
+)
+from graphon_control.validation import n_time_steps, validate_common_config
+
+
+torch.set_num_threads(1)
+
+
+@pytest.mark.parametrize("T,dt,expected", [(1.0, 0.02, 50), (0.3, 0.1, 3), (2, 2, 1)])
+def test_integral_time_grids(T, dt, expected):
+    assert n_time_steps(T, dt) == expected
+
+
+@pytest.mark.parametrize("T,dt", [(0, 0.1), (1, 0), (-1, 0.1), (1, -0.1),
+                                  (1, 2), (1, 0.3), (math.nan, 0.1),
+                                  (1, math.inf), (math.inf, 1)])
+def test_invalid_time_grid_is_rejected(T, dt):
+    with pytest.raises(ValueError):
+        n_time_steps(T, dt)
+
+
+@pytest.mark.parametrize("field,value", [("N", 0), ("batch_size", 1.5), ("sigma", -1),
+                                         ("beta", math.nan), ("lr", 0), ("grad_clip", 0),
+                                         ("evaluation_seed", -1), ("evaluation_seed", True),
+                                         ("score_bound", -1), ("score_bound", math.inf),
+                                         ("jitter", -0.1), ("interaction_alpha", math.nan)])
+def test_common_config_rejects_invalid_fields(field, value):
+    cfg = SimpleNamespace(T=1.0, dt=0.1, N=4, batch_size=2, sigma=0.0, beta=1.0,
+                          lr=0.01, grad_clip=1.0)
+    setattr(cfg, field, value)
+    with pytest.raises(ValueError, match=field):
+        validate_common_config(cfg)
+
+
+def test_optional_time_scale_selects_horizon_and_finite_signed_interaction_is_valid():
+    cfg = SimpleNamespace(T=1.0, dt=0.1, time_scale=None, interaction_alpha=-2.0,
+                          evaluation_seed=1729, score_bound=None)
+    validate_common_config(cfg)
+
+
+@pytest.mark.parametrize("N", [1, 2, 8, 32])
+def test_uniform_graphon_scaling_is_resolution_independent(N):
+    W = mean_field_kernel(N, dtype=torch.float64)
+    assert torch.allclose(W.sum(-1), torch.ones(N, dtype=W.dtype))
+    assert graphon_l2_penalty(W).item() == pytest.approx(1.0)
+    assert graphon_uniform_penalty(W).item() == 0.0
+    assert vxi_discrete(W).item() == 0.0
+
+
+def test_linear_consensus_matches_exact_euler_solution():
+    # For W=1/N and phi(z)=z, the mean is constant, while each deviation is
+    # multiplied by (1-dt). This simultaneously checks sign and N scaling.
+    x0 = torch.tensor([[1.0, 3.0, 5.0], [-2.0, 4.0, 10.0]], dtype=torch.float64)
+    dt, NT = 0.1, 5
+    Ws = mean_field_kernel(3, dtype=x0.dtype).expand(NT, 3, 3)
+    xs = rollout_with_Ws(Ws, x0, dt, lambda y: y)
+    mean = x0.mean(-1, keepdim=True)
+    for n in range(NT + 1):
+        expected = mean + (1 - dt) ** n * (x0 - mean)
+        assert torch.allclose(xs[n], expected, atol=1e-12, rtol=0)
+
+
+def test_directed_receiver_sender_orientation():
+    # Both receivers listen exclusively to sender 1. Sender 1 stays still.
+    W = torch.tensor([[[0.0, 1.0], [0.0, 1.0]]], dtype=torch.float64)
+    x0 = torch.tensor([[0.0, 2.0]], dtype=W.dtype)
+    xs = rollout_with_Ws(W, x0, 0.25, lambda z: z)
+    assert torch.equal(xs[1], torch.tensor([[0.5, 2.0]], dtype=W.dtype))
+
+
+def test_prescribed_brownian_increments_are_not_rescaled():
+    x0 = torch.zeros((1, 2), dtype=torch.float64)
+    Ws = mean_field_kernel(2, dtype=x0.dtype).expand(2, 2, 2)
+    increments = torch.tensor([[[0.2, -0.1]], [[-0.3, 0.4]]], dtype=x0.dtype)
+    xs = rollout_with_Ws(Ws, x0, 0.04, torch.zeros_like, sigma=2.0,
+                         brownian_increments=increments)
+    assert torch.allclose(xs[1:], 2.0 * increments.cumsum(0))
+
+
+def test_generated_brownian_noise_has_sqrt_dt_scaling_and_is_replayable():
+    x0 = torch.zeros((2, 3), dtype=torch.float64)
+    Ws = mean_field_kernel(3, dtype=x0.dtype).expand(2, 3, 3)
+    dt, sigma = 0.09, 1.7
+    generator = torch.Generator().manual_seed(123)
+    xs = rollout_with_Ws(Ws, x0, dt, torch.zeros_like, sigma, generator=generator)
+    replay = torch.Generator().manual_seed(123)
+    normals = torch.stack([torch.randn(x0.shape, dtype=x0.dtype, generator=replay)
+                           for _ in range(2)])
+    assert torch.allclose(xs[1:], sigma * math.sqrt(dt) * normals.cumsum(0))
+
+
+def test_deterministic_rollout_does_not_consume_generator_state():
+    generator = torch.Generator().manual_seed(9)
+    state_before = generator.get_state().clone()
+    x0 = torch.zeros((1, 2))
+    rollout_with_Ws(mean_field_kernel(2).expand(2, 2, 2), x0, 0.1, torch.sin,
+                    sigma=0, generator=generator)
+    assert torch.equal(generator.get_state(), state_before)
+
+
+def test_one_step_precomputation_and_legacy_rollout_agree():
+    x0 = torch.tensor([[0.0, 1.0]], dtype=torch.float64)
+    model = ConstantKernel(mean_field_kernel(2, dtype=x0.dtype))
+    Ws = precompute_Ws(model, 1, 0.1, "cpu", 3)
+    assert Ws.shape == (1, 2, 2)
+    assert Ws.dtype == torch.float64
+    xs, Ws_list = rollout_dynamics_w1(model, x0, 0.1, 0.1, 0, torch.sin)
+    assert len(xs) == 2 and len(Ws_list) == 1
+    assert torch.equal(torch.stack(xs), rollout_with_Ws(Ws, x0, 0.1, torch.sin))
+
+
+def test_paper_and_generic_beta_convention_with_analytic_cost():
+    # Identity interactions preserve x=[-1,1]; variance=1, ||w||²=2.
+    # Integral cost = T*(1+beta*2); terminal mean error is (0-target)².
+    model = ConstantKernel(torch.eye(2, dtype=torch.float64))
+    x0 = torch.tensor([[-1.0, 1.0]], dtype=torch.float64)
+    T, dt, beta = 0.2, 0.1, 3.0
+    Ws = precompute_Ws(model, 2, dt, "cpu", 1)
+    xs = rollout_with_Ws(Ws, x0, dt, lambda y: y)
+    cfg = SimpleNamespace(dt=dt, beta=beta, lambda_terminal=0.5, gamma_target=2.0, x_target=3.0)
+    parts = paper_cost(xs, Ws, cfg)
+    loss, aux = sampled_cost_JNdt(model, x0, T, dt, 0, beta, 0.5, lambda y: y,
+                                 gamma_target=2.0, x_target=3.0)
+    assert parts["l2_reg_cost"].item() == pytest.approx(T * beta * 2)
+    assert parts["total"].item() == pytest.approx(T * (1 + beta * 2) + 0.5 + 18)
+    assert loss.item() == pytest.approx(parts["total"].item())
+    assert aux["total_cost"] == pytest.approx(loss.item())
+
+
+def test_extra_objective_terms_have_explicit_scaling():
+    x = torch.tensor([[[0.0, 2.0]]], dtype=torch.float64)
+    xs = x.expand(3, 1, 2)
+    Ws = torch.eye(2, dtype=x.dtype).expand(2, 2, 2)
+    schedule = resolve_target_schedule(0.2, 0.1, 0, [0, 0.2], [2, 0], [3, 4])
+    parts = trajectory_cost(xs, Ws, dt=0.1, beta=7, eta_uniform=6,
+                            regularizer_mode="uniform", weighted_running_cost=True,
+                            target_schedule=schedule)
+    assert parts["state_cost"].item() == pytest.approx(0.15)
+    assert parts["l2_reg_cost"].item() == 0
+    assert parts["uniform_cost"].item() == pytest.approx(0.6)
+    assert parts["scheduled_target_cost"].item() == pytest.approx(7)
+
+
+@pytest.mark.parametrize("times,values,weights", [([0.1], [], []), ([0.1], [1], [1, 2]),
+    ([math.nan], [1], [1]), ([0.3], [1], [1]), ([0.1], [math.inf], [1]),
+    ([0.1], [1], [-1]), ([0.1], [1], [])])
+def test_invalid_checkpoint_schedules(times, values, weights):
+    with pytest.raises(ValueError):
+        resolve_target_schedule(0.2, 0.1, 0, times, values, weights)
+
+
+def test_checkpoint_default_weights_and_snapping():
+    schedule = resolve_target_schedule(0.3, 0.1, 2, [0.14, 0.3], [1, 2], [])
+    assert [target["index"] for target in schedule] == [1, 3]
+    assert all(target["weight"] == 2 for target in schedule)
+
+
+def test_receiver_variation_uses_graphon_values_and_no_periodic_edge():
+    W = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float64)
+    assert vxi_discrete(W).item() == pytest.approx(2.0)
+    report = kernel_constraint_report(torch.stack([W, W]))
+    assert torch.equal(report["max_row_sum_error"], torch.zeros(2, dtype=W.dtype))
+    assert report["receiver_variation"].shape == (2,)
+
+
+def test_variance_is_stable_under_large_state_translation():
+    x = torch.tensor([[1e12 - 1, 1e12 + 1]], dtype=torch.float64)
+    assert empirical_variance(x).item() == pytest.approx(1.0)
+
+
+def test_differentiable_rollout_and_cost_have_finite_gradients():
+    logits = torch.tensor([[0.2, -0.1], [-0.4, 0.3]], dtype=torch.float64, requires_grad=True)
+    Ws = logits.softmax(-1).expand(3, 2, 2)
+    x0 = torch.tensor([[-1.0, 2.0]], dtype=torch.float64)
+    xs = rollout_with_Ws(Ws, x0, 0.1, torch.tanh)
+    loss = trajectory_cost(xs, Ws, dt=0.1, beta=0.2)["total"]
+    loss.backward()
+    assert torch.isfinite(logits.grad).all()
+    assert logits.grad.abs().sum() > 0
+    # Gradcheck compares autograd to central finite differences of the full path.
+    def objective(theta):
+        kernels = theta.softmax(-1).expand(3, 2, 2)
+        states = rollout_with_Ws(kernels, x0, 0.1, torch.tanh)
+        return trajectory_cost(states, kernels, dt=0.1, beta=0.2)["total"]
+    assert torch.autograd.gradcheck(objective, (logits,), eps=1e-6, atol=1e-5)
+
+
+@pytest.mark.parametrize("NT,chunk", [(0, 1), (1, 0), (1, 1.5)])
+def test_precompute_rejects_empty_or_invalid_batches(NT, chunk):
+    with pytest.raises(ValueError):
+        precompute_Ws(ConstantKernel(mean_field_kernel(2)), NT, 0.1, "cpu", chunk)
+
+
+def test_rollout_rejects_mismatched_noise_shape():
+    with pytest.raises(ValueError, match="brownian_increments"):
+        rollout_with_Ws(mean_field_kernel(2).unsqueeze(0), torch.zeros((1, 2)),
+                        0.1, torch.sin, 0.2, brownian_increments=torch.zeros(2))
diff --git a/tests/test_package.py b/tests/test_package.py
new file mode 100644
index 0000000..e873861
--- /dev/null
+++ b/tests/test_package.py
@@ -0,0 +1,95 @@
+"""Public package imports, compatibility entry points, and legacy checkpoints."""
+
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+
+import torch
+
+from graphon_control.artifacts import load_controller
+from graphon_control.controls import TimeLabelGraphonControl
+
+
+def test_numerical_imports_need_no_plotting_and_create_no_files(tmp_path):
+    # Check in a fresh interpreter: other tests may intentionally import plotting.
+    source_root = Path(__file__).resolve().parents[1]
+    environment = dict(os.environ, PYTHONPATH=str(source_root))
+    program = """
+import importlib.abc
+import pathlib
+import sys
+class NoPlotting(importlib.abc.MetaPathFinder):
+    def find_spec(self, fullname, path=None, target=None):
+        if fullname.split('.')[0] in {'matplotlib', 'PIL'}:
+            raise AssertionError('numerical import attempted plotting: ' + fullname)
+sys.meta_path.insert(0, NoPlotting())
+import graphon_control
+import graphon_control.controls
+import graphon_control.dynamics
+import graphon_control.objectives
+import graphon_control.diagnostics
+import graphon_control.training
+import graphon_control.experiments.generic
+import graphon_control.experiments.kuramoto
+import graphon_control.experiments.bounded_confidence
+import model_1
+import e1_kuramoto
+import e2_bc
+assert not any(m.startswith('matplotlib') for m in sys.modules)
+assert model_1.TimeLabelGraphonControl is graphon_control.TimeLabelGraphonControl
+assert e1_kuramoto.run_e1 is graphon_control.experiments.kuramoto.run_e1
+assert e2_bc.run_e2 is graphon_control.experiments.bounded_confidence.run_e2
+assert list(pathlib.Path.cwd().iterdir()) == []
+print('numerical imports have no plotting or output side effects')
+"""
+    result = subprocess.run([sys.executable, "-c", program], cwd=tmp_path,
+                            env=environment, text=True, capture_output=True)
+    assert result.returncode == 0, result.stderr
+    assert "no plotting or output side effects" in result.stdout
+    assert list(tmp_path.iterdir()) == []
+
+
+def test_legacy_raw_controller_checkpoint_loads_through_public_api(tmp_path):
+    # A handcrafted original-model state dict fixes the old names/shapes without
+    # instantiating the new class to generate the compatibility fixture itself.
+    xi = torch.tensor([0.25, 0.75])
+    legacy = {
+        "xi": xi,
+        "psi_xi_cached": torch.cat((torch.sin(2 * torch.pi * xi[:, None]),
+                                    torch.cos(2 * torch.pi * xi[:, None])), dim=-1),
+        "ff.freqs": torch.tensor([1.0]),
+    }
+    for module, sizes in (("receiver_net", (2, 3, 3, 2)),
+                          ("sender_net", (2, 3, 3, 2)),
+                          ("score_net", (5, 4, 4, 1))):
+        for index, (inputs, outputs) in zip((0, 2, 4), zip(sizes[:-1], sizes[1:])):
+            legacy[f"{module}.net.{index}.weight"] = torch.zeros(outputs, inputs)
+            legacy[f"{module}.net.{index}.bias"] = torch.zeros(outputs)
+    checkpoint = tmp_path / "legacy-model.pt"
+    torch.save(legacy, checkpoint)
+    restored = TimeLabelGraphonControl(xi, 1, 2, 3, 4)
+    assert load_controller(checkpoint, restored) == {}
+    torch.testing.assert_close(restored(torch.tensor([0.0, 0.5])), torch.full((2, 2, 2), 0.5))
+
+
+def test_generic_run_archives_reloadable_numerics_without_plotting(tmp_path):
+    from graphon_control.experiments.generic import TrainConfig, run_experiment
+
+    config = TrainConfig(N=4, T=0.1, dt=0.05, sigma=0, num_steps=1, batch_size=2,
+                         num_fourier_freqs=1, embed_dim=3, hidden_embed=4, hidden_score=4,
+                         seed=2, evaluation_seed=7, device="cpu")
+    result = run_experiment(config, out_root=tmp_path, run_tag="tiny", eval_batch_size=3,
+                            make_plots=False)
+    directory = Path(result["run_dir"])
+    for name in ("config.json", "provenance.json", "history.json", "metrics.json",
+                 "model.pt", "trajectories.npz", "complete.json"):
+        assert (directory / name).is_file()
+    metrics = json.loads((directory / "metrics.json").read_text())
+    assert metrics["l2_reg_cost"] >= config.beta * config.T - 1e-8
+    assert json.loads((directory / "complete.json").read_text())["plots_generated"] is False
+    restored = TimeLabelGraphonControl(result["xi"], 1, 3, 4, 4, time_scale=config.T)
+    load_controller(directory / "model.pt", restored)
+    times = torch.tensor([0.0, config.dt])
+    torch.testing.assert_close(restored(times), result["model"](times))
diff --git a/tests/test_plotting.py b/tests/test_plotting.py
new file mode 100644
index 0000000..e4bfd2b
--- /dev/null
+++ b/tests/test_plotting.py
@@ -0,0 +1,80 @@
+"""Scientific units and callable contracts in graphon visualizations."""
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+import torch
+
+from graphon_control.controls import TimeLabelGraphonControl
+from graphon_control.plotting import plot_graphon_diagnostics, sample_graphon_on_fine_grid
+
+
+def test_uniform_graphon_density_has_unit_row_and_column_mass():
+    axes = plot_graphon_diagnostics(torch.ones(5, 5), show=False)
+    try:
+        row_mass, column_mass = axes[2].lines
+        np.testing.assert_allclose(row_mass.get_ydata(), np.ones(5))
+        np.testing.assert_allclose(column_mass.get_ydata(), np.ones(5))
+    finally:
+        plt.close(axes[0].figure)
+
+
+def test_density_masses_integrate_receivers_and_senders_on_unit_interval():
+    axes = plot_graphon_diagnostics(torch.tensor([[2.0, 0.0], [1.0, 1.0]]), show=False)
+    try:
+        row_mass, column_mass = axes[2].lines
+        np.testing.assert_allclose(row_mass.get_ydata(), [1.0, 1.0])
+        np.testing.assert_allclose(column_mass.get_ydata(), [1.5, 0.5])
+    finally:
+        plt.close(axes[0].figure)
+
+
+def _controller():
+    xi = (torch.arange(5, dtype=torch.float64) + 0.5) / 5
+    return TimeLabelGraphonControl(xi, 1, 3, 4, 4)
+
+
+@pytest.mark.parametrize("explicit_grid", [False, True])
+def test_controller_sampler_returns_density_on_its_own_grid(explicit_grid):
+    model = _controller()
+    kwargs = {"M": model.N} if explicit_grid else {}
+    density = sample_graphon_on_fine_grid(model, t=0.3, **kwargs)
+    expected = model.N * model(0.3)
+    torch.testing.assert_close(density, expected)
+    assert density.shape == (model.N, model.N)
+    assert density.dtype == torch.float64
+    assert density.device.type == "cpu"
+    assert not density.requires_grad
+    torch.testing.assert_close(density.mean(-1), torch.ones(model.N, dtype=density.dtype))
+
+
+def test_controller_sampler_default_time_is_zero():
+    model = _controller()
+    torch.testing.assert_close(sample_graphon_on_fine_grid(model), model.N * model(0.0))
+
+
+def test_controller_sampler_rejects_unsupported_label_grid():
+    with pytest.raises(ValueError, match="arbitrary grids are unsupported"):
+        sample_graphon_on_fine_grid(_controller(), t=0.2, M=11)
+
+
+def test_controller_sampler_rejects_a_time_batch_even_with_one_entry():
+    with pytest.raises(ValueError, match="scalar time"):
+        sample_graphon_on_fine_grid(_controller(), t=torch.tensor([0.2]))
+
+
+def test_callable_sampler_evaluates_midpoint_grid_as_density():
+    density = sample_graphon_on_fine_grid(lambda xi, zeta: xi + 2 * zeta, M=2)
+    torch.testing.assert_close(density, torch.tensor([[0.75, 1.75], [1.25, 2.25]]))
+
+
+def test_callable_internal_typeerror_is_propagated_without_retry():
+    calls = []
+    def broken_kernel(t, x, xi, zeta):
+        calls.append((t, x))
+        raise TypeError("deliberate error within kernel")
+    with pytest.raises(TypeError, match="deliberate error within kernel"):
+        sample_graphon_on_fine_grid(broken_kernel, t=0.2, M=3)
+    assert calls == [(0.2, None)]
diff --git a/tests/test_training.py b/tests/test_training.py
new file mode 100644
index 0000000..7be3fed
--- /dev/null
+++ b/tests/test_training.py
@@ -0,0 +1,65 @@
+"""Optimization behavior and guards against corrupting a training run."""
+
+import pytest
+import torch
+
+from graphon_control.training import train_controller
+
+
+torch.set_num_threads(1)
+
+
+class ScalarController(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.tensor(0.0))
+
+
+def test_training_decreases_known_convex_objective_and_records_metrics():
+    model = ScalarController()
+    def objective():
+        loss = (model.weight - 2).square()
+        return {"total": loss, "squared_error": loss, "constant_metric": 4.0}
+    history, last = train_controller(model, objective, 20, lr=0.1, grad_clip=None)
+    assert history["loss"][-1] < history["loss"][0] / 10
+    assert len(history["squared_error"]) == 20
+    assert history["constant_metric"] == [4.0] * 20
+    assert last["total"] == history["loss"][-1]
+    assert all(isinstance(value, float) for value in last.values())
+
+
+def test_nonfinite_loss_is_rejected_before_update():
+    model = ScalarController()
+    before = model.weight.detach().clone()
+    def objective():
+        return {"total": model.weight * float("nan")}
+    with pytest.raises(FloatingPointError, match="total"):
+        train_controller(model, objective, 1, lr=0.1)
+    assert torch.equal(before, model.weight)
+
+
+def test_nonfinite_gradient_is_rejected_before_update():
+    model = ScalarController()
+    # sqrt(0) has a finite value and an infinite derivative.
+    def objective():
+        return {"total": torch.sqrt(model.weight)}
+    with pytest.raises(FloatingPointError, match="gradient"):
+        train_controller(model, objective, 1, lr=0.1)
+    assert model.weight.item() == 0
+
+
+def test_disconnected_objective_is_reported():
+    model = ScalarController()
+    def objective():
+        return {"total": torch.tensor(1.0, requires_grad=True)}
+    with pytest.raises(ValueError, match="disconnected"):
+        train_controller(model, objective, 1, lr=0.1)
+
+
+@pytest.mark.parametrize("kwargs", [{"num_steps": 0}, {"lr": 0}, {"grad_clip": -1}])
+def test_invalid_training_parameters(kwargs):
+    model = ScalarController()
+    settings = {"num_steps": 1, "lr": 0.1, "grad_clip": 1.0}
+    settings.update(kwargs)
+    with pytest.raises(ValueError):
+        train_controller(model, lambda: {"total": model.weight.square()}, **settings)
diff --git a/tests/test_workflows.py b/tests/test_workflows.py
new file mode 100644
index 0000000..6517bc5
--- /dev/null
+++ b/tests/test_workflows.py
@@ -0,0 +1,133 @@
+"""Full paper artifact workflows, including headless archive-only replotting."""
+
+import hashlib
+import json
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+from graphon_control.artifacts import load_controller
+from graphon_control.controls import TimeLabelGraphonControl
+from graphon_control.dynamics import precompute_Ws
+from graphon_control.experiments import bounded_confidence, kuramoto
+
+
+torch.set_num_threads(1)
+
+
+EXPERIMENTS = [
+    pytest.param(kuramoto, kuramoto.E1Config, "run_e1", "e1", "untwist.gif", id="kuramoto"),
+    pytest.param(bounded_confidence, bounded_confidence.E2Config, "run_e2", "e2", "opinions.gif", id="bounded-confidence"),
+]
+
+
+def _tiny_config(config_type):
+    return config_type(N=8, T=0.1, dt=0.05, num_steps=1, batch_size=2,
+                       num_fourier_freqs=1, embed_dim=3, hidden_embed=4,
+                       hidden_score=4, seed=3, evaluation_seed=29)
+
+
+def _numerical_hashes(run):
+    names = ("config.json", "provenance.json", "history.json", "metrics.json",
+             "model.pt", "trajectories.npz", "run_summary.txt")
+    return {name: hashlib.sha256((run / name).read_bytes()).hexdigest() for name in names}
+
+
+@pytest.mark.parametrize("module,config_type,runner,prefix,trajectory_gif", EXPERIMENTS)
+def test_full_workflow_replot_preserves_archived_numerics(
+    tmp_path, monkeypatch, module, config_type, runner, prefix, trajectory_gif,
+):
+    monkeypatch.delenv("DISPLAY", raising=False)
+    # Absolute output paths and a foreign cwd exercise the public run APIs without
+    # leaving plots in the source checkout or relying on repository-relative paths.
+    working = tmp_path / "foreign-cwd"
+    working.mkdir()
+    monkeypatch.chdir(working)
+    out_root, fig_root = tmp_path / "runs", tmp_path / "unrequested-mirrors"
+    config = _tiny_config(config_type)
+    getattr(module, runner)(config, tag="full", out_root=str(out_root),
+                            fig_root=str(fig_root), make_plots=True)
+    run = out_root / f"{prefix}_full"
+    marker = json.loads((run / "complete.json").read_text())
+    assert marker["status"] == "complete"
+    assert marker.get("initial_plots_generated", marker.get("plots_generated")) is True
+    assert not fig_root.exists()
+    assert list(working.iterdir()) == []
+    assert (run / "state_evolution.pdf").stat().st_size > 0
+    assert (run / "kernel_strip.pdf").stat().st_size > 0
+    assert (run / trajectory_gif).stat().st_size > 0
+    assert (run / "kernel_time.gif").stat().st_size > 0
+    provenance = json.loads((run / "provenance.json").read_text())
+    assert provenance["training_seed"] == config.seed
+    assert provenance["evaluation_seed"] == config.evaluation_seed
+    before = _numerical_hashes(run)
+    with np.load(run / "trajectories.npz", allow_pickle=False) as archive:
+        arrays = {name: archive[name].copy() for name in archive.files}
+    restored = TimeLabelGraphonControl(
+        torch.as_tensor(arrays["xi"]), config.num_fourier_freqs, config.embed_dim,
+        config.hidden_embed, config.hidden_score, time_scale=config.T,
+        score_bound=config.score_bound,
+    )
+    load_controller(run / "model.pt", restored)
+    with torch.no_grad():
+        reloaded_weights = precompute_Ws(restored, config.NT, config.dt, "cpu", config.time_chunk)
+    torch.testing.assert_close(reloaded_weights, torch.as_tensor(arrays["W_trained"]))
+
+    def unexpected_training(*args, **kwargs):
+        raise AssertionError("Replot must not train")
+    monkeypatch.setattr(module, "train_controller", unexpected_training)
+    module.plot_run(run)
+    assert _numerical_hashes(run) == before
+    assert json.loads((run / "complete.json").read_text()) == marker
+    with np.load(run / "trajectories.npz", allow_pickle=False) as archive:
+        assert set(archive.files) == set(arrays)
+        for name, expected in arrays.items():
+            np.testing.assert_array_equal(archive[name], expected)
+    # Existing tags are refused before calling the patched training entry point.
+    with pytest.raises(FileExistsError, match="Run already exists"):
+        getattr(module, runner)(config, tag="full", out_root=str(out_root), make_plots=False)
+    assert _numerical_hashes(run) == before
+
+
+@pytest.mark.parametrize("module,config_type,runner,prefix,trajectory_gif", EXPERIMENTS)
+def test_existing_mirror_is_refused_without_modifying_it(
+    tmp_path, module, config_type, runner, prefix, trajectory_gif,
+):
+    config = _tiny_config(config_type)
+    getattr(module, runner)(config, tag="export", out_root=str(tmp_path), make_plots=False)
+    run = tmp_path / f"{prefix}_export"
+    mirror = tmp_path / "selected-mirror"
+    mirror.mkdir()
+    existing = mirror / "state_evolution.pdf"
+    existing.write_bytes(b"previous selected figure")
+    before = _numerical_hashes(run)
+    with pytest.raises(FileExistsError):
+        module.plot_run(run, fig_dir=mirror)
+    assert existing.read_bytes() == b"previous selected figure"
+    assert _numerical_hashes(run) == before
+    assert not list(run.glob("*.pdf"))
+    assert not list(run.glob("*.gif"))
+    # A new explicit mirror can be generated from a numerical-only run. The
+    # completion marker tracks the resulting plots; scientific artifacts stay fixed.
+    fresh_mirror = tmp_path / "new-mirror"
+    module.plot_run(run, fig_dir=fresh_mirror)
+    assert (fresh_mirror / "state_evolution.pdf").stat().st_size > 0
+    assert (run / trajectory_gif).stat().st_size > 0
+    assert json.loads((run / "complete.json").read_text())["plots_generated"] is True
+    assert _numerical_hashes(run) == before
+
+
+@pytest.mark.parametrize("module,config_type,runner,prefix,trajectory_gif", EXPERIMENTS)
+def test_training_failure_has_no_completion_marker(
+    tmp_path, monkeypatch, module, config_type, runner, prefix, trajectory_gif,
+):
+    def fail(*args, **kwargs):
+        raise RuntimeError("deliberate optimizer failure")
+    monkeypatch.setattr(module, "train_controller", fail)
+    with pytest.raises(RuntimeError, match="deliberate optimizer failure"):
+        getattr(module, runner)(_tiny_config(config_type), tag="failed", out_root=str(tmp_path), make_plots=False)
+    directory = tmp_path / f"{prefix}_failed"
+    assert (directory / "provenance.json").is_file()
+    assert not (directory / "complete.json").exists()

From 03eac609ec2b6cd81176a062e0468a02062f51f3 Mon Sep 17 00:00:00 2001
From: Hydrodynamical <joseph.kent.miller@gmail.com>
Date: Tue, 20 Oct 2026 20:05:19 -0400
Subject: [PATCH 2/2] Use current CI actions and pin the runner image

---
 .github/workflows/tests.yml | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 112bcd8..6c80a65 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -1,15 +1,17 @@
 name: Numerical checks
 on: [push, pull_request]
+permissions:
+  contents: read
 jobs:
   test:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     env:
       MPLBACKEND: Agg
       OMP_NUM_THREADS: '1'
       MKL_NUM_THREADS: '1'
     steps:
-      - uses: actions/checkout@v4
-      - uses: actions/setup-python@v5
+      - uses: actions/checkout@v7.0.1
+      - uses: actions/setup-python@v7.0.0
         with:
           python-version: '3.12'
           cache: pip