diff --git a/giant/analysis/catalog.py b/giant/analysis/catalog.py
index a4a48d2..458965b 100644
--- a/giant/analysis/catalog.py
+++ b/giant/analysis/catalog.py
@@ -166,9 +166,7 @@ def _finalize_counts(merged: dict[str, list], key, nbins: int) -> list[int]:
     return list(merged.get(str(key), [0] * nbins))
 
 
-def _np_hist_pair(
-    r: np.ndarray, t: np.ndarray, nbins: int
-) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+def _np_hist_pair(r: np.ndarray, t: np.ndarray, nbins: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Shared-edge histogram of two small per-event arrays (robust range)."""
     both = np.concatenate([r, t]) if (len(r) or len(t)) else np.array([0.0, 1.0])
     lo, hi = float(np.quantile(both, 0.001)), float(np.quantile(both, 0.999))
@@ -240,9 +238,7 @@ def _marginal_overall_finalize(parts: list[dict], ctx: Context, var: str) -> Red
 
 def _energy_group_expr(lf: pl.LazyFrame, edges: np.ndarray) -> pl.Expr:
     ids, bins = event_energy_bins(lf, edges)
-    return pl.col("event_id").replace_strict(
-        ids, bins, default=-1, return_dtype=pl.Int64
-    )
+    return pl.col("event_id").replace_strict(ids, bins, default=-1, return_dtype=pl.Int64)
 
 
 def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict:
@@ -265,9 +261,7 @@ def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict:
     }
 
 
-def _marginal_grouped_finalize(
-    parts: list[dict], ctx: Context, var: str, axis: str
-) -> Reduced:
+def _marginal_grouped_finalize(parts: list[dict], ctx: Context, var: str, axis: str) -> Reduced:
     label, _ = _var(var)
     edges = _marginal_edges(ctx, var)
     nb = len(edges) - 1
@@ -317,9 +311,7 @@ def _event_scalar_partial(b: Bundle, col: str, use_all: bool) -> dict:
     return {"r": r.tolist(), "t": t.tolist()}
 
 
-def _event_scalar_finalize(
-    parts: list[dict], ctx: Context, spec_id: str, title: str, xlabel: str
-) -> Reduced:
+def _event_scalar_finalize(parts: list[dict], ctx: Context, spec_id: str, title: str, xlabel: str) -> Reduced:
     r = np.concatenate([np.asarray(p["r"], dtype=float) for p in parts])
     t = np.concatenate([np.asarray(p["t"], dtype=float) for p in parts])
     edges, rc, tc = _np_hist_pair(r, t, ctx.n_marginal_bins)
@@ -488,9 +480,7 @@ def _leakage_partial(b: Bundle) -> dict:
 
 def _leakage_finalize(parts: list[dict], ctx: Context) -> Reduced:
     frac = np.concatenate([np.asarray(p["frac"], dtype=float) for p in parts])
-    edges = np.linspace(
-        0.0, max(float(frac.max()) if len(frac) else 1.0, 1e-3), ctx.n_marginal_bins + 1
-    )
+    edges = np.linspace(0.0, max(float(frac.max()) if len(frac) else 1.0, 1e-3), ctx.n_marginal_bins + 1)
     counts = np.histogram(frac, edges)[0]
     return Reduced(
         id="leakage_fraction",
@@ -521,18 +511,8 @@ def _sec_frames(b: Bundle):
 
 def _sec_count_per_event_partial(b: Bundle) -> dict:
     r_sec, t_sec = _sec_frames(b)
-    r = (
-        r_sec.group_by("event_id")
-        .agg(pl.len().alias("n"))
-        .collect(engine="streaming")["n"]
-        .to_numpy()
-    )
-    t = (
-        t_sec.group_by("event_id")
-        .agg(pl.len().alias("n"))
-        .collect(engine="streaming")["n"]
-        .to_numpy()
-    )
+    r = r_sec.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")["n"].to_numpy()
+    t = t_sec.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")["n"].to_numpy()
     return {"r": r.tolist(), "t": t.tolist()}
 
 
@@ -568,9 +548,7 @@ def _sec_count_per_species_partial(b: Bundle) -> dict:
 def _sec_count_per_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
     r = sum_merge([p["r"] for p in parts])
     t = sum_merge([p["t"] for p in parts])
-    keys = sorted(set(r) | set(t), key=lambda k: -(r.get(k, 0) + t.get(k, 0)))[
-        : len(ctx.top_pdgs)
-    ]
+    keys = sorted(set(r) | set(t), key=lambda k: -(r.get(k, 0) + t.get(k, 0)))[: len(ctx.top_pdgs)]
     return Reduced(
         id="sec_count_per_species",
         family="secondaries",
@@ -617,11 +595,9 @@ def _sec_energy_finalize(parts: list[dict], ctx: Context) -> Reduced:
 
 def _sec_cos_angle_partial(b: Bundle) -> dict:
     edges = np.linspace(-1.0, 1.0, b.ctx.n_sec_bins + 1)
-    cos = (
-        pl.col("sdx") * pl.col("axis_x")
-        + pl.col("sdy") * pl.col("axis_y")
-        + pl.col("sdz") * pl.col("axis_z")
-    ).clip(-1.0, 1.0)
+    cos = (pl.col("sdx") * pl.col("axis_x") + pl.col("sdy") * pl.col("axis_y") + pl.col("sdz") * pl.col("axis_z")).clip(
+        -1.0, 1.0
+    )
 
     def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> dict[str, list[int]]:
         ea = entry_axis(steps_lf)
@@ -659,15 +635,13 @@ _router_gating_partial, _router_gating_finalize = _unchunkable(
     lambda b: compute_router_gating(b.checkpoint, b.r_phys, b.t_phys)
 )
 _router_share_pdg_partial, _router_share_pdg_finalize = _unchunkable(
-    lambda b: compute_router_share_by_pdg(
-        b.checkpoint, b.r_phys, b.t_phys, b.ctx.top_pdgs
-    )
+    lambda b: compute_router_share_by_pdg(b.checkpoint, b.r_phys, b.t_phys, b.ctx.top_pdgs)
 )
 _router_share_process_partial, _router_share_process_finalize = _unchunkable(
     lambda b: compute_router_share_by_process(b.checkpoint, b.t_phys)
 )
-_type_embedding_l1_distance_partial, _type_embedding_l1_distance_finalize = (
-    _unchunkable(lambda b: compute_type_embedding_l1_distance(b.type_embedding_l1_dist))
+_type_embedding_l1_distance_partial, _type_embedding_l1_distance_finalize = _unchunkable(
+    lambda b: compute_type_embedding_l1_distance(b.type_embedding_l1_dist)
 )
 
 
@@ -689,9 +663,7 @@ def build_catalog() -> list[PlotSpec]:
                 f"marginal_{var}",
                 "marginals",
                 compute_partial=lambda b, v=var: _marginal_overall_partial(b, v),
-                finalize=lambda parts, ctx, v=var: _marginal_overall_finalize(
-                    parts, ctx, v
-                ),
+                finalize=lambda parts, ctx, v=var: _marginal_overall_finalize(parts, ctx, v),
             )
         )
         for axis in GROUPING_AXES:
@@ -699,12 +671,8 @@ def build_catalog() -> list[PlotSpec]:
                 PlotSpec(
                     f"marginal_{var}_by_{axis}",
                     "marginals",
-                    compute_partial=lambda b, v=var, a=axis: _marginal_grouped_partial(
-                        b, v, a
-                    ),
-                    finalize=lambda parts, ctx, v=var, a=axis: (
-                        _marginal_grouped_finalize(parts, ctx, v, a)
-                    ),
+                    compute_partial=lambda b, v=var, a=axis: _marginal_grouped_partial(b, v, a),
+                    finalize=lambda parts, ctx, v=var, a=axis: _marginal_grouped_finalize(parts, ctx, v, a),
                 )
             )
 
@@ -712,9 +680,7 @@ def build_catalog() -> list[PlotSpec]:
         PlotSpec(
             "event_total_edep",
             "event",
-            compute_partial=lambda b: _event_scalar_partial(
-                b, "total_edep", use_all=True
-            ),
+            compute_partial=lambda b: _event_scalar_partial(b, "total_edep", use_all=True),
             finalize=lambda parts, ctx: _event_scalar_finalize(
                 parts,
                 ctx,
@@ -732,9 +698,7 @@ def build_catalog() -> list[PlotSpec]:
         PlotSpec(
             "event_mean_length",
             "event",
-            compute_partial=lambda b: _event_scalar_partial(
-                b, "mean_length", use_all=False
-            ),
+            compute_partial=lambda b: _event_scalar_partial(b, "mean_length", use_all=False),
             finalize=lambda parts, ctx: _event_scalar_finalize(
                 parts,
                 ctx,
@@ -746,9 +710,7 @@ def build_catalog() -> list[PlotSpec]:
         PlotSpec(
             "event_n_steps",
             "event",
-            compute_partial=lambda b: _event_scalar_partial(
-                b, "n_steps", use_all=False
-            ),
+            compute_partial=lambda b: _event_scalar_partial(b, "n_steps", use_all=False),
             finalize=lambda parts, ctx: _event_scalar_finalize(
                 parts,
                 ctx,
@@ -773,9 +735,7 @@ def build_catalog() -> list[PlotSpec]:
         PlotSpec(
             "shower_transverse",
             "shower",
-            compute_partial=lambda b: _profile_partial(
-                b, transverse_expr, "transverse_edges"
-            ),
+            compute_partial=lambda b: _profile_partial(b, transverse_expr, "transverse_edges"),
             finalize=lambda parts, ctx: _profile_finalize(
                 parts,
                 ctx,
diff --git a/giant/analysis/condor.py b/giant/analysis/condor.py
index 9cd1c64..f512621 100644
--- a/giant/analysis/condor.py
+++ b/giant/analysis/condor.py
@@ -161,9 +161,7 @@ class RunMeta:
         return cls(**json.loads(Path(path).read_text()))
 
 
-def _rows_per_chunk(
-    rollout: str | Path, reference: str | Path, n_chunks: int
-) -> list[int]:
+def _rows_per_chunk(rollout: str | Path, reference: str | Path, n_chunks: int) -> list[int]:
     """Rollout+reference row count of each ``event_id % n_chunks`` chunk.
 
     One cheap streaming ``group_by`` per side (just the ``event_id`` column) —
@@ -268,8 +266,7 @@ def compute_reduced(
     effective_n = n_chunks if spec.chunkable else 1
     if not (0 <= chunk_index < effective_n):
         raise ValueError(
-            f"{spec_id}: chunk_index={chunk_index} out of range for "
-            f"n_chunks={effective_n} (chunkable={spec.chunkable})"
+            f"{spec_id}: chunk_index={chunk_index} out of range for n_chunks={effective_n} (chunkable={spec.chunkable})"
         )
     bundle = Bundle.open(
         rollout,
@@ -327,10 +324,7 @@ def merge_one(spec_id: str, run_dir: str | Path) -> Path:
     effective_n = meta.n_chunks if spec.chunkable else 1
 
     partial_dir = run_path / "reduced_partial"
-    found = {
-        p.chunk: p
-        for p in (Partial.load(jf) for jf in partial_dir.glob(f"{spec_id}__*.json"))
-    }
+    found = {p.chunk: p for p in (Partial.load(jf) for jf in partial_dir.glob(f"{spec_id}__*.json"))}
     missing = sorted(set(range(effective_n)) - set(found))
     if missing:
         raise FileNotFoundError(
@@ -375,11 +369,7 @@ exec {giant_exe} analyze compute-one --id "$1" --chunk "$2" --run-dir {run_dir}
 
 
 def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> str:
-    reqs_attrs = (
-        "+RemoteJob = True\n"
-        if cfg.remote
-        else "requirements = TARGET.ProvidesETPResources\n"
-    )
+    reqs_attrs = "+RemoteJob = True\n" if cfg.remote else "requirements = TARGET.ProvidesETPResources\n"
     return (
         "universe = docker\n"
         f"docker_image = {cfg.docker_image}\n"
@@ -399,9 +389,7 @@ def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> st
     )
 
 
-def _job_walltimes(
-    run_dir: Path, ids: list[str], n_chunks: int
-) -> list[tuple[str, int, int]]:
+def _job_walltimes(run_dir: Path, ids: list[str], n_chunks: int) -> list[tuple[str, int, int]]:
     """``(spec_id, chunk, walltime_s)`` for every job, sized from ``run_meta.json``.
 
     Row counts come from ``prep``'s ``RunMeta.rows_per_chunk``/``total_rows``;
@@ -477,9 +465,7 @@ def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
     (run_dir / "reduced_partial").mkdir(parents=True, exist_ok=True)
 
     wrapper = run_dir / "run_compute.sh"
-    wrapper.write_text(
-        _WRAPPER.format(repo_dir=cfg.repo_dir, giant_exe=giant_exe, run_dir=run_dir)
-    )
+    wrapper.write_text(_WRAPPER.format(repo_dir=cfg.repo_dir, giant_exe=giant_exe, run_dir=run_dir))
     wrapper.chmod(0o755)
 
     jobs = _job_walltimes(run_dir, ids, cfg.n_chunks)
diff --git a/giant/analysis/context.py b/giant/analysis/context.py
index f8cd84b..090ead5 100644
--- a/giant/analysis/context.py
+++ b/giant/analysis/context.py
@@ -74,9 +74,7 @@ def _row_subsample(lf: pl.LazyFrame, sample_rows: int, seed: int) -> pl.LazyFram
     return lf.filter((pl.col("pre_E").hash(seed=seed) % 2**32) < threshold)
 
 
-def _combined_quantiles(
-    r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_q: float
-) -> tuple[float, float]:
+def _combined_quantiles(r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_q: float) -> tuple[float, float]:
     """Robust (lo_q, hi_q) range over the union of two value samples."""
     both = np.concatenate([r_vals, t_vals])
     lo, hi = float(np.quantile(both, lo_q)), float(np.quantile(both, hi_q))
@@ -104,31 +102,15 @@ def build_context(
 
     # Ranged marginal variables: robust ranges over a shared row subsample.
     exprs = [e.alias(n) for n, (_, e) in RANGED_VARS.items()]
-    r_s = (
-        _row_subsample(r_lf, sample_rows, seed)
-        .select(exprs)
-        .collect(engine="streaming")
-    )
-    t_s = (
-        _row_subsample(t_lf, sample_rows, seed)
-        .select(exprs)
-        .collect(engine="streaming")
-    )
+    r_s = _row_subsample(r_lf, sample_rows, seed).select(exprs).collect(engine="streaming")
+    t_s = _row_subsample(t_lf, sample_rows, seed).select(exprs).collect(engine="streaming")
     var_ranges = {
-        name: _combined_quantiles(
-            r_s[name].to_numpy(), t_s[name].to_numpy(), _LO_Q, _HI_Q
-        )
-        for name in RANGED_VARS
+        name: _combined_quantiles(r_s[name].to_numpy(), t_s[name].to_numpy(), _LO_Q, _HI_Q) for name in RANGED_VARS
     }
 
     # Energy-bin edges from exact per-event incident energies (cheap group_by).
     def _incident(lf: pl.LazyFrame) -> np.ndarray:
-        return (
-            lf.group_by("event_id")
-            .agg(pl.col("pre_E").max())
-            .collect(engine="streaming")["pre_E"]
-            .to_numpy()
-        )
+        return lf.group_by("event_id").agg(pl.col("pre_E").max()).collect(engine="streaming")["pre_E"].to_numpy()
 
     r_inc, t_inc = _incident(r_lf), _incident(t_lf)
     energy_edges = energy_bin_edges(np.concatenate([r_inc, t_inc]), n_energy_bins)
@@ -145,8 +127,7 @@ def build_context(
     )
     top_pdgs = [int(x) for x in pdg_counts["pdg"].to_list()[:top_k_pdg]]
     materials = sorted(
-        set(_counts(r_lf, "material")["material"].to_list())
-        | set(_counts(t_lf, "material")["material"].to_list())
+        set(_counts(r_lf, "material")["material"].to_list()) | set(_counts(t_lf, "material")["material"].to_list())
     )
 
     # Shower depth / transverse ranges from a subsampled proxy.
diff --git a/giant/analysis/grouping.py b/giant/analysis/grouping.py
index e5ec095..bc42005 100644
--- a/giant/analysis/grouping.py
+++ b/giant/analysis/grouping.py
@@ -68,9 +68,7 @@ def energy_bin_edges(incident_E: np.ndarray, n_bins: int = 4) -> np.ndarray:
 
 def energy_bin_labels(edges: np.ndarray) -> list[str]:
     """``E in [lo, hi)`` labels for each bin defined by ``edges`` (MeV)."""
-    return [
-        f"E in [{edges[i]:.3g}, {edges[i + 1]:.3g}) MeV" for i in range(len(edges) - 1)
-    ]
+    return [f"E in [{edges[i]:.3g}, {edges[i + 1]:.3g}) MeV" for i in range(len(edges) - 1)]
 
 
 def digitize_expr(value: pl.Expr, edges: np.ndarray) -> pl.Expr:
@@ -88,9 +86,7 @@ def digitize_expr(value: pl.Expr, edges: np.ndarray) -> pl.Expr:
     return idx.clip(0, n_bins - 1)
 
 
-def event_energy_bins(
-    lf: pl.LazyFrame, edges: np.ndarray
-) -> tuple[np.ndarray, np.ndarray]:
+def event_energy_bins(lf: pl.LazyFrame, edges: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     """Per-event incident-energy bin: ``(event_ids, bin_idx)`` numpy arrays.
 
     Incident energy is ``max(pre_E)`` per event (the primary). One bounded
diff --git a/giant/analysis/reduce.py b/giant/analysis/reduce.py
index 15089f6..b050563 100644
--- a/giant/analysis/reduce.py
+++ b/giant/analysis/reduce.py
@@ -142,9 +142,7 @@ def attach_entry_axis(lf: pl.LazyFrame, entry: pl.DataFrame) -> pl.LazyFrame:
     """
     ids = entry["event_id"].to_numpy()
     return lf.with_columns(
-        pl.col("event_id")
-        .replace_strict(ids, entry[col].to_numpy(), return_dtype=pl.Float64)
-        .alias(col)
+        pl.col("event_id").replace_strict(ids, entry[col].to_numpy(), return_dtype=pl.Float64).alias(col)
         for col in _ENTRY_AXIS_COLS
     )
 
@@ -254,10 +252,7 @@ def leakage_fraction(lf: pl.LazyFrame) -> np.ndarray:
         lf.group_by("event_id")
         .agg(
             pl.col("edep").sum().alias("deposited"),
-            pl.col("pre_E")
-            .filter(pl.col("termination_reason") == TERM_ESCAPED)
-            .sum()
-            .alias("escaped"),
+            pl.col("pre_E").filter(pl.col("termination_reason") == TERM_ESCAPED).sum().alias("escaped"),
         )
         .collect(engine="streaming")
     )
diff --git a/giant/analysis/render.py b/giant/analysis/render.py
index 862e023..f5f57ea 100644
--- a/giant/analysis/render.py
+++ b/giant/analysis/render.py
@@ -139,9 +139,7 @@ def _render_overlay(r: Reduced, params: dict):
 def _render_single(r: Reduced, params: dict):
     edges = np.asarray(r.payload["edges"])
     fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
-    ax.stairs(
-        _density(r.payload["rollout"], edges), edges, label=_SERIES_LABELS["rollout"]
-    )
+    ax.stairs(_density(r.payload["rollout"], edges), edges, label=_SERIES_LABELS["rollout"])
     if r.payload.get("log_y"):
         ax.set_yscale("log")
     if r.payload.get("log_x"):
@@ -187,9 +185,7 @@ def _render_profile(r: Reduced, params: dict):
         mean = np.asarray(r.payload[f"{key}_mean"])
         std = np.asarray(r.payload[f"{key}_std"])
         (line,) = ax.plot(centers, mean, label=_SERIES_LABELS[key])
-        ax.fill_between(
-            centers, mean - std, mean + std, alpha=0.2, color=line.get_color()
-        )
+        ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=line.get_color())
     ax.set_xlabel(r.xlabel)
     ax.set_ylabel(r.payload.get("ylabel", "mean deposited energy [MeV]"))
     ps.style_legend(ax, title="source")
@@ -201,9 +197,7 @@ def _render_bar(r: Reduced, params: dict):
     x = np.arange(len(labels))
     width = 0.4
     fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
-    ax.bar(
-        x - width / 2, r.payload["reference"], width, label=_SERIES_LABELS["reference"]
-    )
+    ax.bar(x - width / 2, r.payload["reference"], width, label=_SERIES_LABELS["reference"])
     ax.bar(x + width / 2, r.payload["rollout"], width, label=_SERIES_LABELS["rollout"])
     ax.set_xticks(x)
     ax.set_xticklabels(labels, rotation=45, ha="right")
@@ -215,9 +209,7 @@ def _render_bar(r: Reduced, params: dict):
 def _render_router_gating(r: Reduced, params: dict):
     n_experts = r.payload["n_experts"]
     log_x = r.payload.get("log_x", False)
-    fig, axes = ps.new_figure(
-        "slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False
-    )
+    fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False)
     flat = axes.ravel()
     for ax, key in zip(flat, ("rollout", "reference")):
         side = r.payload.get(key, {})
@@ -226,9 +218,7 @@ def _render_router_gating(r: Reduced, params: dict):
         if len(centers) and means.size:
             cum = np.zeros(len(centers))
             for i in range(n_experts):
-                ax.fill_between(
-                    centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}"
-                )
+                ax.fill_between(centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}")
                 cum = cum + means[:, i]
         if log_x:
             ax.set_xscale("log")
@@ -347,9 +337,7 @@ def render_all(
         families.add(r.family)
         fig = render(r, run_meta)
         ps.savefig(fig, str(family_dir / r.id), formats=("pdf",))
-        (family_dir / f"{r.id}.yaml").write_text(
-            yaml.safe_dump(_plot_metadata(r, run_meta), sort_keys=False)
-        )
+        (family_dir / f"{r.id}.yaml").write_text(yaml.safe_dump(_plot_metadata(r, run_meta), sort_keys=False))
         pdfs.append(family_dir / f"{r.id}.pdf")
         import matplotlib.pyplot as plt
 
@@ -370,9 +358,7 @@ def render_all(
     )
     for fam in families:
         (out_dir / fam / "metadata.yaml").write_text(
-            yaml.safe_dump(
-                {"title": fam, "description": f"{fam} plots."}, sort_keys=False
-            )
+            yaml.safe_dump({"title": fam, "description": f"{fam} plots."}, sort_keys=False)
         )
 
     if run_gallery:
@@ -400,6 +386,4 @@ def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]:
         "reference": meta.reference,
         **meta.plot_meta,
     }
-    return render_all(
-        run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery
-    )
+    return render_all(run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery)
diff --git a/giant/analysis/router_gating.py b/giant/analysis/router_gating.py
index 491a433..a13c287 100644
--- a/giant/analysis/router_gating.py
+++ b/giant/analysis/router_gating.py
@@ -95,9 +95,7 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None:
     # New nested shape (has a "stage1_model" key) vs. a v0.2 checkpoint's
     # flat model_config.
     router_cfg = (
-        (model_cfg.get("stage1_model") or {}).get("router")
-        if "stage1_model" in model_cfg
-        else model_cfg.get("router")
+        (model_cfg.get("stage1_model") or {}).get("router") if "stage1_model" in model_cfg else model_cfg.get("router")
     )
     if not router_cfg or not router_cfg.get("enabled"):
         return None
@@ -124,9 +122,7 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None:
     )
 
 
-def _subsample(
-    lf: pl.LazyFrame, n: int, seed: int, extra_cols: tuple = ()
-) -> pl.DataFrame:
+def _subsample(lf: pl.LazyFrame, n: int, seed: int, extra_cols: tuple = ()) -> pl.DataFrame:
     total = lf.select(pl.len()).collect(engine="streaming").item()
     if total > n:
         threshold = int(n / total * 2**32)
@@ -134,9 +130,7 @@ def _subsample(
     return lf.select(*_COLS, *extra_cols).collect(engine="streaming")
 
 
-def _gate_for_df(
-    handle: _RouterHandle, df: pl.DataFrame
-) -> tuple[pl.DataFrame, np.ndarray]:
+def _gate_for_df(handle: _RouterHandle, df: pl.DataFrame) -> tuple[pl.DataFrame, np.ndarray]:
     """(filtered df, gate_weights) for rows in ``df`` with a known pdg/material.
 
     Rows whose species or material never appeared in the checkpoint's
@@ -160,13 +154,9 @@ def _gate_for_df(
     df = df.filter(pl.Series(known, dtype=pl.Boolean))
 
     data = {
-        "pre_pos": np.column_stack(
-            [df["pre_x"].to_numpy(), df["pre_y"].to_numpy(), df["pre_z"].to_numpy()]
-        ),
+        "pre_pos": np.column_stack([df["pre_x"].to_numpy(), df["pre_y"].to_numpy(), df["pre_z"].to_numpy()]),
         "pre_E": df["pre_E"].to_numpy(),
-        "pre_dir": np.column_stack(
-            [df["pre_dx"].to_numpy(), df["pre_dy"].to_numpy(), df["pre_dz"].to_numpy()]
-        ),
+        "pre_dir": np.column_stack([df["pre_dx"].to_numpy(), df["pre_dy"].to_numpy(), df["pre_dz"].to_numpy()]),
         "layer_id": df["layer_id"].to_numpy(),
         "pdg": df["pdg"].to_numpy(),
         "material": df["material"].to_numpy(),
@@ -180,9 +170,7 @@ def _gate_for_df(
         material_conditioning=handle.material_conditioning,
     )
     with torch.no_grad():
-        gate = handle.router.gate(
-            torch.from_numpy(cond_cont).float(), torch.from_numpy(cond_cat).long()
-        ).numpy()
+        gate = handle.router.gate(torch.from_numpy(cond_cont).float(), torch.from_numpy(cond_cat).long()).numpy()
     return df, gate
 
 
@@ -205,9 +193,7 @@ def _quantile_bins(x: np.ndarray, gate: np.ndarray, n_bins: int) -> dict:
     return {"centers": centers[valid].tolist(), "means": means[valid].tolist()}
 
 
-def _top1_shares(
-    categories: np.ndarray, idx: np.ndarray, order: list, n_experts: int
-) -> dict[str, list[float]]:
+def _top1_shares(categories: np.ndarray, idx: np.ndarray, order: list, n_experts: int) -> dict[str, list[float]]:
     """Fraction of each category's rows hard-dispatched to each expert.
 
     Uses `Router.top1` (argmax), not the soft `gate` mean — grouped top-1
@@ -227,10 +213,7 @@ def _top1_shares(
     return shares
 
 
-_NOTE_NOT_MOE = (
-    "checkpoint has no enabled MoE router (model.router.enabled is "
-    "false/absent) — nothing to show"
-)
+_NOTE_NOT_MOE = "checkpoint has no enabled MoE router (model.router.enabled is false/absent) — nothing to show"
 
 _TITLES = {
     "router_gating": "Router gating (mixture-of-experts decision boundaries)",
@@ -266,9 +249,7 @@ def compute_router_gating(
         df = _subsample(lf, _SAMPLE_ROWS, seed)
         df, gate = _gate_for_df(handle, df)
         x = df["pre_E"].to_numpy()
-        sides[name] = (
-            _quantile_bins(x, gate, _N_BINS) if len(x) else {"centers": [], "means": []}
-        )
+        sides[name] = _quantile_bins(x, gate, _N_BINS) if len(x) else {"centers": [], "means": []}
 
     return Reduced(
         id="router_gating",
@@ -304,9 +285,7 @@ def compute_router_share_by_pdg(
         df, gate = _gate_for_df(handle, df)
         if len(df):
             idx = gate.argmax(axis=1)
-            shares = _top1_shares(
-                df["pdg"].to_numpy(), idx, top_pdgs, handle.router.n_experts
-            )
+            shares = _top1_shares(df["pdg"].to_numpy(), idx, top_pdgs, handle.router.n_experts)
         else:
             shares = {str(p): [0.0] * handle.router.n_experts for p in top_pdgs}
         sides[name] = {labels[i]: shares[str(p)] for i, p in enumerate(top_pdgs)}
@@ -351,9 +330,7 @@ def compute_router_share_by_process(
         counts = df["process"].value_counts().sort("count", descending=True)
         order = counts["process"].to_list()[:top_k]
         idx = gate.argmax(axis=1)
-        shares = _top1_shares(
-            df["process"].to_numpy(), idx, order, handle.router.n_experts
-        )
+        shares = _top1_shares(df["process"].to_numpy(), idx, order, handle.router.n_experts)
     else:
         order, shares = [], {}
 
diff --git a/giant/analysis/runtime_estimate.py b/giant/analysis/runtime_estimate.py
index f5279f7..7f6cc37 100644
--- a/giant/analysis/runtime_estimate.py
+++ b/giant/analysis/runtime_estimate.py
@@ -54,9 +54,7 @@ _FIXED_OVERHEAD_S = 60.0
 # scan. Calibrated from the 3 real router jobs' observed wall times (119, 66,
 # 124s) — max minus _FIXED_OVERHEAD_S, on top of it.
 _ROUTER_FIXED_S = 64.0
-_ROUTER_IDS = frozenset(
-    {"router_gating", "router_share_by_pdg", "router_share_by_process"}
-)
+_ROUTER_IDS = frozenset({"router_gating", "router_share_by_pdg", "router_share_by_process"})
 
 # Conservative fallback for any catalog id not in _COST_MODEL (e.g. a plot
 # added after the last calibration run) — the most expensive fitted per-row
diff --git a/giant/analysis/sources.py b/giant/analysis/sources.py
index 2248c90..abedc74 100644
--- a/giant/analysis/sources.py
+++ b/giant/analysis/sources.py
@@ -93,10 +93,7 @@ def _check_rollout_metadata(path: Path) -> None:
     metadata = pq.read_schema(path).metadata or {}
     coord = metadata.get(PREDICT_COORD_METADATA_KEY.encode())
     if coord is not None and coord.decode() != ROLLOUT_COORD_VALUE:
-        raise ValueError(
-            f"{path} is not a rollout file (coord={coord.decode()!r}); "
-            "expected `giant rollout` output"
-        )
+        raise ValueError(f"{path} is not a rollout file (coord={coord.decode()!r}); expected `giant rollout` output")
 
 
 def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
@@ -121,11 +118,7 @@ def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
     else:
         # The reference (a rollout's seed `dataset`) may be a directory of
         # parquet shards rather than a single file — scan them all.
-        lf = (
-            pl.scan_parquet(str(path / "**/*.parquet"))
-            if path.is_dir()
-            else pl.scan_parquet(path)
-        )
+        lf = pl.scan_parquet(str(path / "**/*.parquet")) if path.is_dir() else pl.scan_parquet(path)
     return lf.with_columns(pl.col("pdg").cast(pl.Int64))
 
 
@@ -137,9 +130,7 @@ def physical_steps(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
     """
     if side is Side.reference:
         return lf
-    return lf.filter(
-        ~pl.col("termination_reason").is_in(list(SYNTHETIC_TERMINATION_REASONS))
-    )
+    return lf.filter(~pl.col("termination_reason").is_in(list(SYNTHETIC_TERMINATION_REASONS)))
 
 
 def secondaries(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
diff --git a/giant/cli.py b/giant/cli.py
index d4a5229..ee49c1d 100644
--- a/giant/cli.py
+++ b/giant/cli.py
@@ -65,9 +65,7 @@ def _router_total_experts(router_cfg: dict) -> int:
     """
     if router_cfg.get("type") == "composed":
         axis_counts = {
-            m.group(1): int(v)
-            for k, v in router_cfg.items()
-            if (m := re.match(r"^axis(\d+)_n_experts$", k))
+            m.group(1): int(v) for k, v in router_cfg.items() if (m := re.match(r"^axis(\d+)_n_experts$", k))
         }
         return math.prod(axis_counts.values()) if axis_counts else 1
     return int(router_cfg.get("n_experts", 1))
@@ -104,11 +102,7 @@ def _ddpm_steps(model_cfg: dict, stage: str) -> int:
 
 
 def _particle_type_other_policy(model_cfg: dict) -> str:
-    return (
-        _stage_cfg(model_cfg, "stage2")
-        .get("particle_type", {})
-        .get("other_policy", "sample")
-    )
+    return _stage_cfg(model_cfg, "stage2").get("particle_type", {}).get("other_policy", "sample")
 
 
 def _load_pdg_topn_map(ckpt: dict):
@@ -129,9 +123,7 @@ def _load_mat_topn_map(ckpt: dict):
     return topnmap_from_json(raw, axis="material") if raw is not None else None
 
 
-def _batch_size_estimate_dims(
-    model_cfg: dict, training: bool, stage: str = "stage1"
-) -> tuple[int, int]:
+def _batch_size_estimate_dims(model_cfg: dict, training: bool, stage: str = "stage1") -> tuple[int, int]:
     """Pick the (hidden_dim, n_blocks) that dominate per-call activation memory.
 
     `model_cfg` is either the new nested shape (has a `f"{stage}_model"` key
@@ -336,14 +328,10 @@ def _load_model_weights(
 
 @app.command()
 def train(
-    data: Annotated[
-        Path, typer.Argument(help="Parquet file or directory of parquet files")
-    ],
+    data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")],
     config: Annotated[
         Optional[Path],
-        typer.Option(
-            "--config", "-c", help="TOML config file (overridden by explicit flags)"
-        ),
+        typer.Option("--config", "-c", help="TOML config file (overridden by explicit flags)"),
     ] = None,
     mode: Annotated[
         Optional[Mode],
@@ -355,8 +343,7 @@ def train(
         typer.Option(
             "--batch-size",
             "-b",
-            help="Integer, or 'auto' to estimate from free GPU memory "
-            "(cuda devices only)",
+            help="Integer, or 'auto' to estimate from free GPU memory (cuda devices only)",
         ),
     ] = None,
     lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None,
@@ -372,17 +359,13 @@ def train(
             "alongside the raw weights in checkpoints (0 disables; default: 0.9999)",
         ),
     ] = None,
-    warmup_epochs: Annotated[
-        Optional[int], typer.Option("--warmup-epochs", "-w")
-    ] = None,
+    warmup_epochs: Annotated[Optional[int], typer.Option("--warmup-epochs", "-w")] = None,
     hidden_dim: Annotated[Optional[int], typer.Option("--hidden-dim", "-H")] = None,
     n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None,
     emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None,
     dropout: Annotated[
         Optional[float],
-        typer.Option(
-            "--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"
-        ),
+        typer.Option("--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"),
     ] = None,
     stage1_generator: Annotated[
         Optional[Mode],
@@ -402,9 +385,7 @@ def train(
     ] = None,
     stage1_n_res_blocks: Annotated[
         Optional[int],
-        typer.Option(
-            "--stage1-n-res-blocks", help="Overrides --n-blocks for stage 1 only"
-        ),
+        typer.Option("--stage1-n-res-blocks", help="Overrides --n-blocks for stage 1 only"),
     ] = None,
     stage1_dropout: Annotated[
         Optional[float],
@@ -452,8 +433,7 @@ def train(
         Optional[int],
         typer.Option(
             "--stage2-context-dim",
-            help="Width of the projected stage-1 outcome fed into stage 2's "
-            "conditioning (default: 64)",
+            help="Width of the projected stage-1 outcome fed into stage 2's conditioning (default: 64)",
         ),
     ] = None,
     stage2_stage1_context: Annotated[
@@ -485,13 +465,9 @@ def train(
     ] = None,
     router_type: Annotated[
         Optional[str],
-        typer.Option(
-            "--router-type", help="Router implementation name (see ROUTER_REGISTRY)"
-        ),
-    ] = None,
-    n_experts: Annotated[
-        Optional[int], typer.Option("--n-experts", help="Number of routed experts")
+        typer.Option("--router-type", help="Router implementation name (see ROUTER_REGISTRY)"),
     ] = None,
+    n_experts: Annotated[Optional[int], typer.Option("--n-experts", help="Number of routed experts")] = None,
     router_axis: Annotated[
         Optional[list[str]],
         typer.Option(
@@ -506,32 +482,28 @@ def train(
         Optional[int],
         typer.Option(
             "--n-critic",
-            help="WGAN-GP (--mode wgan only): critic updates per generator "
-            "update (default: 5)",
+            help="WGAN-GP (--mode wgan only): critic updates per generator update (default: 5)",
         ),
     ] = None,
     gp_weight: Annotated[
         Optional[float],
         typer.Option(
             "--gp-weight",
-            help="WGAN-GP (--mode wgan only): gradient-penalty coefficient "
-            "(default: 10.0)",
+            help="WGAN-GP (--mode wgan only): gradient-penalty coefficient (default: 10.0)",
         ),
     ] = None,
     noise_dim: Annotated[
         Optional[int],
         typer.Option(
             "--noise-dim",
-            help="WGAN (--mode wgan only): generator input noise-vector "
-            "width (default: 64)",
+            help="WGAN (--mode wgan only): generator input noise-vector width (default: 64)",
         ),
     ] = None,
     critic_lr: Annotated[
         Optional[float],
         typer.Option(
             "--critic-lr",
-            help="WGAN-GP (--mode wgan only): critic learning rate "
-            "(default: same as --lr)",
+            help="WGAN-GP (--mode wgan only): critic learning rate (default: same as --lr)",
         ),
     ] = None,
     stage1_n_critic: Annotated[
@@ -540,21 +512,15 @@ def train(
     ] = None,
     stage1_gp_weight: Annotated[
         Optional[float],
-        typer.Option(
-            "--stage1-gp-weight", help="Overrides --gp-weight for stage 1 only"
-        ),
+        typer.Option("--stage1-gp-weight", help="Overrides --gp-weight for stage 1 only"),
     ] = None,
     stage1_noise_dim: Annotated[
         Optional[int],
-        typer.Option(
-            "--stage1-noise-dim", help="Overrides --noise-dim for stage 1 only"
-        ),
+        typer.Option("--stage1-noise-dim", help="Overrides --noise-dim for stage 1 only"),
     ] = None,
     stage1_critic_lr: Annotated[
         Optional[float],
-        typer.Option(
-            "--stage1-critic-lr", help="Overrides --critic-lr for stage 1 only"
-        ),
+        typer.Option("--stage1-critic-lr", help="Overrides --critic-lr for stage 1 only"),
     ] = None,
     stage2_n_critic: Annotated[
         Optional[int],
@@ -562,9 +528,7 @@ def train(
     ] = None,
     stage2_gp_weight: Annotated[
         Optional[float],
-        typer.Option(
-            "--stage2-gp-weight", help="Overrides --gp-weight for stage 2 only"
-        ),
+        typer.Option("--stage2-gp-weight", help="Overrides --gp-weight for stage 2 only"),
     ] = None,
     stage2_noise_dim: Annotated[
         Optional[int],
@@ -576,13 +540,9 @@ def train(
     ] = None,
     stage2_critic_lr: Annotated[
         Optional[float],
-        typer.Option(
-            "--stage2-critic-lr", help="Overrides --critic-lr for stage 2 only"
-        ),
-    ] = None,
-    val_fraction: Annotated[
-        Optional[float], typer.Option("--val-fraction", "-f")
+        typer.Option("--stage2-critic-lr", help="Overrides --critic-lr for stage 2 only"),
     ] = None,
+    val_fraction: Annotated[Optional[float], typer.Option("--val-fraction", "-f")] = None,
     seed: Annotated[
         Optional[int],
         typer.Option("--seed", "-s", help="Random seed for reproducibility"),
@@ -608,15 +568,12 @@ def train(
         Optional[int],
         typer.Option(
             "--max-val-batches",
-            help="Cap the per-epoch val-loss pass to N batches (0 = full "
-            "val set every epoch; default: 200)",
+            help="Cap the per-epoch val-loss pass to N batches (0 = full val set every epoch; default: 200)",
         ),
     ] = None,
     shuffle_buffer: Annotated[
         int,
-        typer.Option(
-            "--shuffle-buffer", "-B", help="Rows held in RAM per worker for shuffling"
-        ),
+        typer.Option("--shuffle-buffer", "-B", help="Rows held in RAM per worker for shuffling"),
     ] = 65536,
     cache_setup: Annotated[
         bool,
@@ -660,8 +617,7 @@ def train(
         Optional[bool],
         typer.Option(
             "--wandb/--no-wandb",
-            help="Log per-epoch training metrics to Weights & Biases "
-            "(requires `uv sync --extra wandb`)",
+            help="Log per-epoch training metrics to Weights & Biases (requires `uv sync --extra wandb`)",
         ),
     ] = None,
     wandb_project: Annotated[
@@ -692,8 +648,7 @@ def train(
                 batch_size_value = int(batch_size)
             except ValueError:
                 typer.echo(
-                    f"error: --batch-size must be an integer or 'auto', "
-                    f"got {batch_size!r}",
+                    f"error: --batch-size must be an integer or 'auto', got {batch_size!r}",
                     err=True,
                 )
                 raise typer.Exit(1)
@@ -753,9 +708,7 @@ def train(
             "decoder": stage2_decoder.value if stage2_decoder is not None else None,
             "k_max": stage2_k_max,
             "context_dim": stage2_context_dim,
-            "stage1_context": stage2_stage1_context.value
-            if stage2_stage1_context is not None
-            else None,
+            "stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
         }.items()
         if v is not None
     }
@@ -834,9 +787,7 @@ def train(
     ):
         stage_wgan = {**shared_wgan_overrides, **stage_specific}
         if stage_wgan:
-            overrides.setdefault(stage_name, {}).setdefault("wgan", {}).update(
-                stage_wgan
-            )
+            overrides.setdefault(stage_name, {}).setdefault("wgan", {}).update(stage_wgan)
 
     cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
     gconfig.validate_config(cfg)
@@ -845,19 +796,13 @@ def train(
     _device = torch.device(device) if device else gconfig.auto_device()
 
     if batch_size_auto:
-        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(
-            cfg, training=True, stage="stage1"
-        )
+        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(cfg, training=True, stage="stage1")
         try:
-            t["batch_size"] = gconfig.estimate_batch_size(
-                est_hidden_dim, est_n_blocks, _device
-            )
+            t["batch_size"] = gconfig.estimate_batch_size(est_hidden_dim, est_n_blocks, _device)
         except ValueError as exc:
             typer.echo(f"error: {exc}", err=True)
             raise typer.Exit(1)
-        typer.echo(
-            f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)"
-        )
+        typer.echo(f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)")
 
     if out is not None:
         out_dir = out
@@ -913,39 +858,19 @@ def new_run(
     n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None,
     emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None,
     dropout: Annotated[Optional[float], typer.Option("--dropout", "-d")] = None,
-    stage1_generator: Annotated[
-        Optional[Mode], typer.Option("--stage1-generator")
-    ] = None,
-    stage1_hidden_dim: Annotated[
-        Optional[int], typer.Option("--stage1-hidden-dim")
-    ] = None,
-    stage1_n_res_blocks: Annotated[
-        Optional[int], typer.Option("--stage1-n-res-blocks")
-    ] = None,
+    stage1_generator: Annotated[Optional[Mode], typer.Option("--stage1-generator")] = None,
+    stage1_hidden_dim: Annotated[Optional[int], typer.Option("--stage1-hidden-dim")] = None,
+    stage1_n_res_blocks: Annotated[Optional[int], typer.Option("--stage1-n-res-blocks")] = None,
     stage1_dropout: Annotated[Optional[float], typer.Option("--stage1-dropout")] = None,
-    stage2_generator: Annotated[
-        Optional[Mode], typer.Option("--stage2-generator")
-    ] = None,
-    stage2_hidden_dim: Annotated[
-        Optional[int], typer.Option("--stage2-hidden-dim")
-    ] = None,
-    stage2_n_res_blocks: Annotated[
-        Optional[int], typer.Option("--stage2-n-res-blocks")
-    ] = None,
+    stage2_generator: Annotated[Optional[Mode], typer.Option("--stage2-generator")] = None,
+    stage2_hidden_dim: Annotated[Optional[int], typer.Option("--stage2-hidden-dim")] = None,
+    stage2_n_res_blocks: Annotated[Optional[int], typer.Option("--stage2-n-res-blocks")] = None,
     stage2_dropout: Annotated[Optional[float], typer.Option("--stage2-dropout")] = None,
-    stage2_decoder: Annotated[
-        Optional[Decoder], typer.Option("--stage2-decoder")
-    ] = None,
+    stage2_decoder: Annotated[Optional[Decoder], typer.Option("--stage2-decoder")] = None,
     stage2_k_max: Annotated[Optional[int], typer.Option("--stage2-k-max")] = None,
-    stage2_context_dim: Annotated[
-        Optional[int], typer.Option("--stage2-context-dim")
-    ] = None,
-    stage2_stage1_context: Annotated[
-        Optional[Stage1Context], typer.Option("--stage2-stage1-context")
-    ] = None,
-    conditioning: Annotated[
-        Optional[Conditioning], typer.Option("--conditioning")
-    ] = None,
+    stage2_context_dim: Annotated[Optional[int], typer.Option("--stage2-context-dim")] = None,
+    stage2_stage1_context: Annotated[Optional[Stage1Context], typer.Option("--stage2-stage1-context")] = None,
+    conditioning: Annotated[Optional[Conditioning], typer.Option("--conditioning")] = None,
     router: Annotated[Optional[bool], typer.Option("--router/--no-router")] = None,
     router_type: Annotated[Optional[str], typer.Option("--router-type")] = None,
     n_experts: Annotated[Optional[int], typer.Option("--n-experts")] = None,
@@ -956,16 +881,13 @@ def new_run(
     ] = None,
     comment: Annotated[
         Optional[str],
-        typer.Option(
-            "--comment", help="Free-text note recorded in config.toml's meta section"
-        ),
+        typer.Option("--comment", help="Free-text note recorded in config.toml's meta section"),
     ] = None,
     data: Annotated[
         Optional[Path],
         typer.Option(
             "--data",
-            help="Dataset path to fill in the printed next-step command "
-            "(not stored in the config)",
+            help="Dataset path to fill in the printed next-step command (not stored in the config)",
         ),
     ] = None,
     force: Annotated[
@@ -977,9 +899,7 @@ def new_run(
     ] = False,
     dry_run: Annotated[
         bool,
-        typer.Option(
-            "--dry-run", help="Print the resolved config without writing anything"
-        ),
+        typer.Option("--dry-run", help="Print the resolved config without writing anything"),
     ] = False,
 ) -> None:
     """Scaffold a new training run: resolve hyperparams to a config.toml and lay out its run dir.
@@ -1030,9 +950,7 @@ def new_run(
             "decoder": stage2_decoder.value if stage2_decoder is not None else None,
             "k_max": stage2_k_max,
             "context_dim": stage2_context_dim,
-            "stage1_context": stage2_stage1_context.value
-            if stage2_stage1_context is not None
-            else None,
+            "stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
         }.items()
         if v is not None
     }
@@ -1116,9 +1034,7 @@ def new_run(
 
 @app.command()
 def predict(
-    data: Annotated[
-        Path, typer.Argument(help="Parquet file or directory of parquet files")
-    ],
+    data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")],
     checkpoint: Annotated[
         Path,
         typer.Option(
@@ -1151,8 +1067,7 @@ def predict(
         typer.Option(
             "--batch-size",
             "-b",
-            help="Inference batch size, or 'auto' to estimate from free GPU "
-            "memory (cuda devices only)",
+            help="Inference batch size, or 'auto' to estimate from free GPU memory (cuda devices only)",
         ),
     ] = "4096",
     steps: Annotated[
@@ -1221,8 +1136,7 @@ def predict(
 
     if "sec_phys" not in ckpt.get("normalizer", {}):
         typer.echo(
-            "error: checkpoint has no normalizer.sec_phys — retrain with the "
-            "current code",
+            "error: checkpoint has no normalizer.sec_phys — retrain with the current code",
             err=True,
         )
         raise typer.Exit(1)
@@ -1235,9 +1149,7 @@ def predict(
     stage2_k_max = _stage_cfg(model_cfg, "stage2").get("k_max", K_MAX)
 
     if batch_size_auto:
-        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(
-            model_cfg, training=False
-        )
+        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(model_cfg, training=False)
         try:
             batch_size_value = gconfig.estimate_batch_size(
                 est_hidden_dim,
@@ -1248,9 +1160,7 @@ def predict(
         except ValueError as exc:
             typer.echo(f"error: {exc}", err=True)
             raise typer.Exit(1)
-        typer.echo(
-            f"batch_size: {batch_size_value} (auto-estimated from free GPU memory)"
-        )
+        typer.echo(f"batch_size: {batch_size_value} (auto-estimated from free GPU memory)")
 
     assert batch_size_value is not None
     bs = batch_size_value
@@ -1311,16 +1221,10 @@ def predict(
             return iter_file_chunks(path, offset=offset, k_max=stage2_k_max)
         return iter_cond_chunks(path, offset=offset)
 
-    cond_pdg_topn = (
-        pdg_topn_map.class_map if particle_conditioning == "onehot" else None
-    )
-    cond_mat_topn = (
-        mat_topn_map.class_map if material_conditioning == "onehot" else None
-    )
+    cond_pdg_topn = pdg_topn_map.class_map if particle_conditioning == "onehot" else None
+    cond_mat_topn = mat_topn_map.class_map if material_conditioning == "onehot" else None
 
-    def _concat(
-        a: dict[str, np.ndarray], b: dict[str, np.ndarray]
-    ) -> dict[str, np.ndarray]:
+    def _concat(a: dict[str, np.ndarray], b: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
         return {k: np.concatenate([a[k], b[k]], axis=0) for k in a}
 
     def _process(piece: dict[str, np.ndarray]) -> None:
@@ -1352,19 +1256,13 @@ def predict(
 
         cc = torch.from_numpy(cond_cont).float().to(_device)
         ck = torch.from_numpy(cond_cat).long().to(_device)
-        stage1_norm, n_sec_pred = sample_stage1(
-            model, cc, ck, steps=steps, ddpm_steps=stage1_ddpm_steps
-        )
+        stage1_norm, n_sec_pred = sample_stage1(model, cc, ck, steps=steps, ddpm_steps=stage1_ddpm_steps)
 
         if coord == Coord.global_:
             # A fresh v0.3.0 Stage1Model owns no n_sec_head —
             # sample_stage1 returns n_sec_pred=None then, so ask stage 2.
-            n_sec_pred = resolve_n_sec(
-                model, sec_decoder, cc, ck, stage1_norm, n_sec_pred
-            )
-            sec_cont, sec_type, _sec_valid_pred = sample_stage2(
-                sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps
-            )
+            n_sec_pred = resolve_n_sec(model, sec_decoder, cc, ck, stage1_norm, n_sec_pred)
+            sec_cont, sec_type, _sec_valid_pred = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
             n_sec_pred_np = n_sec_pred.cpu().numpy()
 
         pred = stage1_norm.cpu().numpy()  # normalised
@@ -1387,14 +1285,8 @@ def predict(
                     "material": piece["material"],
                     "layer_id": piece["layer_id"],
                     "n_sec": piece["n_sec"],
-                    **{
-                        f"pred_{name}": raw[:, j]
-                        for j, name in enumerate(LOCAL_TARGET_NAMES)
-                    },
-                    **{
-                        f"true_{name}": target_raw[:, j]
-                        for j, name in enumerate(LOCAL_TARGET_NAMES)
-                    },
+                    **{f"pred_{name}": raw[:, j] for j, name in enumerate(LOCAL_TARGET_NAMES)},
+                    **{f"true_{name}": target_raw[:, j] for j, name in enumerate(LOCAL_TARGET_NAMES)},
                 }
             )
         else:
@@ -1404,9 +1296,7 @@ def predict(
             # (hence delta_e == edep + e_sec) holds by construction. e_sec_pred
             # doubles as the stick-breaking energy budget for the Stage-2 decode
             # below, since the model has no other source for it at inference.
-            edep, e_sec_pred, _post_E, delta_e = energy_simplex_decode(
-                raw[:, 1:3], piece["pre_E"]
-            )
+            edep, e_sec_pred, _post_E, delta_e = energy_simplex_decode(raw[:, 1:3], piece["pre_E"])
 
             # Normalise predicted direction then rotate back to world frame
             post_dir_local = raw[:, 3:6].copy()
@@ -1419,43 +1309,31 @@ def predict(
             travel_dir_local = raw[:, 6:9].copy()
             norms = np.linalg.norm(travel_dir_local, axis=1, keepdims=True)
             travel_dir_local /= np.where(norms < 1e-8, 1.0, norms)
-            post_pos_world = reconstruct_post_pos(
-                piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local
-            )
+            post_pos_world = reconstruct_post_pos(piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local)
 
             # particle_type.target="physical": sec_pdg_code is a reporting-
             # only nearest-known-PDG label (never fed back into the model —
             # "no snapping at inference"). "onehot"/"embedding": PDG
             # resolution IS the secondary's identity — see
             # decode_secondary_identity's docstring.
-            sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, _l1_dist = (
-                decode_secondary_identity(
-                    sec_decoder,
-                    sec_cont,
-                    sec_type,
-                    n_sec_pred_np,
-                    e_sec_pred,
-                    piece["pre_dir"],
-                    sec_phys_norm,
-                    pdg_map,
-                    pdg_topn_map,
-                    other_policy,
-                    None,
-                )
+            sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, _l1_dist = decode_secondary_identity(
+                sec_decoder,
+                sec_cont,
+                sec_type,
+                n_sec_pred_np,
+                e_sec_pred,
+                piece["pre_dir"],
+                sec_phys_norm,
+                pdg_map,
+                pdg_topn_map,
+                other_policy,
+                None,
             )
-            sec_pdg_list = [
-                sec_pdg_code[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)
-            ]
+            sec_pdg_list = [sec_pdg_code[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)]
             sec_E_list = [sec_E[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)]
-            sec_dx_list = [
-                sec_dir_world[i, :n, 0].tolist() for i, n in enumerate(n_sec_pred_np)
-            ]
-            sec_dy_list = [
-                sec_dir_world[i, :n, 1].tolist() for i, n in enumerate(n_sec_pred_np)
-            ]
-            sec_dz_list = [
-                sec_dir_world[i, :n, 2].tolist() for i, n in enumerate(n_sec_pred_np)
-            ]
+            sec_dx_list = [sec_dir_world[i, :n, 0].tolist() for i, n in enumerate(n_sec_pred_np)]
+            sec_dy_list = [sec_dir_world[i, :n, 1].tolist() for i, n in enumerate(n_sec_pred_np)]
+            sec_dz_list = [sec_dir_world[i, :n, 2].tolist() for i, n in enumerate(n_sec_pred_np)]
 
             table = pa.table(
                 {
@@ -1537,9 +1415,7 @@ def predict(
     typer.echo(f"reference: {ref_path}")
 
     if skipped:
-        codes = ", ".join(
-            f"{pdg} ({count})" for pdg, count in sorted(unknown_pdg_counts.items())
-        )
+        codes = ", ".join(f"{pdg} ({count})" for pdg, count in sorted(unknown_pdg_counts.items()))
         typer.echo(
             f"warning: skipped {skipped:,} row(s) with unknown PDG code(s): {codes}",
             err=True,
@@ -1587,9 +1463,7 @@ def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.nda
 
 @app.command()
 def rollout(
-    data: Annotated[
-        Path, typer.Argument(help="Parquet file/dir to seed showers from (real events)")
-    ],
+    data: Annotated[Path, typer.Argument(help="Parquet file/dir to seed showers from (real events)")],
     checkpoint: Annotated[
         Path,
         typer.Option("--checkpoint", "-c", help="Checkpoint .pt (best.pt/last.pt)"),
@@ -1609,9 +1483,7 @@ def rollout(
             help="Stop a track when its energy drops below this [MeV]",
         ),
     ] = 0.1,
-    max_steps: Annotated[
-        int, typer.Option("--max-steps", help="Max steps per individual track")
-    ] = 1000,
+    max_steps: Annotated[int, typer.Option("--max-steps", help="Max steps per individual track")] = 1000,
     steps: Annotated[
         int,
         typer.Option(
@@ -1629,9 +1501,7 @@ def rollout(
             "requires a checkpoint trained with EMA enabled.",
         ),
     ] = Weights.raw,
-    batch_size: Annotated[
-        int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward")
-    ] = 4096,
+    batch_size: Annotated[int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward")] = 4096,
     max_tracks_per_event: Annotated[
         Optional[int],
         typer.Option(
@@ -1646,15 +1516,9 @@ def rollout(
             help="Override the oracle's NN-distance escape threshold [mm]",
         ),
     ] = None,
-    n_events: Annotated[
-        Optional[int], typer.Option("--n-events", help="Cap number of seed events")
-    ] = None,
-    device: Annotated[
-        Optional[str], typer.Option("--device", "-d", help="cpu | cuda | mps (auto)")
-    ] = None,
-    out: Annotated[
-        Optional[Path], typer.Option("--out", "-o", help="Output steps parquet")
-    ] = None,
+    n_events: Annotated[Optional[int], typer.Option("--n-events", help="Cap number of seed events")] = None,
+    device: Annotated[Optional[str], typer.Option("--device", "-d", help="cpu | cuda | mps (auto)")] = None,
+    out: Annotated[Optional[Path], typer.Option("--out", "-o", help="Output steps parquet")] = None,
     seed: Annotated[
         Optional[int],
         typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
@@ -1679,8 +1543,7 @@ def rollout(
 
     if "sec_phys" not in ckpt.get("normalizer", {}):
         typer.echo(
-            "error: checkpoint has no normalizer.sec_phys — retrain with the "
-            "current code",
+            "error: checkpoint has no normalizer.sec_phys — retrain with the current code",
             err=True,
         )
         raise typer.Exit(1)
@@ -1730,10 +1593,7 @@ def rollout(
     typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})")
 
     oracle = GeometryOracle.load(geometry)
-    typer.echo(
-        f"loaded geometry oracle: {geometry} "
-        f"(escape_threshold={oracle.escape_threshold:.3f})"
-    )
+    typer.echo(f"loaded geometry oracle: {geometry} (escape_threshold={oracle.escape_threshold:.3f})")
 
     files = find_parquet_files(data)
     seeds = _seed_from_data(files, n_events)
@@ -1821,9 +1681,7 @@ def rollout(
             # stage2_model.particle_type.target="embedding"; omitted (not
             # written as null) otherwise, so giant.analysis can tell "not
             # applicable to this checkpoint" apart from "collector empty".
-            **(
-                {"type_embedding_l1_dist": l1_summary} if l1_summary is not None else {}
-            ),
+            **({"type_embedding_l1_dist": l1_summary} if l1_summary is not None else {}),
             # Full architecture spec baked into the checkpoint — includes the
             # entire router sub-dict, not just a hand-picked subset, so any
             # model knob (router type/n_experts, noise_dim, vocab sizes, ...)
@@ -1855,9 +1713,7 @@ app.add_typer(analyze_app, name="analyze")
 def analyze_prep(
     rollout_yaml: Annotated[
         Path,
-        typer.Argument(
-            help="giant rollout YAML sidecar (names the rollout + reference files)"
-        ),
+        typer.Argument(help="giant rollout YAML sidecar (names the rollout + reference files)"),
     ],
     run_dir: Annotated[
         Optional[Path],
@@ -1872,9 +1728,7 @@ def analyze_prep(
     top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6,
     chunks: Annotated[
         int,
-        typer.Option(
-            "--chunks", help="Split each plot's data into this many event_id chunks"
-        ),
+        typer.Option("--chunks", help="Split each plot's data into this many event_id chunks"),
     ] = 1,
 ) -> None:
     """Read the rollout YAML → shared.json + run_meta.json in the run directory."""
@@ -1894,15 +1748,9 @@ def analyze_prep(
 
 @analyze_app.command("compute-one")
 def analyze_compute_one(
-    id: Annotated[
-        str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")
-    ],
-    run_dir: Annotated[
-        Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")
-    ],
-    chunk: Annotated[
-        int, typer.Option("--chunk", help="Chunk index (see `analyze prep --chunks`)")
-    ] = 0,
+    id: Annotated[str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")],
+    run_dir: Annotated[Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")],
+    chunk: Annotated[int, typer.Option("--chunk", help="Chunk index (see `analyze prep --chunks`)")] = 0,
 ) -> None:
     """Run one (plot, chunk)'s streaming reduction (this is what each condor job runs)."""
     from giant.analysis import compute_one
@@ -1913,12 +1761,8 @@ def analyze_compute_one(
 
 @analyze_app.command("merge-one")
 def analyze_merge_one(
-    id: Annotated[
-        str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")
-    ],
-    run_dir: Annotated[
-        Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")
-    ],
+    id: Annotated[str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")],
+    run_dir: Annotated[Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")],
 ) -> None:
     """Merge one plot's chunk partials into its final reduced JSON.
 
@@ -1942,14 +1786,10 @@ def analyze_list() -> None:
 
 @analyze_app.command("render")
 def analyze_render(
-    run_dir: Annotated[
-        Path, typer.Argument(help="Run directory from `analyze prep` (holds reduced/)")
-    ],
+    run_dir: Annotated[Path, typer.Argument(help="Run directory from `analyze prep` (holds reduced/)")],
     gallery: Annotated[
         bool,
-        typer.Option(
-            "--gallery/--no-gallery", help="Run `gallery generate` after rendering"
-        ),
+        typer.Option("--gallery/--no-gallery", help="Run `gallery generate` after rendering"),
     ] = False,
 ) -> None:
     """Render reduced artifacts to styled PDFs + gallery metadata (local; needs LaTeX)."""
@@ -1971,9 +1811,7 @@ def analyze_submit(
             help="Override the run directory (default: <cwd>/analysis_runs/analysis_<id>)",
         ),
     ] = None,
-    docker_image: Annotated[
-        str, typer.Option("--docker-image")
-    ] = "cverstege/alma9-gridjob",
+    docker_image: Annotated[str, typer.Option("--docker-image")] = "cverstege/alma9-gridjob",
     request_memory: Annotated[int, typer.Option("--request-memory", help="MB")] = 8192,
     remote: Annotated[
         bool,
@@ -1989,9 +1827,7 @@ def analyze_submit(
     n_energy_bins: Annotated[int, typer.Option("--energy-bins")] = 4,
     n_marginal_bins: Annotated[int, typer.Option("--bins")] = 50,
     top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6,
-    dry_run: Annotated[
-        bool, typer.Option("--dry-run", help="Write files but don't condor_submit")
-    ] = False,
+    dry_run: Annotated[bool, typer.Option("--dry-run", help="Write files but don't condor_submit")] = False,
 ) -> None:
     """prep + write the HTCondor submit description (one job per plot x chunk), then submit."""
     import subprocess
diff --git a/giant/config.py b/giant/config.py
index c779969..92c42d9 100644
--- a/giant/config.py
+++ b/giant/config.py
@@ -299,13 +299,7 @@ DEFAULT_CONFIG: dict = {
 
 def git_hash() -> str:
     try:
-        return (
-            subprocess.check_output(
-                ["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL
-            )
-            .decode()
-            .strip()
-        )
+        return subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL).decode().strip()
     except Exception:
         return "unknown"
 
@@ -357,12 +351,8 @@ def estimate_batch_size(
     calibration since there's no backward graph or optimizer state.
     """
     if device.type != "cuda":
-        raise ValueError(
-            f"--batch-size auto is only supported on cuda devices, got {device.type!r}"
-        )
-    device_index = (
-        device.index if device.index is not None else torch.cuda.current_device()
-    )
+        raise ValueError(f"--batch-size auto is only supported on cuda devices, got {device.type!r}")
+    device_index = device.index if device.index is not None else torch.cuda.current_device()
     free_bytes, _total_bytes = torch.cuda.mem_get_info(device_index)
     if training:
         ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = (
@@ -690,11 +680,7 @@ def validate_config(cfg: dict) -> None:
 
     for stage_name in ("stage1_model", "stage2_model"):
         router = _get_path(cfg, f"{stage_name}.router") or {}
-        if (
-            router.get("enabled")
-            and router.get("type") in ("pdg", "process")
-            and particle_type == "physical"
-        ):
+        if router.get("enabled") and router.get("type") in ("pdg", "process") and particle_type == "physical":
             raise ValueError(
                 f"{stage_name}.router.type = {router['type']!r} builds its "
                 "own training-vocab-scoped embedding, incompatible with "
@@ -704,9 +690,7 @@ def validate_config(cfg: dict) -> None:
                 "conditioning.particle.type"
             )
 
-    if _get_path(cfg, "stage2_model.router.tie_to_stage1") and not _get_path(
-        cfg, "stage1_model.active"
-    ):
+    if _get_path(cfg, "stage2_model.router.tie_to_stage1") and not _get_path(cfg, "stage1_model.active"):
         raise ValueError(
             "stage2_model.router.tie_to_stage1 = true requires "
             "stage1_model.active = true (there is no stage-1 router to tie to)"
@@ -737,10 +721,7 @@ def validate_config(cfg: dict) -> None:
     if _get_path(cfg, "stage2_model.decoder") == "autoregressive":
         history = _get_path(cfg, "stage2_model.autoregressive.history")
         if history not in ("markov", "attention"):
-            raise ValueError(
-                f"stage2_model.autoregressive.history = {history!r} — must "
-                "be 'markov' or 'attention'"
-            )
+            raise ValueError(f"stage2_model.autoregressive.history = {history!r} — must be 'markov' or 'attention'")
         teacher_forcing = _get_path(cfg, "stage2_model.autoregressive.teacher_forcing")
         if teacher_forcing not in ("always", "scheduled", "never"):
             raise ValueError(
diff --git a/giant/data/dataset.py b/giant/data/dataset.py
index 13c78fc..6e2e282 100644
--- a/giant/data/dataset.py
+++ b/giant/data/dataset.py
@@ -123,9 +123,7 @@ class StreamingStepsDataset(IterableDataset):
         buf_n = 0
 
         for path in files:
-            for chunk in iter_file_chunks(
-                path, offset=self._offsets[path], k_max=self.k_max
-            ):
+            for chunk in iter_file_chunks(path, offset=self._offsets[path], k_max=self.k_max):
                 mask = sorted_membership(chunk["event_id"], self._events_arr)
                 if not mask.any():
                     continue
diff --git a/giant/data/loader.py b/giant/data/loader.py
index 43a1a44..94acfba 100644
--- a/giant/data/loader.py
+++ b/giant/data/loader.py
@@ -115,9 +115,7 @@ def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndar
     return out
 
 
-def _df_to_dict(
-    df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX
-) -> dict[str, np.ndarray]:
+def _df_to_dict(df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
     has_sec_lists = "sec_E_list" in df.columns
 
     d: dict[str, np.ndarray] = {
@@ -136,9 +134,7 @@ def _df_to_dict(
         # / ProcessRouter). Guarded like has_sec_lists: older parquet
         # conversions predating this column still load fine.
         "process": (
-            df["process"].to_numpy(dtype=object)
-            if "process" in df.columns
-            else np.full(len(df), "", dtype=object)
+            df["process"].to_numpy(dtype=object) if "process" in df.columns else np.full(len(df), "", dtype=object)
         ),
         "step_length": df["step_length"].to_numpy(dtype=np.float32),
         "post_E": df["post_E"].to_numpy(dtype=np.float32),
@@ -151,16 +147,12 @@ def _df_to_dict(
     if has_sec_lists:
         d["sec_E_list"] = _pad_list_col(df["sec_E_list"], k_max)
         d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], k_max)
-        d["sec_dir_list"] = _pad_dir_col(
-            df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max
-        )
+        d["sec_dir_list"] = _pad_dir_col(df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max)
 
     return d
 
 
-def load_steps(
-    path: str | Path, offset: int = 0, k_max: int = K_MAX
-) -> dict[str, np.ndarray]:
+def load_steps(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
     return _df_to_dict(pd.read_parquet(path), offset=offset, k_max=k_max)
 
 
@@ -170,9 +162,7 @@ def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
     return _offset_event_id(ids, offset)
 
 
-def iter_file_chunks(
-    path: str | Path, offset: int = 0, k_max: int = K_MAX
-) -> Iterator[dict[str, np.ndarray]]:
+def iter_file_chunks(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> Iterator[dict[str, np.ndarray]]:
     """Yield one parquet row-group at a time so a large file never fully loads.
 
     `k_max` sets the padded width of the sec_*_list columns (should match
@@ -214,15 +204,11 @@ def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]
     }
 
 
-def iter_cond_chunks(
-    path: str | Path, offset: int = 0
-) -> Iterator[dict[str, np.ndarray]]:
+def iter_cond_chunks(path: str | Path, offset: int = 0) -> Iterator[dict[str, np.ndarray]]:
     """Yield conditioning-only row-groups (no post-step columns read from disk)."""
     pf = pq.ParquetFile(path)
     for i in range(pf.num_row_groups):
-        yield _cond_df_to_dict(
-            pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset
-        )
+        yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset)
 
 
 def build_index_maps(
@@ -315,9 +301,7 @@ class TopNMap:
     other_members: dict
 
 
-def build_topn_map_from_files(
-    files: list[Path], column: str, n_classes: int, cast=str
-) -> TopNMap:
+def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, cast=str) -> TopNMap:
     """Scan `column` and build a frequency-capped value->index map, structurally
     identical to `build_process_map_from_files` (shares its ranking core via
     `_topn_plus_other_map`), generalized over the source column and key type.
diff --git a/giant/data/setup_cache.py b/giant/data/setup_cache.py
index a0a3981..b1a613e 100644
--- a/giant/data/setup_cache.py
+++ b/giant/data/setup_cache.py
@@ -108,10 +108,7 @@ def normalizer_key(
     # affect which cond_cont columns are computed for real vs. zero-filled
     # (giant.data.transforms._physical_cond_columns), so both must be part of
     # the key or two mixed-axis runs could collide on the same cache entry.
-    return (
-        f"valfrac={val_fraction:.6g}_seed={seed}_pcond={particle_conditioning}"
-        f"_mcond={material_conditioning}"
-    )
+    return f"valfrac={val_fraction:.6g}_seed={seed}_pcond={particle_conditioning}_mcond={material_conditioning}"
 
 
 # Top-N-map axes: "pdg" keys match pdg_map's int
@@ -126,10 +123,7 @@ def topn_key(axis: str, n_classes: int) -> str:
     sidecar stays reusable across runs with different emb_dim (see the
     dict[int, dict] precedent `proc_maps` sets, keyed by n_experts)."""
     if axis not in _TOPN_AXIS_CASTS:
-        raise ValueError(
-            f"unknown top-N map axis {axis!r}, expected one of "
-            f"{sorted(_TOPN_AXIS_CASTS)}"
-        )
+        raise ValueError(f"unknown top-N map axis {axis!r}, expected one of {sorted(_TOPN_AXIS_CASTS)}")
     return f"{axis}:{n_classes}"
 
 
@@ -165,9 +159,7 @@ class NormalizerEntry:
             "tgt_norm": self.tgt_norm.to_dict(),
             "sec_phys_norm": self.sec_phys_norm.to_dict(),
             "n_train_steps": self.n_train_steps,
-            "energy_quantiles": np.asarray(
-                self.energy_quantiles, dtype=np.float32
-            ).tolist(),
+            "energy_quantiles": np.asarray(self.energy_quantiles, dtype=np.float32).tolist(),
         }
 
     @classmethod
@@ -234,13 +226,8 @@ class SetupCache:
                 np.array(d["event_index"]["counts"], dtype=np.int64),
             )
         proc_maps = {int(k): v for k, v in d.get("proc_maps", {}).items()}
-        normalizers = {
-            k: NormalizerEntry.from_json(v) for k, v in d.get("normalizers", {}).items()
-        }
-        topn_maps = {
-            k: topnmap_from_json(v, axis=k.split(":", 1)[0])
-            for k, v in d.get("topn_maps", {}).items()
-        }
+        normalizers = {k: NormalizerEntry.from_json(v) for k, v in d.get("normalizers", {}).items()}
+        topn_maps = {k: topnmap_from_json(v, axis=k.split(":", 1)[0]) for k, v in d.get("topn_maps", {}).items()}
         return cls(
             fingerprint=d["fingerprint"],
             git_hash=d.get("git_hash", "unknown"),
@@ -263,18 +250,14 @@ class SetupCache:
             fingerprint=other.fingerprint,
             git_hash=other.git_hash,
             vocab=other.vocab if other.vocab is not None else self.vocab,
-            event_index=(
-                other.event_index if other.event_index is not None else self.event_index
-            ),
+            event_index=(other.event_index if other.event_index is not None else self.event_index),
             proc_maps={**self.proc_maps, **other.proc_maps},
             normalizers={**self.normalizers, **other.normalizers},
             topn_maps={**self.topn_maps, **other.topn_maps},
         )
 
 
-def load(
-    data: str | Path, files: list[Path], echo=lambda *a, **k: None
-) -> SetupCache | None:
+def load(data: str | Path, files: list[Path], echo=lambda *a, **k: None) -> SetupCache | None:
     """Load and validate the sidecar for `data`; `None` on any miss (never raises).
 
     A missing file, corrupt JSON, format-version mismatch, dimension-constant
@@ -298,9 +281,7 @@ def load(
             echo("setup cache: format version changed — ignoring stale cache")
             return None
         if raw.get("dims") != _DIMS:
-            echo(
-                "setup cache: model dimension constants changed — ignoring stale cache"
-            )
+            echo("setup cache: model dimension constants changed — ignoring stale cache")
             return None
         fp = fingerprint_files(files)
         if raw.get("fingerprint") != fp:
@@ -343,9 +324,7 @@ def save(
         with open(lock_path, "a") as lock_file:
             fcntl.flock(lock_file, fcntl.LOCK_EX)
             try:
-                base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(
-                    files
-                )
+                base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(files)
                 merged = base.merge(sections)
                 payload = json.dumps(merged.to_json(), separators=(",", ":"))
                 tmp.write_text(payload)
@@ -353,9 +332,7 @@ def save(
             finally:
                 fcntl.flock(lock_file, fcntl.LOCK_UN)
     except OSError as exc:
-        echo(
-            f"setup cache: could not write {path} ({exc}) — continuing without caching"
-        )
+        echo(f"setup cache: could not write {path} ({exc}) — continuing without caching")
         try:
             tmp.unlink(missing_ok=True)
         except OSError:
@@ -366,16 +343,12 @@ def compute_event_index_from_files(files: list[Path]) -> tuple[np.ndarray, np.nd
     """Unique event ids + per-event row (step) counts, across all `files`."""
     if not files:
         return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)
-    all_ids = np.concatenate(
-        [load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)]
-    )
+    all_ids = np.concatenate([load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)])
     unique_ids, counts = np.unique(all_ids, return_counts=True)
     return unique_ids, counts
 
 
-def n_train_steps_for_split(
-    unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray
-) -> int:
+def n_train_steps_for_split(unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray) -> int:
     """Row (step) count summed over whichever `unique_ids` fall in `train_events_arr`.
 
     `train_events_arr` must be ascending and duplicate-free (as produced by
diff --git a/giant/data/transforms.py b/giant/data/transforms.py
index bfe7234..4dfc672 100644
--- a/giant/data/transforms.py
+++ b/giant/data/transforms.py
@@ -84,9 +84,7 @@ def energy_simplex_encode(
     return z.astype(np.float32)
 
 
-def energy_simplex_decode(
-    z: np.ndarray, pre_E: np.ndarray
-) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+def energy_simplex_decode(z: np.ndarray, pre_E: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
     """Inverse of `energy_simplex_encode`: ALR coords + pre_E → physical energies.
 
     A softmax over `[z_edep, z_sec, 0]` recovers the three simplex fractions, so
@@ -118,9 +116,7 @@ def _rodrigues_axis(pre_dir: np.ndarray) -> np.ndarray:
     arbitrary second operand) on every row; profiling on a 114M-row file
     showed `np.cross` as the single hottest call inside this rotation.
     """
-    axis = np.stack(
-        [pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1
-    )
+    axis = np.stack([pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1)
     axis_norm = np.linalg.norm(axis, axis=1, keepdims=True)  # (N,1)
     # axis_norm ~ 0 happens at BOTH poles: pre_dir ~ +ẑ (forward) and
     # pre_dir ~ -ẑ (near-exact backscatter) — ‖pre_dir × ẑ‖ = sin(angle to
@@ -199,9 +195,7 @@ def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarra
     kxv = _cross_with_z_axis(axis, post_dir)  # (N,3)
     kdv = (axis * post_dir).sum(axis=1, keepdims=True)  # (N,1)
 
-    return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
-        np.float32
-    )
+    return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(np.float32)
 
 
 class Normalizer:
@@ -341,9 +335,7 @@ def sorted_membership(values: np.ndarray, sorted_arr: np.ndarray) -> np.ndarray:
     return sorted_arr[idx] == values
 
 
-def _vectorized_map_lookup(
-    values: np.ndarray, mapping: dict, strict: bool = True, default: int = 0
-) -> np.ndarray:
+def _vectorized_map_lookup(values: np.ndarray, mapping: dict, strict: bool = True, default: int = 0) -> np.ndarray:
     """Vectorized equivalent of `np.array([mapping[v] for v in values], dtype=np.int64)`.
 
     Replaces a per-element Python dict lookup with one `searchsorted` call.
@@ -387,9 +379,7 @@ def travel_direction(pre_pos: np.ndarray, post_pos: np.ndarray) -> np.ndarray:
     disp = post_pos - pre_pos
     norm = np.linalg.norm(disp, axis=1, keepdims=True)
     safe_norm = np.where(norm < 1e-7, 1.0, norm)
-    return np.where(norm < 1e-7, np.array([[0.0, 0.0, 1.0]]), disp / safe_norm).astype(
-        np.float32
-    )
+    return np.where(norm < 1e-7, np.array([[0.0, 0.0, 1.0]]), disp / safe_norm).astype(np.float32)
 
 
 def reconstruct_post_pos(
@@ -408,9 +398,7 @@ def reconstruct_post_pos(
     return (pre_pos + step_length.reshape(-1, 1) * travel_dir_world).astype(np.float32)
 
 
-def inv_local_frame_rotation(
-    pre_dir: np.ndarray, post_dir_local: np.ndarray
-) -> np.ndarray:
+def inv_local_frame_rotation(pre_dir: np.ndarray, post_dir_local: np.ndarray) -> np.ndarray:
     """Inverse of local_frame_rotation: rotate from local frame back to world frame.
 
     Applies R^T (same axis, negative angle) to post_dir_local.
@@ -424,9 +412,7 @@ def inv_local_frame_rotation(
     kdv = (axis * post_dir_local).sum(axis=1, keepdims=True)
 
     # Negative angle: sin_t → -sin_t
-    return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
-        np.float32
-    )
+    return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(np.float32)
 
 
 _STICK_LOGIT_CLIP = 10.0  # logit value used for the last valid secondary slot
@@ -491,14 +477,10 @@ def encode_secondaries(
                 remaining_raw = e_sec - cumsum[:, i - 1]
             shortfall_flagged |= sec_valid[:, i] & (remaining_raw < -_SHORTFALL_TOL)
             remaining = np.maximum(remaining_raw, _EPS)
-            f = np.clip(
-                sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS
-            )
+            f = np.clip(sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS)
             logit = np.log(f / (1.0 - f)).astype(np.float32)
             # Last valid slot: give it the full remaining budget
-            is_last = sec_valid[:, i] & ~(
-                sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool)
-            )
+            is_last = sec_valid[:, i] & ~(sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool))
             logit = np.where(is_last, _STICK_LOGIT_CLIP, logit)
             logit = np.where(
                 sec_valid[:, i],
@@ -525,9 +507,7 @@ def encode_secondaries(
             # Only rotate valid slots; leave padded slots as (0,0,1) or whatever.
             valid_mask = sec_valid[:, i]
             if valid_mask.any():
-                dir_local[valid_mask, i] = local_frame_rotation(
-                    pre_dir[valid_mask], sec_dir_list[valid_mask, i]
-                )
+                dir_local[valid_mask, i] = local_frame_rotation(pre_dir[valid_mask], sec_dir_list[valid_mask, i])
 
     if sec_pdg_list is not None:
         from giant.particles import particle_phys_array
@@ -553,9 +533,7 @@ def encode_secondaries(
     return sec_cont.astype(np.float32)
 
 
-def encode_secondary_type_idx(
-    sec_pdg_list: np.ndarray, sec_valid: np.ndarray, class_map: dict
-) -> np.ndarray:
+def encode_secondary_type_idx(sec_pdg_list: np.ndarray, sec_valid: np.ndarray, class_map: dict) -> np.ndarray:
     """Per-secondary-slot class index into `class_map` — (N, K_MAX) int64.
 
     `class_map` is either a top-N-plus-other map's `class_map`
@@ -583,9 +561,7 @@ def encode_secondary_type_idx(
     # isn't guaranteed to be a key — an arbitrary present one always is).
     dummy = next(iter(class_map))
     safe_pdg = np.where(sec_valid, sec_pdg_list, dummy)
-    idx = _vectorized_map_lookup(safe_pdg.reshape(-1), class_map, strict=True).reshape(
-        N, K
-    )
+    idx = _vectorized_map_lookup(safe_pdg.reshape(-1), class_map, strict=True).reshape(N, K)
     return np.where(sec_valid, idx, 0).astype(np.int64)
 
 
@@ -659,9 +635,7 @@ def decode_secondary_cont(
     for i in range(K):
         valid = sec_valid[:, i]
         if valid.any():
-            sec_dir_world[valid, i] = inv_local_frame_rotation(
-                pre_dir[valid], dir_local[valid, i]
-            )
+            sec_dir_world[valid, i] = inv_local_frame_rotation(pre_dir[valid], dir_local[valid, i])
 
     return sec_E, sec_dir_world, sec_valid
 
@@ -704,9 +678,7 @@ def decode_secondaries(
         sec_cont = sec_cont.copy()
         sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2)
 
-    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(
-        sec_cont, n_sec, e_sec, pre_dir
-    )
+    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(sec_cont, n_sec, e_sec, pre_dir)
 
     log_mass = sec_cont[:, :, 4]  # (N, K)
     charge = sec_cont[:, :, 5]  # (N, K)
@@ -760,16 +732,12 @@ def _physical_cond_columns(
     elif particle_conditioning in ("embedding", "onehot"):
         particle_cols = np.zeros((n, PARTICLE_PHYS_DIM), dtype=np.float32)
     else:
-        raise ValueError(
-            f"unknown conditioning.particle.type {particle_conditioning!r}"
-        )
+        raise ValueError(f"unknown conditioning.particle.type {particle_conditioning!r}")
 
     if material_conditioning == "physical":
         from giant.materials import material_properties_array
 
-        z_eff, a_eff, density, x0, lambda_int = material_properties_array(
-            data["material"]
-        ).T
+        z_eff, a_eff, density, x0, lambda_int = material_properties_array(data["material"]).T
         material_cols = np.column_stack(
             [
                 z_eff,
@@ -782,9 +750,7 @@ def _physical_cond_columns(
     elif material_conditioning in ("embedding", "onehot"):
         material_cols = np.zeros((n, MATERIAL_PHYS_DIM), dtype=np.float32)
     else:
-        raise ValueError(
-            f"unknown conditioning.material.type {material_conditioning!r}"
-        )
+        raise ValueError(f"unknown conditioning.material.type {material_conditioning!r}")
 
     return np.column_stack([particle_cols, material_cols]).astype(np.float32)
 
@@ -848,9 +814,7 @@ def build_cond_features(
     cond_cat = np.column_stack(cat_cols)
 
     if cond_normalizer is not None:
-        cond_cont = _cond_normalizer_transform(
-            cond_cont, cond_normalizer, particle_conditioning, material_conditioning
-        )
+        cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, particle_conditioning, material_conditioning)
 
     return cond_cont, cond_cat
 
@@ -972,13 +936,9 @@ def build_features(
     """
 
     post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])
-    travel_dir_local = local_frame_rotation(
-        data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"])
-    )
+    travel_dir_local = local_frame_rotation(data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"]))
 
-    energy_z = energy_simplex_encode(
-        data["edep"], data["e_sec"], data["post_E"], data["pre_E"]
-    )  # (N, 2)
+    energy_z = energy_simplex_encode(data["edep"], data["e_sec"], data["post_E"], data["pre_E"])  # (N, 2)
 
     target_s1 = np.column_stack(
         [
@@ -1014,9 +974,7 @@ def build_features(
         cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
     cond_cat = np.column_stack(cat_cols)  # (N, 2/3/4)
 
-    n_sec_raw = data["n_sec"].astype(
-        np.int64
-    )  # (N,) unclamped, for the valid-slot mask
+    n_sec_raw = data["n_sec"].astype(np.int64)  # (N,) unclamped, for the valid-slot mask
 
     # Secondary continuous targets
     sec_E_list = data.get("sec_E_list")
diff --git a/giant/geometry.py b/giant/geometry.py
index ffe27da..3d137e8 100644
--- a/giant/geometry.py
+++ b/giant/geometry.py
@@ -31,10 +31,7 @@ import numpy as np
 import pandas as pd
 import pyarrow.parquet as pq
 
-_INSTALL_HINT = (
-    "the geometry oracle needs scikit-learn — install it with "
-    "`uv sync --extra cpu --extra geometry`"
-)
+_INSTALL_HINT = "the geometry oracle needs scikit-learn — install it with `uv sync --extra cpu --extra geometry`"
 
 
 def _require_sklearn():
@@ -63,9 +60,7 @@ class _SlabLookup:
     layer_ids: np.ndarray  # (n_segments,) int64, layer_id of each segment
     radius_max: float  # largest transverse radius seen in training data
 
-    def query(
-        self, pos: np.ndarray, margin: float
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    def query(self, pos: np.ndarray, margin: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
         other = [i for i in range(3) if i != self.axis]
         z = pos[:, self.axis]
         radius = np.sqrt(pos[:, other[0]] ** 2 + pos[:, other[1]] ** 2)
@@ -75,11 +70,7 @@ class _SlabLookup:
         material = self.materials[idx]
         layer_id = self.layer_ids[idx]
 
-        escaped = (
-            (z < self.z_edges[0] - margin)
-            | (z > self.z_edges[-1] + margin)
-            | (radius > self.radius_max + margin)
-        )
+        escaped = (z < self.z_edges[0] - margin) | (z > self.z_edges[-1] + margin) | (radius > self.radius_max + margin)
         return material, layer_id, escaped
 
 
@@ -300,16 +291,12 @@ def _fit_slab_lookup(
     """
     other = [i for i in range(3) if i != axis]
     z = pos[:, axis].astype(np.float64)
-    radius = np.sqrt(
-        pos[:, other[0]].astype(np.float64) ** 2
-        + pos[:, other[1]].astype(np.float64) ** 2
-    )
+    radius = np.sqrt(pos[:, other[0]].astype(np.float64) ** 2 + pos[:, other[1]].astype(np.float64) ** 2)
 
     z_min, z_max = float(z.min()), float(z.max())
     if z_min == z_max:
         raise ValueError(
-            "all points share the same depth-axis coordinate — pick a "
-            "different `depth_axis` or use method='knn'/'svm'"
+            "all points share the same depth-axis coordinate — pick a different `depth_axis` or use method='knn'/'svm'"
         )
     edges = np.linspace(z_min, z_max, n_bins + 1)
     bin_idx = np.clip(np.searchsorted(edges, z, side="right") - 1, 0, n_bins - 1)
@@ -344,12 +331,7 @@ def _fit_slab_lookup(
     bin_layer = bin_layer[fill_from]
 
     # Run-length-encode consecutive bins sharing a label into segments.
-    changed = (
-        np.flatnonzero(
-            (bin_material[1:] != bin_material[:-1]) | (bin_layer[1:] != bin_layer[:-1])
-        )
-        + 1
-    )
+    changed = np.flatnonzero((bin_material[1:] != bin_material[:-1]) | (bin_layer[1:] != bin_layer[:-1])) + 1
     seg_starts = np.concatenate([[0], changed])
     z_edges = np.concatenate([edges[seg_starts], edges[-1:]])
     materials = bin_material[seg_starts]
@@ -460,11 +442,7 @@ def build_geometry_oracle(
     # Escape threshold from the reference point spacing. Sample a subset for the
     # median 2-NN distance (the 1st neighbour of a training point is itself).
     nn = NearestNeighbors(n_neighbors=2).fit(X)
-    probe = (
-        X
-        if len(X) <= 20_000
-        else X[np.random.default_rng(seed).choice(len(X), 20_000, replace=False)]
-    )
+    probe = X if len(X) <= 20_000 else X[np.random.default_rng(seed).choice(len(X), 20_000, replace=False)]
     d2, _ = nn.kneighbors(probe, n_neighbors=2)
     median_nn = float(np.median(d2[:, 1]))
     escape_threshold = escape_factor * median_nn
diff --git a/giant/materials.py b/giant/materials.py
index a15c8d2..1aed5de 100644
--- a/giant/materials.py
+++ b/giant/materials.py
@@ -64,18 +64,10 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
     "G4_CESIUM_IODIDE": MaterialProperties(
         z_eff=54.0, a_eff=129.904539, density=4.51, x0=1.860288, lambda_int=39.305990
     ),
-    "G4_Pb": MaterialProperties(
-        z_eff=82.0, a_eff=207.216962, density=11.35, x0=0.561253, lambda_int=18.247950
-    ),
-    "G4_W": MaterialProperties(
-        z_eff=74.0, a_eff=183.841648, density=19.30, x0=0.350418, lambda_int=10.311580
-    ),
-    "G4_Cu": MaterialProperties(
-        z_eff=29.0, a_eff=63.545648, density=8.96, x0=1.435578, lambda_int=15.587940
-    ),
-    "G4_Fe": MaterialProperties(
-        z_eff=26.0, a_eff=55.845113, density=7.874, x0=1.757493, lambda_int=16.990300
-    ),
+    "G4_Pb": MaterialProperties(z_eff=82.0, a_eff=207.216962, density=11.35, x0=0.561253, lambda_int=18.247950),
+    "G4_W": MaterialProperties(z_eff=74.0, a_eff=183.841648, density=19.30, x0=0.350418, lambda_int=10.311580),
+    "G4_Cu": MaterialProperties(z_eff=29.0, a_eff=63.545648, density=8.96, x0=1.435578, lambda_int=15.587940),
+    "G4_Fe": MaterialProperties(z_eff=26.0, a_eff=55.845113, density=7.874, x0=1.757493, lambda_int=16.990300),
     "G4_BRASS": MaterialProperties(
         z_eff=30.939130,
         a_eff=68.500857,
@@ -83,9 +75,7 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
         x0=1.367465,
         lambda_int=16.947420,
     ),
-    "G4_POLYSTYRENE": MaterialProperties(
-        z_eff=3.5, a_eff=6.509339, density=1.06, x0=41.312510, lambda_int=68.749880
-    ),
+    "G4_POLYSTYRENE": MaterialProperties(z_eff=3.5, a_eff=6.509339, density=1.06, x0=41.312510, lambda_int=68.749880),
     "G4_PLASTIC_SC_VINYLTOLUENE": MaterialProperties(
         z_eff=3.368421,
         a_eff=6.219791,
@@ -108,20 +98,15 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
         x0=30392.070000,
         lambda_int=71009.500000,
     ),
-    "G4_lAr": MaterialProperties(
-        z_eff=18.0, a_eff=39.947692, density=1.396, x0=14.003440, lambda_int=85.706400
-    ),
+    "G4_lAr": MaterialProperties(z_eff=18.0, a_eff=39.947692, density=1.396, x0=14.003440, lambda_int=85.706400),
 }
 
 
-def get_material_properties(
-    name: str, table: dict[str, MaterialProperties] | None = None
-) -> MaterialProperties:
+def get_material_properties(name: str, table: dict[str, MaterialProperties] | None = None) -> MaterialProperties:
     t = MATERIAL_PROPERTIES if table is None else table
     if name not in t:
         raise UnknownMaterialError(
-            f"material {name!r} is not in giant.materials.MATERIAL_PROPERTIES "
-            f"-- add it (known: {sorted(t)})"
+            f"material {name!r} is not in giant.materials.MATERIAL_PROPERTIES -- add it (known: {sorted(t)})"
         )
     props = t[name]
     if any(v is None for v in props):
@@ -134,9 +119,7 @@ def get_material_properties(
     return props
 
 
-def material_properties_array(
-    names: np.ndarray, table: dict[str, MaterialProperties] | None = None
-) -> np.ndarray:
+def material_properties_array(names: np.ndarray, table: dict[str, MaterialProperties] | None = None) -> np.ndarray:
     """(N,) str material names -> (N, 5) float32 [z_eff, a_eff, density, x0, lambda_int]."""
     out = np.array(
         [get_material_properties(str(m), table) for m in np.asarray(names)],
diff --git a/giant/model/network.py b/giant/model/network.py
index 8aeadfb..bed790c 100644
--- a/giant/model/network.py
+++ b/giant/model/network.py
@@ -30,11 +30,7 @@ class SinusoidalEmbedding(nn.Module):
         super().__init__()
         assert dim % 2 == 0, "dim must be even"
         half = dim // 2
-        freqs = torch.exp(
-            -math.log(10000)
-            * torch.arange(half, dtype=torch.float32)
-            / max(half - 1, 1)
-        )
+        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32) / max(half - 1, 1))
         self.register_buffer("freqs", freqs)
 
     def forward(self, t: torch.Tensor) -> torch.Tensor:
@@ -43,9 +39,7 @@ class SinusoidalEmbedding(nn.Module):
         return torch.cat([args.sin(), args.cos()], dim=-1)  # (B, dim)
 
 
-def cat_col_layout(
-    particle_type: str, material_type: str
-) -> tuple[int | None, int | None]:
+def cat_col_layout(particle_type: str, material_type: str) -> tuple[int | None, int | None]:
     """`cond_cat` column indices for each axis's top-N-onehot index, or
     `None` if that axis isn't `"onehot"`.
 
@@ -119,18 +113,14 @@ class ConditionEncoder(nn.Module):
         super().__init__()
         self.particle_cfg = dict(particle_cfg)
         self.material_cfg = dict(material_cfg)
-        self._particle_topn_col, self._material_topn_col = cat_col_layout(
-            particle_cfg["type"], material_cfg["type"]
-        )
+        self._particle_topn_col, self._material_topn_col = cat_col_layout(particle_cfg["type"], material_cfg["type"])
 
         p_type = particle_cfg["type"]
         p_emb_dim = particle_cfg["emb_dim"]
         if p_type == "embedding":
             self.pdg_emb = nn.Embedding(pdg_vocab, p_emb_dim)
         elif p_type == "physical":
-            self.particle_mlp = _make_axis_mlp(
-                PARTICLE_PHYS_DIM, p_emb_dim, particle_cfg.get("n_layers", 1)
-            )
+            self.particle_mlp = _make_axis_mlp(PARTICLE_PHYS_DIM, p_emb_dim, particle_cfg.get("n_layers", 1))
         elif p_type != "onehot":
             raise ValueError(f"unknown conditioning.particle.type {p_type!r}")
 
@@ -139,9 +129,7 @@ class ConditionEncoder(nn.Module):
         if m_type == "embedding":
             self.mat_emb = nn.Embedding(mat_vocab, m_emb_dim)
         elif m_type == "physical":
-            self.material_mlp = _make_axis_mlp(
-                MATERIAL_PHYS_DIM, m_emb_dim, material_cfg.get("n_layers", 1)
-            )
+            self.material_mlp = _make_axis_mlp(MATERIAL_PHYS_DIM, m_emb_dim, material_cfg.get("n_layers", 1))
         elif m_type != "onehot":
             raise ValueError(f"unknown conditioning.material.type {m_type!r}")
 
@@ -157,9 +145,7 @@ class ConditionEncoder(nn.Module):
         if p_type == "embedding":
             return self.pdg_emb(cond_cat[:, 0])
         if p_type == "physical":
-            particle_phys = cond_cont[
-                :, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM
-            ]
+            particle_phys = cond_cont[:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM]
             return self.particle_mlp(particle_phys)
         assert self._particle_topn_col is not None
         return F.one_hot(
@@ -245,9 +231,7 @@ class Router(nn.Module):
         """(B, n_experts) soft weights, rows summing to 1."""
         raise NotImplementedError
 
-    def combine_weights(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> torch.Tensor:
+    def combine_weights(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         """(B, n_experts) train-time expert-combination weights.
 
         Default (`gumbel=False`): identical to `gate()`. Opt-in
@@ -265,16 +249,12 @@ class Router(nn.Module):
         """(B,) hard expert index, used for eval-time grouped dispatch."""
         return self.gate(cond_cont, cond_cat).argmax(dim=-1)
 
-    def balance_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> torch.Tensor:
+    def balance_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         """Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017)."""
         importance = self.gate(cond_cont, cond_cat).sum(dim=0)  # (n_experts,)
         return (importance.std() / (importance.mean() + 1e-8)) ** 2
 
-    def classify_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
-    ) -> torch.Tensor:
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
         """Optional supervised auxiliary loss shaping the router's own belief.
 
         Default: none (a scalar 0). Routers gating on an unobservable
@@ -282,16 +262,12 @@ class Router(nn.Module):
         """
         return torch.zeros((), device=cond_cont.device)
 
-    def entropy_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> torch.Tensor:
+    def entropy_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         """Optional auxiliary loss rewarding sharper (lower-entropy) routing."""
         norm_entropy, _ = self.gate_stats(cond_cont, cond_cat)
         return norm_entropy
 
-    def gate_stats(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    def gate_stats(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         """Diagnostics: `(norm_entropy, importance)` — see v0.2 docstring for
         the full explanation, unchanged in v0.3.0."""
         gate = self.gate(cond_cont, cond_cat)  # (B, n_experts)
@@ -321,9 +297,7 @@ def build_router(name: str, n_experts: int, **kwargs) -> Router:
     in one config without special-casing.
     """
     if name not in ROUTER_REGISTRY:
-        raise ValueError(
-            f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}"
-        )
+        raise ValueError(f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}")
     cls = ROUTER_REGISTRY[name]
     accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
     filtered = {k: v for k, v in kwargs.items() if k in accepted}
@@ -374,8 +348,7 @@ class EnergyRouter(Router):
         if learn_width or learn_temperature:
             if not (width_min_ratio < 1.0 < width_max_ratio):
                 raise ValueError(
-                    f"width_min_ratio ({width_min_ratio}) and width_max_ratio "
-                    f"({width_max_ratio}) must bracket 1.0"
+                    f"width_min_ratio ({width_min_ratio}) and width_max_ratio ({width_max_ratio}) must bracket 1.0"
                 )
             self._width_lo = width_min_ratio * temperature
             self._width_hi = width_max_ratio * temperature
@@ -388,10 +361,7 @@ class EnergyRouter(Router):
             centers = torch.linspace(-2.0, 2.0, n_experts)
         else:
             if len(centers_init) != n_experts:
-                raise ValueError(
-                    f"centers_init has {len(centers_init)} values, "
-                    f"expected n_experts={n_experts}"
-                )
+                raise ValueError(f"centers_init has {len(centers_init)} values, expected n_experts={n_experts}")
             centers = torch.tensor(list(centers_init), dtype=torch.float32)
         if learn_centers:
             self.centers = nn.Parameter(centers)
@@ -436,9 +406,7 @@ class PdgRouter(Router):
 
     def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         e = self.pdg_emb(cond_cat[:, 0])  # (B, emb_dim)
-        d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(
-            -1
-        )  # (B, n_experts)
+        d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(-1)  # (B, n_experts)
         return torch.softmax(-d2 / self.temperature, dim=-1)
 
 
@@ -477,9 +445,7 @@ class ProcessRouter(Router):
     def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1)
 
-    def classify_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
-    ) -> torch.Tensor:
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
         return F.cross_entropy(self.logits(cond_cont, cond_cat), labels)
 
 
@@ -500,14 +466,10 @@ class ComposedRouter(Router):
         joint = self.routers[0].gate(cond_cont, cond_cat)  # (B, n_0)
         for router in self.routers[1:]:
             g = router.gate(cond_cont, cond_cat)  # (B, n_i)
-            joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(
-                1
-            )  # (B, prod so far)
+            joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(1)  # (B, prod so far)
         return joint
 
-    def classify_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
-    ) -> torch.Tensor:
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
         total = torch.zeros((), device=cond_cont.device)
         for router in self.routers:
             total = total + router.classify_loss(cond_cont, cond_cat, labels)
@@ -560,9 +522,7 @@ def _parse_composed_axes(router_cfg: dict) -> list[dict]:
 _VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process")
 
 
-def _check_router_conditioning_compat(
-    router_types: list[str], particle_conditioning: str
-) -> None:
+def _check_router_conditioning_compat(router_types: list[str], particle_conditioning: str) -> None:
     """Reject a router axis that reintroduces a training-vocab PDG lookup
     under `conditioning.particle.type = "physical"`.
 
@@ -596,16 +556,12 @@ def _build_router_from_cfg(
     shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
     if router_cfg["type"] == "composed":
         axes = _parse_composed_axes(router_cfg)
-        _check_router_conditioning_compat(
-            [a["type"] for a in axes], particle_conditioning
-        )
+        _check_router_conditioning_compat([a["type"] for a in axes], particle_conditioning)
         router = build_composed_router(axes, **shared_vocab)
         router.gumbel = bool(router_cfg.get("gumbel", False))
         return router
     _check_router_conditioning_compat([router_cfg["type"]], particle_conditioning)
-    router_kwargs = {
-        k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")
-    }
+    router_kwargs = {k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")}
     router_kwargs.setdefault("pdg_vocab", pdg_vocab)
     router_kwargs.setdefault("mat_vocab", mat_vocab)
     router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
@@ -637,9 +593,7 @@ class ExpertTrunk(nn.Module):
     ) -> None:
         super().__init__()
         self.input_proj = nn.Linear(in_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_blocks)]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_proj = nn.Linear(hidden_dim, out_dim)
 
     def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
@@ -710,12 +664,7 @@ class MonolithicTrunk(Trunk):
     ) -> None:
         super().__init__()
         self.input_proj = nn.Linear(in_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, cond_dim, dropout=dropout)
-                for _ in range(n_res_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_res_blocks)])
         self.out_proj = nn.Linear(hidden_dim, out_dim)
 
     def forward(
@@ -745,12 +694,7 @@ class RoutedTrunk(Trunk):
         super().__init__()
         self.router = router
         self.experts = nn.ModuleList(
-            [
-                ExpertTrunk(
-                    in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout
-                )
-                for _ in range(router.n_experts)
-            ]
+            [ExpertTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) for _ in range(router.n_experts)]
         )
 
     def forward(
@@ -760,9 +704,7 @@ class RoutedTrunk(Trunk):
         cond_cont: torch.Tensor,
         cond_cat: torch.Tensor,
     ) -> torch.Tensor:
-        return _route_forward(
-            self.experts, self.router, x, cond, cond_cont, cond_cat, self.training
-        )
+        return _route_forward(self.experts, self.router, x, cond, cond_cont, cond_cat, self.training)
 
 
 def build_trunk(
@@ -775,9 +717,7 @@ def build_trunk(
     dropout: float = 0.0,
 ) -> Trunk:
     if router is not None:
-        return RoutedTrunk(
-            router, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout
-        )
+        return RoutedTrunk(router, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
     return MonolithicTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
 
 
@@ -845,13 +785,9 @@ class _CausalAttnBlock(nn.Module):
     def __init__(self, dim: int, n_heads: int, dropout: float = 0.0) -> None:
         super().__init__()
         self.norm1 = nn.LayerNorm(dim)
-        self.attn = nn.MultiheadAttention(
-            dim, n_heads, dropout=dropout, batch_first=True
-        )
+        self.attn = nn.MultiheadAttention(dim, n_heads, dropout=dropout, batch_first=True)
         self.norm2 = nn.LayerNorm(dim)
-        self.mlp = nn.Sequential(
-            nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim)
-        )
+        self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))
 
     def forward(self, x: torch.Tensor, causal_mask: torch.Tensor) -> torch.Tensor:
         h = self.norm1(x)
@@ -860,9 +796,7 @@ class _CausalAttnBlock(nn.Module):
         x = x + self.mlp(self.norm2(x))
         return x
 
-    def step(
-        self, x_new: torch.Tensor, kv_cache: torch.Tensor | None
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    def step(self, x_new: torch.Tensor, kv_cache: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]:
         """`x_new`: `(B, 1, dim)`, this position's input. `kv_cache`: `None`
         (first position) or `(B, T, dim)` — `norm1(x)` of every earlier
         position at this same block. Returns `(out, new_kv_cache)`, `out`
@@ -900,15 +834,11 @@ class AttentionHistory(HistoryEncoder):
     `Stage2Autoregressive.history_step`.
     """
 
-    def __init__(
-        self, in_dim: int, out_dim: int, n_heads: int = 4, n_layers: int = 2
-    ) -> None:
+    def __init__(self, in_dim: int, out_dim: int, n_heads: int = 4, n_layers: int = 2) -> None:
         super().__init__()
         self.start = nn.Parameter(torch.zeros(in_dim))
         self.in_proj = nn.Linear(in_dim, out_dim)
-        self.blocks = nn.ModuleList(
-            [_CausalAttnBlock(out_dim, n_heads) for _ in range(n_layers)]
-        )
+        self.blocks = nn.ModuleList([_CausalAttnBlock(out_dim, n_heads) for _ in range(n_layers)])
 
     def _embed(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
         start = self.start.view(1, 1, -1).expand_as(feat)
@@ -959,9 +889,7 @@ def stage2_type_dim(particle_type_cfg: dict, emb_dim: int) -> int:
     return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
 
 
-def stage2_trunk_sec_dim(
-    particle_type_cfg: dict, generator: str, k_max: int, emb_dim: int
-) -> int:
+def stage2_trunk_sec_dim(particle_type_cfg: dict, generator: str, k_max: int, emb_dim: int) -> int:
     """`Stage2OneShot`'s trunk output width.
 
     `target = "physical"` is untouched from v0.2/today:
@@ -1020,17 +948,13 @@ class Stage1Model(nn.Module):
         self.cond_enc = (
             cond_enc
             if cond_enc is not None
-            else ConditionEncoder(
-                pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
-            )
+            else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
         )
         has_time = generator in ("flow", "ddpm")
         self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
         merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
         in_dim = noise_dim if generator == "wgan" else x_dim
-        self.trunk = build_trunk(
-            router, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout
-        )
+        self.trunk = build_trunk(router, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout)
         self.n_sec_head = None
         if n_sec_head_k_max is not None:
             self.n_sec_head = nn.Sequential(
@@ -1047,16 +971,10 @@ class Stage1Model(nn.Module):
         t: torch.Tensor | None = None,
     ) -> torch.Tensor:
         c_emb = self.cond_enc(cond_cont, cond_cat)
-        cond = (
-            torch.cat([self.time_emb(t), c_emb], dim=-1)
-            if self.time_emb is not None
-            else c_emb
-        )
+        cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb
         return self.trunk(x_t, cond, cond_cont, cond_cat)
 
-    def predict_n_sec(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> torch.Tensor:
+    def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         """Return n_sec logits (B, K_MAX+1) from conditioning alone. Only
         valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0
         configs predict n_sec from Stage2OneShot instead."""
@@ -1127,9 +1045,7 @@ class Stage2OneShot(nn.Module):
         self.cond_enc = (
             cond_enc
             if cond_enc is not None
-            else ConditionEncoder(
-                pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
-            )
+            else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
         )
         self.context_adapter = ContextAdapter(x_dim, context_dim)
         self.fuse = nn.Sequential(
@@ -1140,9 +1056,7 @@ class Stage2OneShot(nn.Module):
         self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
         merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
         in_dim = noise_dim if generator == "wgan" else sec_dim
-        self.trunk = build_trunk(
-            router, in_dim, sec_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout
-        )
+        self.trunk = build_trunk(router, in_dim, sec_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout)
         self.n_sec_head = None
         if build_n_sec_head:
             self.n_sec_head = nn.Sequential(
@@ -1162,9 +1076,7 @@ class Stage2OneShot(nn.Module):
             self._type_k_max = k_max
             self._type_emb_dim = emb_dim
 
-    def _cond_embed(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor
-    ) -> torch.Tensor:
+    def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
         base = self.cond_enc(cond_cont, cond_cat)
         ctx = self.context_adapter(stage1_out)
         return self.fuse(torch.cat([base, ctx], dim=-1))
@@ -1178,11 +1090,7 @@ class Stage2OneShot(nn.Module):
         t: torch.Tensor | None = None,
     ) -> torch.Tensor:
         c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
-        cond = (
-            torch.cat([self.time_emb(t), c_emb], dim=-1)
-            if self.time_emb is not None
-            else c_emb
-        )
+        cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb
         return self.trunk(x_t, cond, cond_cont, cond_cat)
 
     def predict_n_sec(
@@ -1274,10 +1182,7 @@ class Stage2Autoregressive(nn.Module):
     ) -> None:
         super().__init__()
         if history not in ("markov", "attention"):
-            raise ValueError(
-                f"stage2_model.autoregressive.history={history!r} — must be "
-                "'markov' or 'attention'"
-            )
+            raise ValueError(f"stage2_model.autoregressive.history={history!r} — must be 'markov' or 'attention'")
         self.history_kind = history
         self.generator_kind = generator
         self.noise_dim = noise_dim
@@ -1289,9 +1194,7 @@ class Stage2Autoregressive(nn.Module):
         self.cond_enc = (
             cond_enc
             if cond_enc is not None
-            else ConditionEncoder(
-                pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
-            )
+            else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
         )
         self.context_adapter = ContextAdapter(x_dim, context_dim)
         self.base_fuse = nn.Sequential(
@@ -1305,15 +1208,11 @@ class Stage2Autoregressive(nn.Module):
         history_dim = cond_out_dim
         hist_in_dim = CONT_SLOT_DIM + self.type_dim
         self.history_encoder: HistoryEncoder = (
-            AttentionHistory(
-                hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers
-            )
+            AttentionHistory(hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers)
             if history == "attention"
             else MarkovHistory(hist_in_dim, history_dim)
         )
-        token_fuse_in = (
-            cond_out_dim + context_dim + history_dim + 2
-        )  # +2: remaining_frac, slot_idx
+        token_fuse_in = cond_out_dim + context_dim + history_dim + 2  # +2: remaining_frac, slot_idx
         self.token_fuse = nn.Sequential(
             nn.Linear(token_fuse_in, cond_out_dim),
             nn.SiLU(),
@@ -1350,9 +1249,7 @@ class Stage2Autoregressive(nn.Module):
                 nn.Linear(hidden_dim // 2, self.type_dim),
             )
 
-    def _base_cond(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor
-    ) -> torch.Tensor:
+    def _base_cond(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
         base = self.cond_enc(cond_cont, cond_cat)
         ctx = self.context_adapter(stage1_out)
         return self.base_fuse(torch.cat([base, ctx], dim=-1))
@@ -1392,9 +1289,7 @@ class Stage2Autoregressive(nn.Module):
             return self.history_encoder.init_cache()
         return None
 
-    def history_step(
-        self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache
-    ) -> tuple[torch.Tensor, object]:
+    def history_step(self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache) -> tuple[torch.Tensor, object]:
         """One inference slot's worth of history encoding: advances `cache`
         (from `init_history_cache`, or a previous `history_step` call) by
         `token_feat`/`has_prev` (`(B, 1, ...)` — the just-emitted previous
@@ -1445,9 +1340,7 @@ class Stage2Autoregressive(nn.Module):
         out = self.trunk(x_flat, cond_flat, cond_cont_flat, cond_cat_flat)
         return out.view(B, K, -1)
 
-    def predict_n_sec(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor
-    ) -> torch.Tensor:
+    def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
         if self.n_sec_head is None:
             raise RuntimeError(
                 "this Stage2Autoregressive has no n_sec_head — it belongs to "
@@ -1514,9 +1407,7 @@ class CriticModel(nn.Module):
         if stage not in ("stage1", "stage2"):
             raise ValueError(f"stage must be 'stage1' or 'stage2', got {stage!r}")
         self.stage = stage
-        self.cond_enc = ConditionEncoder(
-            pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
-        )
+        self.cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
         if stage == "stage2":
             self.context_adapter = ContextAdapter(context_in_dim, context_dim)
             self.fuse = nn.Sequential(
@@ -1524,12 +1415,7 @@ class CriticModel(nn.Module):
                 nn.SiLU(),
             )
         self.input_proj = nn.Linear(in_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
-                for _ in range(n_res_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_res_blocks)])
         self.out_norm = nn.LayerNorm(hidden_dim)
         self.out_proj = nn.Linear(hidden_dim, 1)
 
@@ -1641,9 +1527,7 @@ def _migrate_legacy_model_config(model_config: dict) -> dict:
     }
 
 
-def migrate_legacy_state_dict(
-    old_stage1_sd: dict, old_stage2_sd: dict
-) -> tuple[dict, dict]:
+def migrate_legacy_state_dict(old_stage1_sd: dict, old_stage2_sd: dict) -> tuple[dict, dict]:
     """Remap a v0.2 checkpoint's (`DenoisingMLP`-or-`WGANGenerator`,
     `SecondaryDecoder`-or-`WGANSecondaryGenerator`) state dicts onto the new
     `(Stage1Model, Stage2OneShot)` module structure produced by
@@ -1702,11 +1586,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
     common representation. `false` (default) keeps v0.2 behaviour:
     independent instances with identical config but independent weights.
     """
-    cfg = (
-        model_config
-        if "stage1_model" in model_config
-        else _migrate_legacy_model_config(model_config)
-    )
+    cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config)
     pdg_vocab = cfg["pdg_vocab"]
     mat_vocab = cfg["mat_vocab"]
     conditioning = cfg["conditioning"]
@@ -1718,9 +1598,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
     cond_out_dim = conditioning.get("out_dim", 128)
     shared_cond_enc: ConditionEncoder | None = None
     if conditioning.get("share_stages"):
-        shared_cond_enc = ConditionEncoder(
-            pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
-        )
+        shared_cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
 
     result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}
 
@@ -1728,15 +1606,11 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
     if s1cfg.get("active", True):
         router_cfg = s1cfg.get("router") or {}
         if router_cfg.get("enabled"):
-            stage1_router = _build_router_from_cfg(
-                router_cfg, pdg_vocab, mat_vocab, particle_conditioning
-            )
+            stage1_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
         generator = s1cfg.get("generator", "flow")
         gen_sub = s1cfg.get(generator, {}) or {}
         legacy_owner = (s2cfg.get("n_sec") or {}).get("legacy_owner")
-        n_sec_head_k_max = (
-            s2cfg.get("k_max", K_MAX) if legacy_owner == "stage1" else None
-        )
+        n_sec_head_k_max = s2cfg.get("k_max", K_MAX) if legacy_owner == "stage1" else None
         result["stage1"] = Stage1Model(
             pdg_vocab=pdg_vocab,
             mat_vocab=mat_vocab,
@@ -1762,9 +1636,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
             if router_cfg.get("tie_to_stage1") and stage1_router is not None:
                 stage2_router = stage1_router
             else:
-                stage2_router = _build_router_from_cfg(
-                    router_cfg, pdg_vocab, mat_vocab, particle_conditioning
-                )
+                stage2_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
         generator = s2cfg.get("generator", "wgan")
         gen_sub = s2cfg.get(generator, {}) or {}
         legacy_owner = (s2cfg.get("n_sec") or {}).get("legacy_owner")
@@ -1796,9 +1668,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 cond_enc=shared_cond_enc,
             )
         else:
-            sec_dim = stage2_trunk_sec_dim(
-                particle_type_cfg, generator, k_max, particle_cfg["emb_dim"]
-            )
+            sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator, k_max, particle_cfg["emb_dim"])
             result["stage2"] = Stage2OneShot(
                 pdg_vocab=pdg_vocab,
                 mat_vocab=mat_vocab,
@@ -1828,11 +1698,7 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
     "wgan"` training. Training-only — never persisted for inference the way
     `build_models`'s pair is. `None` for a stage that's inactive or not
     WGAN."""
-    cfg = (
-        model_config
-        if "stage1_model" in model_config
-        else _migrate_legacy_model_config(model_config)
-    )
+    cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config)
     pdg_vocab = cfg["pdg_vocab"]
     mat_vocab = cfg["mat_vocab"]
     conditioning = cfg["conditioning"]
@@ -1861,9 +1727,7 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
     if s2cfg.get("active", True) and s2cfg.get("generator") == "wgan":
         k_max = s2cfg.get("k_max", K_MAX)
         particle_type_cfg = s2cfg.get("particle_type") or {"target": "physical"}
-        in_dim = stage2_trunk_sec_dim(
-            particle_type_cfg, "wgan", k_max, particle_cfg["emb_dim"]
-        )
+        in_dim = stage2_trunk_sec_dim(particle_type_cfg, "wgan", k_max, particle_cfg["emb_dim"])
         result["stage2"] = CriticModel(
             pdg_vocab=pdg_vocab,
             mat_vocab=mat_vocab,
diff --git a/giant/model/schedule.py b/giant/model/schedule.py
index b8b3891..c63e018 100644
--- a/giant/model/schedule.py
+++ b/giant/model/schedule.py
@@ -11,9 +11,7 @@ class CosineSchedule:
         steps = np.arange(T + 1, dtype=np.float64)
         f = np.cos(((steps / T + s) / (1.0 + s)) * np.pi / 2.0) ** 2
         alpha_bars = (f / f[0]).astype(np.float32)
-        betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(
-            np.float32
-        )
+        betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(np.float32)
 
         self.betas = torch.from_numpy(betas)
         self.alphas = torch.from_numpy((1.0 - betas))
diff --git a/giant/particles.py b/giant/particles.py
index c6be4e9..d8fcdc1 100644
--- a/giant/particles.py
+++ b/giant/particles.py
@@ -58,9 +58,7 @@ def particle_mass_charge(pdg: int) -> tuple[float, float]:
         if _pdgid.is_nucleus(pdg):
             z, a = _pdgid.Z(pdg), _pdgid.A(pdg)
             if z is None or a is None:
-                raise ValueError(
-                    f"PDG {pdg}: is_nucleus but Z/A decode failed"
-                ) from None
+                raise ValueError(f"PDG {pdg}: is_nucleus but Z/A decode failed") from None
             return float(a) * _AMU_MEV, float(z)
         raise ValueError(
             f"PDG code {pdg} could not be resolved via the `particle` package "
@@ -109,9 +107,7 @@ def nearest_known_pdg(mass: np.ndarray, charge: np.ndarray, candidates) -> np.nd
     if len(resolved) == 0:
         raise ValueError("nearest_known_pdg: no resolvable candidates")
     codes = np.array([r[0] for r in resolved], dtype=np.int64)
-    table_log_mass = np.log(
-        np.array([r[1] for r in resolved], dtype=np.float64) + _LOG_EPS
-    )
+    table_log_mass = np.log(np.array([r[1] for r in resolved], dtype=np.float64) + _LOG_EPS)
     table_charge = np.array([r[2] for r in resolved], dtype=np.float64)
 
     mass = np.asarray(mass, dtype=np.float64)
@@ -180,10 +176,7 @@ def decode_topn_class(
     n_other = int(other_mask.sum())
     if n_other:
         if not topn_map.other_members:
-            raise ValueError(
-                "decode_topn_class: 'other' class predicted but "
-                "topn_map.other_members is empty"
-            )
+            raise ValueError("decode_topn_class: 'other' class predicted but topn_map.other_members is empty")
         members = np.array(list(topn_map.other_members.keys()), dtype=np.int64)
         counts = np.array(list(topn_map.other_members.values()), dtype=np.float64)
         if other_policy == "drop":
diff --git a/giant/pipeline.py b/giant/pipeline.py
index 843140d..b6367ee 100644
--- a/giant/pipeline.py
+++ b/giant/pipeline.py
@@ -73,21 +73,14 @@ def _seed_energy_router(
     if not active:
         return
     if energy_quantiles.size == 0:
-        echo(
-            "  warning: no energy samples collected — EnergyRouter falls back to "
-            "default centers"
-        )
+        echo("  warning: no energy samples collected — EnergyRouter falls back to default centers")
         return
     assert cond_norm.mean is not None and cond_norm.std is not None
     levels = np.linspace(0.0, 1.0, router_cfg["n_experts"])
     raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels)
-    centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[
-        energy_idx
-    ]
+    centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[energy_idx]
     router_cfg["centers_init"] = centers_init.astype(np.float32).tolist()
-    echo(
-        f"  seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}"
-    )
+    echo(f"  seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}")
 
 
 def run_setup_stage(
@@ -142,23 +135,14 @@ def run_setup_stage(
         if cache is not None:
             cache.event_index = (unique_ids, counts)
 
-    train_events, val_events = make_event_split(
-        unique_ids, val_fraction=val_fraction, seed=seed
-    )
+    train_events, val_events = make_event_split(unique_ids, val_fraction=val_fraction, seed=seed)
     events_arr = np.array(sorted(train_events))
     n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr)
-    echo(
-        f"  {int(counts.sum()):,} steps  |  "
-        f"{len(train_events)} train events  |  "
-        f"{len(val_events)} val events"
-    )
+    echo(f"  {int(counts.sum()):,} steps  |  {len(train_events)} train events  |  {len(val_events)} val events")
 
     if cache is not None and cache.vocab is not None:
         pdg_map, mat_map = cache.vocab
-        echo(
-            f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, "
-            f"{len(mat_map)} materials)"
-        )
+        echo(f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, {len(mat_map)} materials)")
     else:
         echo("building vocabulary maps …")
         pdg_map, mat_map = build_index_maps_from_files(files)
@@ -174,11 +158,7 @@ def run_setup_stage(
     # need that generality.
     proc_map: dict[str, int] | None = None
     process_router_cfg = next(
-        (
-            r
-            for r in (stage1_router, stage2_router)
-            if r.get("enabled") and r.get("type") == "process"
-        ),
+        (r for r in (stage1_router, stage2_router) if r.get("enabled") and r.get("type") == "process"),
         None,
     )
     if process_router_cfg is not None:
@@ -186,10 +166,7 @@ def run_setup_stage(
         cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None
         if cached_proc_map is not None:
             proc_map = cached_proc_map
-            echo(
-                f"process vocabulary: cache hit ({len(proc_map)} labels, "
-                f"{n_experts} experts)"
-            )
+            echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {n_experts} experts)")
         else:
             echo("building process vocabulary …")
             proc_map = build_process_map_from_files(files, n_experts=n_experts)
@@ -214,16 +191,11 @@ def run_setup_stage(
         cached = cache.topn_maps.get(cache_key) if cache is not None else None
         if cached is not None:
             pdg_topn_map = cached
-            echo(
-                f"pdg top-N map: cache hit ({len(pdg_topn_map.class_map)} codes, "
-                f"{n_classes} classes)"
-            )
+            echo(f"pdg top-N map: cache hit ({len(pdg_topn_map.class_map)} codes, {n_classes} classes)")
         else:
             echo("building pdg top-N map …")
             pdg_topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
-            echo(
-                f"  {len(pdg_topn_map.class_map)} pdg codes mapped to {n_classes} classes"
-            )
+            echo(f"  {len(pdg_topn_map.class_map)} pdg codes mapped to {n_classes} classes")
             if cache is not None:
                 cache.topn_maps[cache_key] = pdg_topn_map
 
@@ -234,29 +206,17 @@ def run_setup_stage(
         cached = cache.topn_maps.get(cache_key) if cache is not None else None
         if cached is not None:
             mat_topn_map = cached
-            echo(
-                f"material top-N map: cache hit ({len(mat_topn_map.class_map)} "
-                f"materials, {n_classes} classes)"
-            )
+            echo(f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {n_classes} classes)")
         else:
             echo("building material top-N map …")
-            mat_topn_map = build_topn_map_from_files(
-                files, "material", n_classes=n_classes, cast=str
-            )
-            echo(
-                f"  {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes"
-            )
+            mat_topn_map = build_topn_map_from_files(files, "material", n_classes=n_classes, cast=str)
+            echo(f"  {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes")
             if cache is not None:
                 cache.topn_maps[cache_key] = mat_topn_map
 
-    energy_router_active = any(
-        r.get("enabled") and r.get("type") == "energy"
-        for r in (stage1_router, stage2_router)
-    )
+    energy_router_active = any(r.get("enabled") and r.get("type") == "energy" for r in (stage1_router, stage2_router))
     energy_idx = 3
-    norm_key = setup_cache.normalizer_key(
-        val_fraction, seed, particle_conditioning, material_conditioning
-    )
+    norm_key = setup_cache.normalizer_key(val_fraction, seed, particle_conditioning, material_conditioning)
     entry = cache.normalizers.get(norm_key) if cache is not None else None
 
     if entry is not None:
@@ -281,27 +241,23 @@ def run_setup_stage(
         # router against this same (val_fraction, seed, conditioning) key
         # never needs to rescan just to seed centers.
         collect_energy_sample = energy_router_active or cache is not None
-        energy_sampler = (
-            _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
-        )
+        energy_sampler = _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
         for i, path in enumerate(files):
             for chunk in iter_file_chunks(path, offset=event_id_offset(i), k_max=k_max):
                 mask = sorted_membership(chunk["event_id"], events_arr)
                 if not mask.any():
                     continue
                 chunk_tr = {k: v[mask] for k, v in chunk.items()}
-                cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _, _ = (
-                    build_features(
-                        chunk_tr,
-                        pdg_map,
-                        mat_map,
-                        proc_map=proc_map,
-                        require_secondaries=True,
-                        particle_conditioning=particle_conditioning,
-                        material_conditioning=material_conditioning,
-                        sec_phys_only=True,
-                        k_max=k_max,
-                    )
+                cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _, _ = build_features(
+                    chunk_tr,
+                    pdg_map,
+                    mat_map,
+                    proc_map=proc_map,
+                    require_secondaries=True,
+                    particle_conditioning=particle_conditioning,
+                    material_conditioning=material_conditioning,
+                    sec_phys_only=True,
+                    k_max=k_max,
                 )
                 cond_acc.update(cond_cont)
                 tgt_acc.update(target_s1)
@@ -423,9 +379,7 @@ def run_train_job(
     # The secondary type-index map depends on stage2_model.particle_type.target,
     # independently of conditioning's own onehot/embedding choice above
     # (physical stays untouched/None).
-    particle_type_target = (
-        cfg["stage2_model"].get("particle_type", {}).get("target", "physical")
-    )
+    particle_type_target = cfg["stage2_model"].get("particle_type", {}).get("target", "physical")
     if particle_type_target == "onehot":
         assert setup.pdg_topn_map is not None
         sec_type_class_map = setup.pdg_topn_map.class_map
diff --git a/giant/rollout.py b/giant/rollout.py
index bb6b9fd..087473d 100644
--- a/giant/rollout.py
+++ b/giant/rollout.py
@@ -120,9 +120,7 @@ def decode_secondary_identity(
     pdg_topn_map: "TopNMap | None",
     other_policy: str,
     rng: np.random.Generator | None,
-) -> tuple[
-    np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None
-]:
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None]:
     """Decode Stage 2's raw (sec_cont, sec_type) output into physical
     secondary attributes, branching on `sec_decoder.particle_type_cfg`:
 
@@ -150,14 +148,12 @@ def decode_secondary_identity(
         sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid = decode_secondaries(
             sec_full, n_sec_np, e_sec, pre_dir, sec_phys_normalizer=sec_phys_norm
         )
-        sec_pdg = nearest_known_pdg(
-            sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()
-        ).reshape(sec_mass.shape)
+        sec_pdg = nearest_known_pdg(sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()).reshape(
+            sec_mass.shape
+        )
         return sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, None
 
-    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(
-        sec_cont.cpu().numpy(), n_sec_np, e_sec, pre_dir
-    )
+    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(sec_cont.cpu().numpy(), n_sec_np, e_sec, pre_dir)
     sec_type_np = sec_type.cpu().numpy()
     l1_dist = None
 
@@ -182,12 +178,8 @@ def decode_secondary_identity(
         l1_dist = np.where(sec_valid, l1_dist, 0.0).astype(np.float32)
 
     sec_mass, sec_charge = particle_phys_array(sec_pdg.reshape(-1)).T
-    sec_mass = np.where(sec_valid, sec_mass.reshape(sec_pdg.shape), 0.0).astype(
-        np.float32
-    )
-    sec_charge = np.where(sec_valid, sec_charge.reshape(sec_pdg.shape), 0.0).astype(
-        np.float32
-    )
+    sec_mass = np.where(sec_valid, sec_mass.reshape(sec_pdg.shape), 0.0).astype(np.float32)
+    sec_charge = np.where(sec_valid, sec_charge.reshape(sec_pdg.shape), 0.0).astype(np.float32)
     sec_pdg = np.where(sec_valid, sec_pdg, 0).astype(np.int64)
     return sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, l1_dist
 
@@ -299,13 +291,9 @@ class _Recorder:
     RAM until the very end.
     """
 
-    def __init__(
-        self, sink: Callable[[dict[str, np.ndarray]], None] | None = None
-    ) -> None:
+    def __init__(self, sink: Callable[[dict[str, np.ndarray]], None] | None = None) -> None:
         self._sink = sink
-        self._cols: dict[str, list] | None = (
-            None if sink is not None else {k: [] for k in _RECORD_KEYS}
-        )
+        self._cols: dict[str, list] | None = None if sink is not None else {k: [] for k in _RECORD_KEYS}
         self.n_rows = 0
         self.termination_reason_counts: Counter[str] = Counter()
 
@@ -313,10 +301,7 @@ class _Recorder:
         n = len(cols["event_id"])
         if n == 0:
             return
-        row = {
-            k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n)
-            for k in _RECORD_KEYS
-        }
+        row = {k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n) for k in _RECORD_KEYS}
         self.n_rows += n
         reasons = row["termination_reason"]
         nonempty = reasons[reasons != ""]
@@ -333,8 +318,7 @@ class _Recorder:
 
     def to_dict(self) -> dict[str, np.ndarray]:
         assert self._cols is not None, (
-            "to_dict() is unavailable when streaming to a sink — use "
-            "n_rows/termination_reason_counts instead"
+            "to_dict() is unavailable when streaming to a sink — use n_rows/termination_reason_counts instead"
         )
         out = {}
         for k, chunks in self._cols.items():
@@ -619,33 +603,19 @@ def _step_chunk(
     # --- Pre-step termination gates (in priority order; each track picks one) ---
     stop = np.zeros(n, dtype=bool)
     escaped_sel = escaped & ~stop
-    rec.add(
-        **_terminal_rows(
-            tr, escaped_sel, TERM_ESCAPED, edep=np.zeros(int(escaped_sel.sum()))
-        )
-    )
+    rec.add(**_terminal_rows(tr, escaped_sel, TERM_ESCAPED, edep=np.zeros(int(escaped_sel.sum()))))
     stop |= escaped_sel
 
     unknown_sel = ~known_pdg & ~stop
-    rec.add(
-        **_terminal_rows(
-            tr, unknown_sel, TERM_UNKNOWN_PDG, edep=tr["pre_E"][unknown_sel]
-        )
-    )
+    rec.add(**_terminal_rows(tr, unknown_sel, TERM_UNKNOWN_PDG, edep=tr["pre_E"][unknown_sel]))
     stop |= unknown_sel
 
     cutoff_sel = (tr["pre_E"] < energy_cutoff) & ~stop
-    rec.add(
-        **_terminal_rows(
-            tr, cutoff_sel, TERM_ENERGY_CUTOFF, edep=tr["pre_E"][cutoff_sel]
-        )
-    )
+    rec.add(**_terminal_rows(tr, cutoff_sel, TERM_ENERGY_CUTOFF, edep=tr["pre_E"][cutoff_sel]))
     stop |= cutoff_sel
 
     maxstep_sel = (tr["step_in_track"] >= max_steps) & ~stop
-    rec.add(
-        **_terminal_rows(tr, maxstep_sel, TERM_MAX_STEPS, edep=tr["pre_E"][maxstep_sel])
-    )
+    rec.add(**_terminal_rows(tr, maxstep_sel, TERM_MAX_STEPS, edep=tr["pre_E"][maxstep_sel]))
     stop |= maxstep_sel
 
     active = ~stop
@@ -679,41 +649,27 @@ def _step_chunk(
         cond_norm,
         particle_conditioning=particle_conditioning,
         material_conditioning=material_conditioning,
-        pdg_topn_map=pdg_topn_map.class_map
-        if particle_conditioning == "onehot"
-        else None,
-        mat_topn_map=mat_topn_map.class_map
-        if material_conditioning == "onehot"
-        else None,
+        pdg_topn_map=pdg_topn_map.class_map if particle_conditioning == "onehot" else None,
+        mat_topn_map=mat_topn_map.class_map if material_conditioning == "onehot" else None,
     )
     cc = torch.from_numpy(cond_cont).float().to(device)
     ck = torch.from_numpy(cond_cat).long().to(device)
 
-    stage1_norm, n_sec_pred_stage1 = sample_stage1(
-        stage1_model, cc, ck, steps, stage1_ddpm_steps
-    )
+    stage1_norm, n_sec_pred_stage1 = sample_stage1(stage1_model, cc, ck, steps, stage1_ddpm_steps)
     raw = tgt_norm.inverse_transform(stage1_norm.cpu().numpy())
 
     step_length = inv_log_transform(raw[:, 0])
     edep, e_sec, post_E, _delta = energy_simplex_decode(raw[:, 1:3], tr["pre_E"])
 
     post_dir_local = raw[:, 3:6].copy()
-    post_dir_local /= np.clip(
-        np.linalg.norm(post_dir_local, axis=1, keepdims=True), 1e-8, None
-    )
+    post_dir_local /= np.clip(np.linalg.norm(post_dir_local, axis=1, keepdims=True), 1e-8, None)
     post_dir_world = inv_local_frame_rotation(tr["pre_dir"], post_dir_local)
 
     travel_dir_local = raw[:, 6:9].copy()
-    travel_dir_local /= np.clip(
-        np.linalg.norm(travel_dir_local, axis=1, keepdims=True), 1e-8, None
-    )
-    post_pos = reconstruct_post_pos(
-        tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local
-    )
+    travel_dir_local /= np.clip(np.linalg.norm(travel_dir_local, axis=1, keepdims=True), 1e-8, None)
+    post_pos = reconstruct_post_pos(tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local)
 
-    n_sec_pred = resolve_n_sec(
-        stage1_model, sec_decoder, cc, ck, stage1_norm, n_sec_pred_stage1
-    )
+    n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cc, ck, stage1_norm, n_sec_pred_stage1)
     n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64)
 
     # --- Secondaries ---
@@ -724,23 +680,19 @@ def _step_chunk(
     # decode_secondary_identity's docstring for how each
     # particle_type.target differs on whether PDG resolution is a real
     # identity decision or just a reporting label.
-    sec_cont, sec_type, _valid = sample_stage2(
-        sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps
-    )
-    sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, sec_type_l1_dist = (
-        decode_secondary_identity(
-            sec_decoder,
-            sec_cont,
-            sec_type,
-            n_sec_np,
-            e_sec,
-            tr["pre_dir"],
-            sec_phys_norm,
-            pdg_map,
-            pdg_topn_map,
-            other_policy,
-            rng,
-        )
+    sec_cont, sec_type, _valid = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
+    sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, sec_type_l1_dist = decode_secondary_identity(
+        sec_decoder,
+        sec_cont,
+        sec_type,
+        n_sec_np,
+        e_sec,
+        tr["pre_dir"],
+        sec_phys_norm,
+        pdg_map,
+        pdg_topn_map,
+        other_policy,
+        rng,
     )
     sec_valid = np.arange(sec_E.shape[1])[None, :] < n_sec_np[:, None]
     if l1_dist_collector is not None and sec_type_l1_dist is not None:
diff --git a/giant/sample.py b/giant/sample.py
index 3683384..d4486b2 100644
--- a/giant/sample.py
+++ b/giant/sample.py
@@ -67,9 +67,7 @@ def sample_ddpm(
         alpha = schedule.alphas[i]
         alpha_bar = schedule.alpha_bars[i]
         z = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
-        x = (1.0 / alpha.sqrt()) * (
-            x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred
-        ) + beta.sqrt() * z
+        x = (1.0 / alpha.sqrt()) * (x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred) + beta.sqrt() * z
     return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)
 
 
@@ -173,9 +171,7 @@ def _decode_stage2_flat(
     else:
         sec_cont = x.view(B, k_max, CONT_SLOT_DIM)
         sec_type = sec_decoder.predict_type(cond_cont, cond_cat, stage1_out)
-    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
-        1
-    )
+    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
     return sec_cont, sec_type, sec_valid
 
 
@@ -206,9 +202,7 @@ def sample_secondaries(
         v = sec_decoder(x, cond_cont, cond_cat, stage1_out, t=t)
         x = x + v * dt
 
-    return _decode_stage2_flat(
-        sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred
-    )
+    return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)
 
 
 @torch.no_grad()
@@ -226,9 +220,7 @@ def sample_secondaries_wgan(
     B = cond_cont.size(0)
     z = torch.randn(B, sec_decoder.noise_dim, device=cond_cont.device)
     x = sec_decoder(z, cond_cont, cond_cat, stage1_out)
-    return _decode_stage2_flat(
-        sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred
-    )
+    return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)
 
 
 @torch.no_grad()
@@ -306,12 +298,8 @@ def sample_secondaries_ar(
         has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
         history_feat = prev_repr.unsqueeze(1)  # (B, 1, CONT_SLOT_DIM + type_dim)
         remaining_frac = remaining.unsqueeze(1)  # (B, 1)
-        slot_idx = torch.full(
-            (B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32
-        )
-        hist, history_cache = sec_decoder.history_step(
-            history_feat, has_prev, history_cache
-        )
+        slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
+        hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)
 
         if generator == "wgan":
             z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
@@ -366,21 +354,15 @@ def sample_secondaries_ar(
         sec_type[:, k] = type_k
 
         if target == "onehot":
-            type_for_history = F.one_hot(
-                type_k.argmax(dim=-1), num_classes=type_dim
-            ).float()
+            type_for_history = F.one_hot(type_k.argmax(dim=-1), num_classes=type_dim).float()
         else:
             type_for_history = type_k
 
         stick_fraction = torch.sigmoid(cont_k[:, 0])
-        prev_repr = torch.cat(
-            [stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1
-        )
+        prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
         remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)
 
-    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
-        1
-    )
+    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
     return sec_cont, sec_type, sec_valid
 
 
@@ -426,16 +408,10 @@ def sample_stage2(
     case, so there's nothing to dispatch to here.
     """
     if isinstance(sec_decoder, Stage2Autoregressive):
-        return sample_secondaries_ar(
-            sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps
-        )
+        return sample_secondaries_ar(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
     if sec_decoder.generator_kind == "wgan":
-        return sample_secondaries_wgan(
-            sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
-        )
-    return sample_secondaries(
-        sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps
-    )
+        return sample_secondaries_wgan(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
+    return sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
 
 
 def resolve_n_sec(
diff --git a/giant/training/loop.py b/giant/training/loop.py
index b2cb192..ca43651 100644
--- a/giant/training/loop.py
+++ b/giant/training/loop.py
@@ -78,9 +78,7 @@ def _try_validate_marginals(trainer: StageTrainer, val_loader, device, **kwargs)
     return validate_marginals(model, val_loader, device=device, **kwargs)
 
 
-def _marginal_kl(
-    trainers: dict[str, StageTrainer], val_loader, device, **kwargs
-) -> float:
+def _marginal_kl(trainers: dict[str, StageTrainer], val_loader, device, **kwargs) -> float:
     """Mean marginal KL over the stage-1 sampling chain, or NaN when stage 1
     is inactive or `validate_marginals` declined to produce a result."""
     stage1 = trainers.get("stage1")
@@ -90,9 +88,7 @@ def _marginal_kl(
         stage1,
         val_loader,
         device,
-        sec_decoder=trainers["stage2"].sampling_model()
-        if "stage2" in trainers
-        else None,
+        sec_decoder=trainers["stage2"].sampling_model() if "stage2" in trainers else None,
         **kwargs,
     )
     if result is None:
@@ -137,9 +133,7 @@ def train(
 
     trainers = build_stage_trainers(cfg, models, critics, device, total_train_batches)
     if not trainers:
-        raise ValueError(
-            "no active stage — stage1_model.active and stage2_model.active are both false"
-        )
+        raise ValueError("no active stage — stage1_model.active and stage2_model.active are both false")
     has_adversarial = any(not tr.supports_val_loss for tr in trainers.values())
 
     checkpoint_extras = {
@@ -147,12 +141,8 @@ def train(
         "pdg_map": pdg_map,
         "mat_map": mat_map,
         "proc_map": proc_map,
-        "pdg_topn_map": topnmap_to_json(pdg_topn_map)
-        if pdg_topn_map is not None
-        else None,
-        "mat_topn_map": topnmap_to_json(mat_topn_map)
-        if mat_topn_map is not None
-        else None,
+        "pdg_topn_map": topnmap_to_json(pdg_topn_map) if pdg_topn_map is not None else None,
+        "mat_topn_map": topnmap_to_json(mat_topn_map) if mat_topn_map is not None else None,
         "model_config": model_config,
     }
 
@@ -166,10 +156,7 @@ def train(
         best_val_loss = ckpt.get("best_val_loss", float("inf"))
         global_step = ckpt.get("global_step", 0)
         if start_epoch > epochs:
-            print(
-                f"checkpoint already completed epoch {start_epoch - 1} "
-                f"(>= --epochs {epochs}) — nothing to train"
-            )
+            print(f"checkpoint already completed epoch {start_epoch - 1} (>= --epochs {epochs}) — nothing to train")
             return
 
     collector = MetricsCollector.create(
@@ -206,10 +193,7 @@ def train(
             for batch in bar:
                 B = batch[0].size(0)
                 collector.add_train_batch(
-                    {
-                        name: trainer.step(batch, device, global_step)
-                        for name, trainer in trainers.items()
-                    },
+                    {name: trainer.step(batch, device, global_step) for name, trainer in trainers.items()},
                     B,
                 )
                 bar.set_postfix_str(collector.postfix(), refresh=False)
@@ -220,9 +204,7 @@ def train(
             bar.close()
 
             if shutdown.requested:
-                ckpt = build_checkpoint(
-                    trainers, epoch - 1, global_step, best_val_loss, checkpoint_extras
-                )
+                ckpt = build_checkpoint(trainers, epoch - 1, global_step, best_val_loss, checkpoint_extras)
                 torch.save(ckpt, out_dir / "last.pt")
                 last_completed_epoch = epoch - 1
                 print(
@@ -244,15 +226,10 @@ def train(
                             break
                         B = batch[0].size(0)
                         collector.add_val_batch(
-                            {
-                                name: tr.val_loss(batch, device)
-                                for name, tr in scored.items()
-                            },
+                            {name: tr.val_loss(batch, device) for name, tr in scored.items()},
                             B,
                         )
-                        collector.observe_routers(
-                            batch[0].to(device), batch[1].to(device), B
-                        )
+                        collector.observe_routers(batch[0].to(device), batch[1].to(device), B)
 
             # An adversarial stage has no averageable validation loss, so it
             # needs the marginal-KL signal every epoch to pick a best
@@ -264,10 +241,7 @@ def train(
             elif validate_every > 0 and epoch % validate_every == 0:
                 stage1 = trainers.get("stage1")
                 ddpm_steps = 1000
-                if (
-                    isinstance(stage1, FlowDDPMStageTrainer)
-                    and stage1.ddpm_schedule is not None
-                ):
+                if isinstance(stage1, FlowDDPMStageTrainer) and stage1.ddpm_schedule is not None:
                     ddpm_steps = stage1.ddpm_schedule.T
                 marginal_kl = _marginal_kl(
                     trainers,
@@ -292,22 +266,16 @@ def train(
             collector.set("val/marginal_kl", marginal_kl)
             collector.set(
                 "gpu_mem_mb",
-                torch.cuda.max_memory_allocated(device) / (1024 * 1024)
-                if device.type == "cuda"
-                else 0.0,
-            )
-            collector.set(
-                "samples_per_sec", collector.train_samples / max(epoch_time, 1e-8)
+                torch.cuda.max_memory_allocated(device) / (1024 * 1024) if device.type == "cuda" else 0.0,
             )
+            collector.set("samples_per_sec", collector.train_samples / max(epoch_time, 1e-8))
             collector.set("is_best", int(is_best))
             collector.set("epoch_time_s", epoch_time)
 
             print(collector.summary_line(val_loss, epoch_time, is_best))
             collector.write_epoch(global_step)
 
-            ckpt = build_checkpoint(
-                trainers, epoch, global_step, best_val_loss, checkpoint_extras
-            )
+            ckpt = build_checkpoint(trainers, epoch, global_step, best_val_loss, checkpoint_extras)
             if is_best:
                 best_val_loss = val_loss
                 ckpt["best_val_loss"] = best_val_loss
diff --git a/giant/training/metrics.py b/giant/training/metrics.py
index 42975ee..d1fb0dd 100644
--- a/giant/training/metrics.py
+++ b/giant/training/metrics.py
@@ -116,11 +116,7 @@ class _RouterAccumulator:
 
     def add(self, entropy: torch.Tensor, importance: torch.Tensor, n: int) -> None:
         self.entropy += entropy.item() * n
-        self.importance = (
-            importance.clone()
-            if self.importance is None
-            else self.importance + importance
-        )
+        self.importance = importance.clone() if self.importance is None else self.importance + importance
         self.n += n
 
     def stats(self) -> dict[str, float]:
@@ -168,22 +164,18 @@ class MetricsCollector:
         self._train = {name: _Accumulator() for name in trainers}
         self._val = {name: _Accumulator() for name in trainers}
         self._routers = {
-            name: _RouterAccumulator(tr.router.n_experts)
-            for name, tr in trainers.items()
-            if tr.router is not None
+            name: _RouterAccumulator(tr.router.n_experts) for name, tr in trainers.items() if tr.router is not None
         }
         # Only "mean" specs need summing; "last" specs are read straight off
         # the accumulator's most recent stats dict. "grad_norm" is always
         # summed — it feeds the run-level `grad_norm` column whether or not
         # a trainer reports it per stage.
         self._train_keys = {
-            name: {spec.key for spec in tr.train_metrics if spec.reduce == "mean"}
-            | {"grad_norm"}
+            name: {spec.key for spec in tr.train_metrics if spec.reduce == "mean"} | {"grad_norm"}
             for name, tr in trainers.items()
         }
         self._val_keys = {
-            name: {spec.key for spec in tr.val_metrics if spec.reduce == "mean"}
-            for name, tr in trainers.items()
+            name: {spec.key for spec in tr.val_metrics if spec.reduce == "mean"} for name, tr in trainers.items()
         }
         self._run_values: dict[str, float] = {}
         self._epoch = 0
@@ -224,13 +216,9 @@ class MetricsCollector:
                 import wandb
             except ImportError as exc:
                 raise RuntimeError(
-                    "train.wandb = true (--wandb) requires the 'wandb' package — "
-                    "install it via `uv sync --extra wandb`"
+                    "train.wandb = true (--wandb) requires the 'wandb' package — install it via `uv sync --extra wandb`"
                 ) from exc
-            param_counts = {
-                name: sum(p.numel() for p in tr.model.parameters())
-                for name, tr in trainers.items()
-            }
+            param_counts = {name: sum(p.numel() for p in tr.model.parameters()) for name, tr in trainers.items()}
             param_counts["total"] = sum(param_counts.values())
             wandb_run = wandb.init(
                 project=wandb_project,
@@ -290,9 +278,7 @@ class MetricsCollector:
             self._batch_grad_norm += stage_stats.get("grad_norm", 0.0)
         if self._ema_seeded:
             self._ema_loss += _EMA_ALPHA * (self._batch_loss - self._ema_loss)
-            self._ema_grad_norm += _EMA_ALPHA * (
-                self._batch_grad_norm - self._ema_grad_norm
-            )
+            self._ema_grad_norm += _EMA_ALPHA * (self._batch_grad_norm - self._ema_grad_norm)
         else:
             self._ema_loss = self._batch_loss
             self._ema_grad_norm = self._batch_grad_norm
@@ -303,9 +289,7 @@ class MetricsCollector:
             self._val[name].add(stage_stats, self._val_keys[name], batch_size)
 
     @torch.no_grad()
-    def observe_routers(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, batch_size: int
-    ) -> None:
+    def observe_routers(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, batch_size: int) -> None:
         """Record gate diagnostics for every routed stage on this batch.
 
         Called from the validation pass only (as in v0.2/v0.3.0), so a stage
@@ -339,9 +323,7 @@ class MetricsCollector:
             payload[f"batch/{name}/lr"] = trainer.optimizer.param_groups[0]["lr"]
             if trainer.router is not None:
                 with torch.no_grad():
-                    entropy, _ = trainer.router.gate_stats(
-                        batch[0].to(device), batch[1].to(device)
-                    )
+                    entropy, _ = trainer.router.gate_stats(batch[0].to(device), batch[1].to(device))
                 payload[f"batch/{name}/router/entropy"] = entropy.item()
         self.wandb_run.log(payload, step=global_step)
 
@@ -367,10 +349,7 @@ class MetricsCollector:
         self._run_values[column] = value
 
     def summary_line(self, val_loss: float, epoch_time: float, is_best: bool) -> str:
-        bits = [
-            trainer.summary(self.train_means(name))
-            for name, trainer in self.trainers.items()
-        ]
+        bits = [trainer.summary(self.train_means(name)) for name, trainer in self.trainers.items()]
         marker = "  [best]" if is_best else ""
         return (
             f"epoch {self._epoch:{self.epoch_width}d}/{self.epochs}  "
diff --git a/giant/training/stage2_inputs.py b/giant/training/stage2_inputs.py
index b23ac5d..e2f8367 100644
--- a/giant/training/stage2_inputs.py
+++ b/giant/training/stage2_inputs.py
@@ -100,9 +100,9 @@ def _assemble_stage2_real(
     `Stage2Autoregressive`'s per-token target also uses; the two must stay in
     lockstep. See `_assemble_stage2_ar_target`'s docstring for the
     (target, generator) width rules."""
-    return _assemble_stage2_ar_target(
-        sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim
-    ).flatten(1)
+    return _assemble_stage2_ar_target(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim).flatten(
+        1
+    )
 
 
 def _stick_fraction(sec_cont: torch.Tensor) -> torch.Tensor:
@@ -137,18 +137,14 @@ def _ar_has_prev(k_max: int, device: torch.device) -> torch.Tensor:
     return (torch.arange(k_max, device=device) >= 1).unsqueeze(0)
 
 
-def _ar_meta(
-    k_max: int, batch: int, device: torch.device, fraction: torch.Tensor
-) -> dict[str, torch.Tensor]:
+def _ar_meta(k_max: int, batch: int, device: torch.device, fraction: torch.Tensor) -> dict[str, torch.Tensor]:
     """`has_prev`/`remaining_frac`/`slot_idx` — the three per-token AR
     conditioning tensors that don't depend on *which* history representation
     (ground truth vs. the scheduled-sampling mix) produced `fraction`.
     Shared by `_assemble_stage2_ar_inputs` and
     `_assemble_stage2_ar_inputs_scheduled`, which differ only in
     `history_feat`."""
-    slot_idx = (
-        torch.arange(k_max, device=device).float() / max(k_max - 1, 1)
-    ).unsqueeze(0)
+    slot_idx = (torch.arange(k_max, device=device).float() / max(k_max - 1, 1)).unsqueeze(0)
     return {
         "has_prev": _ar_has_prev(k_max, device).expand(batch, -1),
         "remaining_frac": _remaining_energy_fraction(fraction),
@@ -182,9 +178,7 @@ def _assemble_stage2_ar_inputs(
     return {"history_feat": history_feat, **_ar_meta(K, B, device, fraction)}
 
 
-def _stage2_tf_prob(
-    mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int
-) -> float:
+def _stage2_tf_prob(mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int) -> float:
     """P(condition slot k+1 on the TRUE token k rather than the model's own
     prediction), for the current epoch
     (`stage2_model.autoregressive.teacher_forcing`).
@@ -260,9 +254,7 @@ def _assemble_stage2_ar_inputs_scheduled(
     device = sec_cont.device
     B, K = sec_cont.shape[0], sec_cont.shape[1]
     if p_tf >= 1.0:
-        return _assemble_stage2_ar_inputs(
-            sec_cont, sec_type_idx, particle_type_cfg, cond_enc, emb_dim
-        )
+        return _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, particle_type_cfg, cond_enc, emb_dim)
 
     was_training = model.training
     sec_cont_pred, sec_type_pred, _ = sample_secondaries_ar(
@@ -273,9 +265,7 @@ def _assemble_stage2_ar_inputs_scheduled(
 
     fraction_gt = _stick_fraction(sec_cont)
     dir_gt = sec_cont[..., 1:CONT_SLOT_DIM]
-    type_repr_gt = _type_repr(
-        sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim
-    )
+    type_repr_gt = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim)
     fraction_pred, dir_pred, type_repr_pred = _history_repr_from_ar_sample(
         sec_cont_pred, sec_type_pred, particle_type_cfg
     )
@@ -323,8 +313,6 @@ def _relax_onehot_type_slice(
     cont, type_logits = x[..., :cont_dim], x[..., cont_dim:]
     if grad_probe is not None:
         cont.register_hook(lambda g: grad_probe.__setitem__("cont", g.norm().item()))
-        type_logits.register_hook(
-            lambda g: grad_probe.__setitem__("type", g.norm().item())
-        )
+        type_logits.register_hook(lambda g: grad_probe.__setitem__("type", g.norm().item()))
     type_soft = F.gumbel_softmax(type_logits, tau=tau, hard=True, dim=-1)
     return torch.cat([cont, type_soft], dim=-1).reshape(B, -1)
diff --git a/giant/training/trainers.py b/giant/training/trainers.py
index e1ac249..8473a87 100644
--- a/giant/training/trainers.py
+++ b/giant/training/trainers.py
@@ -41,9 +41,7 @@ from giant.training.stage2_inputs import (
 
 
 @torch.no_grad()
-def _update_ema(
-    ema_model: torch.nn.Module, model: torch.nn.Module, decay: float
-) -> None:
+def _update_ema(ema_model: torch.nn.Module, model: torch.nn.Module, decay: float) -> None:
     for ema_p, p in zip(ema_model.parameters(), model.parameters()):
         ema_p.mul_(decay).add_(p, alpha=1 - decay)
 
@@ -129,9 +127,7 @@ class StageSpec:
     type_gumbel_tau_end: float = 0.1
 
     @classmethod
-    def from_config(
-        cls, cfg: dict, name: str, is_stage2: bool, steps_per_epoch: int
-    ) -> "StageSpec":
+    def from_config(cls, cfg: dict, name: str, is_stage2: bool, steps_per_epoch: int) -> "StageSpec":
         t = cfg["train"]
         stage_cfg = cfg[f"{name}_model"]
         router_cfg = stage_cfg.get("router") or {}
@@ -144,8 +140,7 @@ class StageSpec:
             decoder=stage_cfg.get("decoder", "one_shot") if is_stage2 else "one_shot",
             lambda_weight=stage_cfg.get("lambda", 1.0),
             n_sec_lambda=cfg["stage2_model"].get("n_sec", {}).get("lambda", 0.1),
-            particle_type=cfg["stage2_model"].get("particle_type")
-            or {"target": "physical"},
+            particle_type=cfg["stage2_model"].get("particle_type") or {"target": "physical"},
             particle_type_emb_dim=cfg["conditioning"]["particle"]["emb_dim"],
             # train.* keys are all guaranteed by DEFAULT_CONFIG's deep-merge
             # (giant/config.py), so they read directly; the field defaults
@@ -244,9 +239,7 @@ class StageTrainer:
 
     # --- schedule -------------------------------------------------------
 
-    def _init_lr_schedule(
-        self, optimizer: optim.Optimizer, warmup_steps: int, total_steps: int
-    ) -> None:
+    def _init_lr_schedule(self, optimizer: optim.Optimizer, warmup_steps: int, total_steps: int) -> None:
         self._lr_lambda = _cosine_warmup_lambda(warmup_steps, total_steps)
         self.total_steps = total_steps
         self.lr_sched = optim.lr_scheduler.LambdaLR(optimizer, self._lr_lambda)
@@ -270,9 +263,7 @@ class StageTrainer:
         """This stage's fragment of the end-of-epoch console line."""
         raise NotImplementedError
 
-    def val_objective(
-        self, train_means: dict, val_means: dict, marginal_kl: float
-    ) -> float:
+    def val_objective(self, train_means: dict, val_means: dict, marginal_kl: float) -> float:
         """This stage's contribution to the best-checkpoint selection score."""
         raise NotImplementedError
 
@@ -359,9 +350,7 @@ class StageTrainer:
         return target.flatten(1) if flatten else target
 
     @staticmethod
-    def _sec_mask(
-        n_sec: torch.Tensor, k_max: int, device: torch.device
-    ) -> torch.Tensor:
+    def _sec_mask(n_sec: torch.Tensor, k_max: int, device: torch.device) -> torch.Tensor:
         """`(B, K_MAX)` bool prefix mask: slot k is valid iff `k < n_sec`."""
         return torch.arange(k_max, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
 
@@ -397,9 +386,7 @@ class StageTrainer:
         return l_nsec, nsec_acc
 
     @staticmethod
-    def _step_optimizer(
-        optimizer: optim.Optimizer, loss: torch.Tensor, params: list
-    ) -> float:
+    def _step_optimizer(optimizer: optim.Optimizer, loss: torch.Tensor, params: list) -> float:
         """`zero_grad -> backward -> clip_grad_norm_(1.0) -> step`, returning
         the pre-clip grad norm. The one place the grad-clip constant lives."""
         optimizer.zero_grad()
@@ -450,9 +437,7 @@ class StageTrainer:
 class FlowDDPMStageTrainer(StageTrainer):
     """flow or ddpm generator for a single stage."""
 
-    def __init__(
-        self, spec: StageSpec, model: torch.nn.Module, device: torch.device
-    ) -> None:
+    def __init__(self, spec: StageSpec, model: torch.nn.Module, device: torch.device) -> None:
         if spec.is_stage2 and spec.generator not in ("flow",):
             raise NotImplementedError(
                 f"stage2_model.generator={spec.generator!r} is accepted by the "
@@ -466,26 +451,16 @@ class FlowDDPMStageTrainer(StageTrainer):
         # NotImplementedError above): "physical" keeps it folded in
         # (PARTICLE_PHYS_DIM wide, unchanged from v0.2); "onehot"/"embedding"
         # pull it out into model.type_head instead (0 here).
-        self._flow_type_dim = (
-            None
-            if self.particle_type_cfg.get("target", "physical") == "physical"
-            else 0
-        )
+        self._flow_type_dim = None if self.particle_type_cfg.get("target", "physical") == "physical" else 0
 
         self.params = list(self.model.parameters())
-        self.optimizer = optim.AdamW(
-            self.params, lr=spec.lr, weight_decay=spec.weight_decay
-        )
+        self.optimizer = optim.AdamW(self.params, lr=spec.lr, weight_decay=spec.weight_decay)
         self._init_lr_schedule(
             self.optimizer,
             warmup_steps=spec.warmup_epochs * spec.steps_per_epoch,
             total_steps=max(spec.epochs * spec.steps_per_epoch, 1),
         )
-        self.ddpm_schedule = (
-            CosineSchedule(T=spec.ddpm_n_steps).to(device)
-            if spec.generator == "ddpm"
-            else None
-        )
+        self.ddpm_schedule = CosineSchedule(T=spec.ddpm_n_steps).to(device) if spec.generator == "ddpm" else None
 
         self.train_metrics = [
             train_metric(key)
@@ -515,9 +490,7 @@ class FlowDDPMStageTrainer(StageTrainer):
         ]
         self.stage_metrics = [stage_metric("lr")]
 
-    def _generator_loss(
-        self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=None
-    ):
+    def _generator_loss(self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=None):
         if not self.is_stage2:
             if self.generator == "flow":
                 return flow_matching_loss(self.model, x1_s1, cond_cont, cond_cat)
@@ -584,22 +557,16 @@ class FlowDDPMStageTrainer(StageTrainer):
         mask = sec_mask.float()
         denom = mask.sum().clamp(min=1)
         if self.particle_type_cfg.get("target") == "onehot":
-            ce = F.cross_entropy(
-                type_out.transpose(1, 2), sec_type_idx, reduction="none"
-            )
+            ce = F.cross_entropy(type_out.transpose(1, 2), sec_type_idx, reduction="none")
             l_type = (ce * mask).sum() / denom
-            type_acc = (
-                (type_out.argmax(-1) == sec_type_idx).float() * mask
-            ).sum() / denom
+            type_acc = ((type_out.argmax(-1) == sec_type_idx).float() * mask).sum() / denom
         else:  # "embedding"
             target_vec = self.model.cond_enc.pdg_emb(sec_type_idx).detach()
             se = ((type_out - target_vec) ** 2).mean(-1)
             l_type = (se * mask).sum() / denom
         return l_type, type_acc
 
-    def _compute(
-        self, batch: tuple, device: torch.device, epoch: int | None = None
-    ) -> dict:
+    def _compute(self, batch: tuple, device: torch.device, epoch: int | None = None) -> dict:
         """`epoch=None` (the `val_loss` path) always uses full teacher
         forcing (`p_tf=1.0`) regardless of `spec.teacher_forcing` — validation
         should stay a stable, non-stochastic ground-truth comparison; only
@@ -619,23 +586,13 @@ class FlowDDPMStageTrainer(StageTrainer):
         x1_s2 = None
         ar_inputs = None
         if self.is_stage2 and self.decoder == "autoregressive":
-            ar_inputs = self._ar_inputs(
-                cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch
-            )
-            x1_s2 = self._sec_target(
-                sec_cont, sec_type_idx, self.generator, flatten=False
-            )
+            ar_inputs = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch)
+            x1_s2 = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=False)
         elif self.is_stage2:
-            x1_s2 = self._sec_target(
-                sec_cont, sec_type_idx, self.generator, flatten=True
-            )
+            x1_s2 = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=True)
 
-        l_gen = self._generator_loss(
-            cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=ar_inputs
-        )
-        l_nsec, nsec_acc = self._n_sec_loss(
-            cond_cont, cond_cat, stage1_ctx, n_sec, device
-        )
+        l_gen = self._generator_loss(cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=ar_inputs)
+        l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
 
         l_type, type_acc = self._type_loss(
             cond_cont,
@@ -653,11 +610,7 @@ class FlowDDPMStageTrainer(StageTrainer):
             l_proc = self.router.classify_loss(cond_cont, cond_cat, proc_idx)
             l_entropy = self.router.entropy_loss(cond_cont, cond_cat)
 
-        total = (
-            self.spec.lambda_weight * l_gen
-            + self.spec.n_sec_lambda * l_nsec
-            + self.particle_type_lambda * l_type
-        )
+        total = self.spec.lambda_weight * l_gen + self.spec.n_sec_lambda * l_nsec + self.particle_type_lambda * l_type
         if self.spec.lambda_balance > 0:
             total = total + self.spec.lambda_balance * l_balance
         if self.spec.lambda_proc > 0:
@@ -698,9 +651,7 @@ class FlowDDPMStageTrainer(StageTrainer):
 
     @torch.no_grad()
     def val_loss(self, batch: tuple, device: torch.device) -> dict:
-        return {
-            key: value.item() for key, value in self._compute(batch, device).items()
-        }
+        return {key: value.item() for key, value in self._compute(batch, device).items()}
 
     # --- reporting ------------------------------------------------------
 
@@ -710,9 +661,7 @@ class FlowDDPMStageTrainer(StageTrainer):
     def summary(self, means: dict) -> str:
         return f"{self.name}[loss={means.get('loss', 0.0):.3f}]"
 
-    def val_objective(
-        self, train_means: dict, val_means: dict, marginal_kl: float
-    ) -> float:
+    def val_objective(self, train_means: dict, val_means: dict, marginal_kl: float) -> float:
         return val_means.get("loss", 0.0)
 
 
@@ -799,15 +748,8 @@ class WGANStageTrainer(StageTrainer):
 
         if self.decoder == "autoregressive":
             epoch = global_step // self.spec.steps_per_epoch
-            ar = self._ar_inputs(
-                cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch
-            )
-            real = (
-                self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=False).reshape(
-                    B, -1
-                )
-                * mask
-            )
+            ar = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch)
+            real = self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=False).reshape(B, -1) * mask
             z = torch.randn(B, k_max, self.model.noise_dim, device=device)
             fake_raw = self.model(
                 z,
@@ -891,9 +833,7 @@ class WGANStageTrainer(StageTrainer):
 
         # --- generator (+ n_sec) step ---
         did_g_step = global_step % self.n_critic == 0
-        l_nsec, nsec_acc = self._n_sec_loss(
-            cond_cont, cond_cat, stage1_ctx, n_sec, device
-        )
+        l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
 
         # On a non-generator-step batch with no n_sec_head on this stage
         # (n_sec now defaults to stage 2), there's nothing for
@@ -902,9 +842,7 @@ class WGANStageTrainer(StageTrainer):
         skip_g_step = not did_g_step and self.model.n_sec_head is None
         if did_g_step:
             g_loss_adv = generator_loss(critic_fn, fake)
-            g_loss = (
-                self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * l_nsec
-            )
+            g_loss = self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * l_nsec
         else:
             g_loss_adv = torch.zeros((), device=device)
             g_loss = self.spec.n_sec_lambda * l_nsec
@@ -941,14 +879,9 @@ class WGANStageTrainer(StageTrainer):
         return stats["d_loss"] + stats["g_loss"]
 
     def summary(self, means: dict) -> str:
-        return (
-            f"{self.name}[d={means.get('d_loss', 0.0):.3f} "
-            f"g={means.get('g_loss', 0.0):.3f}]"
-        )
+        return f"{self.name}[d={means.get('d_loss', 0.0):.3f} g={means.get('g_loss', 0.0):.3f}]"
 
-    def val_objective(
-        self, train_means: dict, val_means: dict, marginal_kl: float
-    ) -> float:
+    def val_objective(self, train_means: dict, val_means: dict, marginal_kl: float) -> float:
         """No monotone per-batch WGAN loss fit for averaging, so
         best-checkpoint selection uses the real marginal-KL signal when
         `validate_marginals` produced one, and falls back to this epoch's own
@@ -1002,8 +935,7 @@ def build_stage_trainers(
         if spec.generator == "wgan":
             critic = critics.get(name)
             assert critic is not None, (
-                f"{name}_model.generator='wgan' requires a critic (see "
-                "giant.model.network.build_critics)"
+                f"{name}_model.generator='wgan' requires a critic (see giant.model.network.build_critics)"
             )
             trainers[name] = WGANStageTrainer(spec, model, critic, device)
         else:
diff --git a/giant/validate.py b/giant/validate.py
index c193f4f..1d3b5a0 100644
--- a/giant/validate.py
+++ b/giant/validate.py
@@ -8,9 +8,7 @@ from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
 _SEC_PHYS_NAMES = ["log_mass", "charge"]
 
 
-def _histogram_kl(
-    p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8
-) -> float:
+def _histogram_kl(p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8) -> float:
     """KL(P || Q) between two 1D samples, estimated via a shared histogram."""
     lo = min(p_samples.min(), q_samples.min())
     hi = max(p_samples.max(), q_samples.max())
@@ -33,9 +31,7 @@ def _bincount_frac(x: np.ndarray, minlength: int) -> np.ndarray:
     return counts / total if total > 0 else counts
 
 
-def _categorical_kl(
-    real_idx: np.ndarray, gen_idx: np.ndarray, n_classes: int, eps: float = 1e-8
-) -> float:
+def _categorical_kl(real_idx: np.ndarray, gen_idx: np.ndarray, n_classes: int, eps: float = 1e-8) -> float:
     """KL(P_real || Q_gen) between two class-index samples over `n_classes`
     categories, estimated from bincount fractions. NaN if either side has no
     valid samples (mirrors `_histogram_kl`'s empty-input handling)."""
@@ -48,9 +44,7 @@ def _categorical_kl(
     return float(np.sum(p * np.log(p / q)))
 
 
-def _embedding_nearest_class(
-    vectors: torch.Tensor, emb_weight: torch.Tensor
-) -> np.ndarray:
+def _embedding_nearest_class(vectors: torch.Tensor, emb_weight: torch.Tensor) -> np.ndarray:
     """Nearest row index (L1) of `vectors` (..., emb_dim) against `emb_weight`
     (vocab, emb_dim) — same computation as
     `giant.particles.decode_embedding_nearest`, but returning the raw class
@@ -109,11 +103,7 @@ def validate_marginals(
         sec_decoder.eval()
 
     k_max = sec_decoder.k_max if sec_decoder is not None else 0
-    target = (
-        sec_decoder.particle_type_cfg.get("target", "physical")
-        if sec_decoder is not None
-        else "physical"
-    )
+    target = sec_decoder.particle_type_cfg.get("target", "physical") if sec_decoder is not None else "physical"
 
     all_real, all_gen = [], []
     all_n_sec_real, all_n_sec_pred = [], []
@@ -131,9 +121,7 @@ def validate_marginals(
         cond_cont = cond_cont.to(device)
         cond_cat = cond_cat.to(device)
 
-        gen, n_sec_pred = sample_stage1(
-            stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps
-        )
+        gen, n_sec_pred = sample_stage1(stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps)
 
         all_real.append(x1.numpy())
         all_gen.append(gen.cpu().numpy())
@@ -141,9 +129,7 @@ def validate_marginals(
         if sec_decoder is None:
             continue
 
-        n_sec_pred = resolve_n_sec(
-            stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred
-        )
+        n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred)
         n_sec_pred_np = n_sec_pred.cpu().numpy()
         n_sec_np = n_sec.numpy()
         all_n_sec_real.append(n_sec_np)
@@ -155,9 +141,7 @@ def validate_marginals(
         sec_cont_pred, sec_type_pred, sec_valid_pred = sample_stage2(
             sec_decoder, cond_cont, cond_cat, gen, n_sec_pred, steps=steps
         )
-        gen_frac = 1.0 / (
-            1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64))
-        )
+        gen_frac = 1.0 / (1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64)))
         gen_valid = sec_valid_pred.cpu().numpy()
 
         if target == "physical":
@@ -182,17 +166,9 @@ def validate_marginals(
     real = np.concatenate(all_real, axis=0)
     generated = np.concatenate(all_gen, axis=0)
 
-    kl_divergence = np.array(
-        [
-            _histogram_kl(real[:, j], generated[:, j], bins=kl_bins)
-            for j in range(real.shape[1])
-        ]
-    )
+    kl_divergence = np.array([_histogram_kl(real[:, j], generated[:, j], bins=kl_bins) for j in range(real.shape[1])])
 
-    header = (
-        f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} "
-        f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
-    )
+    header = f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} {'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
     print(f"\n{header}")
     print("-" * len(header))
     for j, name in enumerate(LOCAL_TARGET_NAMES):
@@ -215,10 +191,7 @@ def validate_marginals(
     n_sec_accuracy = float((n_sec_real == n_sec_pred_all).mean())
 
     energy_fraction_kl = np.full(k_max, np.nan)
-    print(
-        f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f}  "
-        f"mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}"
-    )
+    print(f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f}  mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}")
     n_sec_dist_header = f"{'n_sec value':<20} {'real_frac':>10} {'gen_frac':>10}"
     print(n_sec_dist_header)
     print("-" * len(n_sec_dist_header))
@@ -240,10 +213,7 @@ def validate_marginals(
             continue
         kl = _histogram_kl(r, g, bins=kl_bins)
         energy_fraction_kl[j] = kl
-        print(
-            f"{j:<24} {r.mean():>10.4f} {g.mean():>10.4f} "
-            f"{r.std():>10.4f} {g.std():>10.4f} {kl:>14.4f}"
-        )
+        print(f"{j:<24} {r.mean():>10.4f} {g.mean():>10.4f} {r.std():>10.4f} {g.std():>10.4f} {kl:>14.4f}")
 
     result.update(
         {
@@ -259,12 +229,7 @@ def validate_marginals(
         phys_gen = np.concatenate(all_phys_gen, axis=0)  # (M, 2)
 
         if len(phys_real) > 0 and len(phys_gen) > 0:
-            phys_kl = np.array(
-                [
-                    _histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins)
-                    for j in range(2)
-                ]
-            )
+            phys_kl = np.array([_histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins) for j in range(2)])
         else:
             phys_kl = np.full(2, np.nan)
 
@@ -278,21 +243,14 @@ def validate_marginals(
             if len(r) == 0 or len(g) == 0:
                 continue
             print(
-                f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} "
-                f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
+                f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} {r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
             )
 
-        result.update(
-            {"phys_real": phys_real, "phys_generated": phys_gen, "phys_kl": phys_kl}
-        )
+        result.update({"phys_real": phys_real, "phys_generated": phys_gen, "phys_kl": phys_kl})
     else:
         type_class_real = np.concatenate(all_type_class_real, axis=0)
         type_class_gen = np.concatenate(all_type_class_gen, axis=0)
-        n_classes = (
-            sec_decoder.type_dim
-            if target == "onehot"
-            else sec_decoder.cond_enc.pdg_emb.weight.size(0)
-        )
+        n_classes = sec_decoder.type_dim if target == "onehot" else sec_decoder.cond_enc.pdg_emb.weight.size(0)
         type_class_kl = _categorical_kl(type_class_real, type_class_gen, n_classes)
 
         print(
diff --git a/pyproject.toml b/pyproject.toml
index a8c08f3..1d7538b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -52,6 +52,9 @@ analysis = [
 giant = "giant.cli:app"
 dwarf = "scripts.dwarf:app"
 
+[tool.ruff]
+line-length = 120
+
 [tool.coverage.run]
 source = ["giant", "scripts"]
 omit = ["*/legacy/*"]
diff --git a/scripts/bump_dataset_version.py b/scripts/bump_dataset_version.py
index 216748a..27ecabe 100644
--- a/scripts/bump_dataset_version.py
+++ b/scripts/bump_dataset_version.py
@@ -82,9 +82,7 @@ def _max_index(parent: Path, pattern: re.Pattern) -> int:
 
 def _git_user_name() -> str | None:
     try:
-        out = subprocess.run(
-            ["git", "config", "user.name"], capture_output=True, text=True, timeout=2
-        )
+        out = subprocess.run(["git", "config", "user.name"], capture_output=True, text=True, timeout=2)
     except (OSError, subprocess.SubprocessError):
         # OSError (e.g. git not on PATH) and subprocess.SubprocessError
         # (e.g. TimeoutExpired) are unrelated hierarchies — TimeoutExpired
@@ -142,9 +140,7 @@ def plan_bump_schema(
     raw_gen_dir = root / "raw" / kind / gen_tag
     processed_gen_dir = root / "processed" / kind / gen_tag
     if not raw_gen_dir.is_dir() and not processed_gen_dir.is_dir():
-        raise SystemExit(
-            f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first"
-        )
+        raise SystemExit(f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first")
     if target is not None:
         if not SCHEMA_RE.match(target):
             raise SystemExit(f"error: --to must look like 'schemaN', got {target!r}")
@@ -154,9 +150,7 @@ def plan_bump_schema(
         schema_tag = f"schema{next_schema}"
     new_dirs = [processed_gen_dir / schema_tag]
     by_suffix = f" ({by})" if by else ""
-    log_line = (
-        f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
-    )
+    log_line = f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
     return new_dirs, log_line
 
 
@@ -212,9 +206,7 @@ def _manifest_referenced_files(pools_root: Path) -> set[Path]:
     return referenced
 
 
-def _referenced_root_count(
-    raw_gen_dir: Path, processed_gen_dir: Path
-) -> tuple[int, int]:
+def _referenced_root_count(raw_gen_dir: Path, processed_gen_dir: Path) -> tuple[int, int]:
     """(total .root files, count with a same-named .parquet under any schema) for one gen."""
     if not raw_gen_dir.is_dir():
         return 0, 0
@@ -243,9 +235,7 @@ def _referenced_root_count(
     return total, referenced
 
 
-def _referenced_parquet_count(
-    schema_dir: Path, manifest_referenced: set[Path]
-) -> tuple[int, int]:
+def _referenced_parquet_count(schema_dir: Path, manifest_referenced: set[Path]) -> tuple[int, int]:
     """(total .parquet files, count listed in at least one manifest) for one schema dir."""
     if not schema_dir.is_dir():
         return 0, 0
@@ -342,11 +332,7 @@ def print_status(root: Path) -> None:
     grand_files = 0
     for kind_dir in sorted(p for p in raw_root.iterdir() if p.is_dir()):
         kind = kind_dir.name
-        gens = sorted(
-            int(m.group(1))
-            for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir())
-            if m
-        )
+        gens = sorted(int(m.group(1)) for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir()) if m)
         print(_colorize(f"{kind}/", "kind"))
         kind_total = 0
         kind_files = 0
@@ -359,25 +345,18 @@ def print_status(root: Path) -> None:
             schemas = sorted(
                 int(m.group(1))
                 for m in (
-                    SCHEMA_RE.match(p.name)
-                    for p in (schema_dir.iterdir() if schema_dir.is_dir() else [])
-                    if p.is_dir()
+                    SCHEMA_RE.match(p.name) for p in (schema_dir.iterdir() if schema_dir.is_dir() else []) if p.is_dir()
                 )
                 if m
             )
             schema_sizes = {s: _du(schema_dir / f"schema{s}") for s in schemas}
             schema_counts = {
-                s: _referenced_parquet_count(
-                    schema_dir / f"schema{s}", manifest_referenced
-                )
-                for s in schemas
+                s: _referenced_parquet_count(schema_dir / f"schema{s}", manifest_referenced) for s in schemas
             }
             processed_size = sum(schema_sizes.values())
             processed_files = sum(c[0] for c in schema_counts.values())
             processed_referenced = sum(c[1] for c in schema_counts.values())
-            raw_files, raw_referenced = _referenced_root_count(
-                raw_gen_dir, processed_gen_dir
-            )
+            raw_files, raw_referenced = _referenced_root_count(raw_gen_dir, processed_gen_dir)
             gen_total = raw_size + processed_size
             gen_files = raw_files + processed_files
             kind_total += gen_total
@@ -425,9 +404,7 @@ def print_status(root: Path) -> None:
                         print(_reason_line(schema_reason, indent=4))
             else:
                 print(_colorize("      (none)", "schema"))
-        print(
-            _row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files)
-        )
+        print(_row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files))
         print()
         grand_total += kind_total
         grand_files += kind_files
@@ -526,13 +503,8 @@ def plan_update_manifest(
     return result, missing
 
 
-def apply_update_manifest(
-    manifest_path: Path, lines: list[tuple[str, str | None]]
-) -> None:
-    out = [
-        replacement if replacement is not None else original
-        for original, replacement in lines
-    ]
+def apply_update_manifest(manifest_path: Path, lines: list[tuple[str, str | None]]) -> None:
+    out = [replacement if replacement is not None else original for original, replacement in lines]
     manifest_path.write_text("\n".join(out) + "\n")
 
 
@@ -552,9 +524,7 @@ def _resolve_manifest_files(manifest_path: Path) -> list[Path]:
     return files
 
 
-def plan_create_manifest(
-    output_path: Path, parquet_files: list[Path]
-) -> tuple[list[str], list[Path], list[Path]]:
+def plan_create_manifest(output_path: Path, parquet_files: list[Path]) -> tuple[list[str], list[Path], list[Path]]:
     """Return (relative_lines, missing_files, resolved_abs_paths)."""
     manifest_dir = output_path.resolve().parent
     lines: list[str] = []
@@ -569,9 +539,7 @@ def plan_create_manifest(
     return lines, missing, resolved
 
 
-def check_holdout_overlap(
-    output_path: Path, resolved_new_files: list[Path]
-) -> list[tuple[str, Path]]:
+def check_holdout_overlap(output_path: Path, resolved_new_files: list[Path]) -> list[tuple[str, Path]]:
     """Return (other_manifest_name, file) pairs where new files clash with existing manifests.
 
     The check is triggered when output_path is (or will be) holdout.manifest, or when a
@@ -641,9 +609,7 @@ def _run_bump(
     if gen is None:
         new_dirs, log_line = plan_bump_gen(root_path, kind, reason, by, date, to)
     else:
-        new_dirs, log_line = plan_bump_schema(
-            root_path, kind, gen, reason, by, date, to
-        )
+        new_dirs, log_line = plan_bump_schema(root_path, kind, gen, reason, by, date, to)
 
     print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
     print("new directories:")
diff --git a/scripts/check_migration_v02_v03.py b/scripts/check_migration_v02_v03.py
index 3136d11..3386a91 100644
--- a/scripts/check_migration_v02_v03.py
+++ b/scripts/check_migration_v02_v03.py
@@ -84,19 +84,12 @@ def main() -> int:
     mode = model_config.get("mode", "flow")
     routed = bool((model_config.get("router") or {}).get("enabled"))
     print(f"checkpoint: {args.checkpoint}")
-    print(
-        f"  mode={mode!r}  conditioning={model_config.get('conditioning')!r}  "
-        f"routed={routed}  ema={args.ema}"
-    )
+    print(f"  mode={mode!r}  conditioning={model_config.get('conditioning')!r}  routed={routed}  ema={args.ema}")
 
     stage1_key = "model_ema" if args.ema and "model_ema" in ckpt else "model"
-    stage2_key = (
-        "sec_decoder_ema" if args.ema and "sec_decoder_ema" in ckpt else "sec_decoder"
-    )
+    stage2_key = "sec_decoder_ema" if args.ema and "sec_decoder_ema" in ckpt else "sec_decoder"
     if args.ema and stage1_key == "model":
-        print(
-            "  warning: --ema requested but no model_ema in checkpoint, using raw weights"
-        )
+        print("  warning: --ema requested but no model_ema in checkpoint, using raw weights")
 
     # --- old side: the frozen v0.2 snapshot, loaded with the checkpoint's own weights ---
     old_stage1, old_stage2 = legacy.build_models(model_config)
@@ -119,9 +112,7 @@ def main() -> int:
         print("PASS (construction only, routed checkpoint)")
         return 0
 
-    remapped1, remapped2 = net.migrate_legacy_state_dict(
-        ckpt[stage1_key], ckpt[stage2_key]
-    )
+    remapped1, remapped2 = net.migrate_legacy_state_dict(ckpt[stage1_key], ckpt[stage2_key])
     missing1, unexpected1 = new_stage1.load_state_dict(remapped1, strict=True)
     missing2, unexpected2 = new_stage2.load_state_dict(remapped2, strict=True)
     if missing1 or unexpected1 or missing2 or unexpected2:
@@ -132,9 +123,7 @@ def main() -> int:
     new_stage1.eval()
     new_stage2.eval()
 
-    cond_cont, cond_cat, x1, x2, t, z1, z2 = _random_batch(
-        model_config, args.batch, args.seed
-    )
+    cond_cont, cond_cat, x1, x2, t, z1, z2 = _random_batch(model_config, args.batch, args.seed)
 
     ok = True
     with torch.no_grad():
diff --git a/scripts/create_root_files.py b/scripts/create_root_files.py
index d960c78..66ba11a 100644
--- a/scripts/create_root_files.py
+++ b/scripts/create_root_files.py
@@ -64,9 +64,7 @@ def parse_detector_spec(spec: str) -> tuple[str, str | None]:
     if ":" in spec:
         label, config = spec.split(":", 1)
         if not label or not config:
-            raise PlanError(
-                f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG"
-            )
+            raise PlanError(f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG")
         return label, config
     return spec, None
 
@@ -94,9 +92,7 @@ def plan_jobs(
         raise PlanError(f"--gen must look like 'genN', got {gen!r}")
     gen_dir = dataset_root / "raw" / kind / gen
     if not gen_dir.is_dir():
-        raise PlanError(
-            f"{gen_dir} doesn't exist — run bump_dataset_version.py bump-gen first"
-        )
+        raise PlanError(f"{gen_dir} doesn't exist — run bump_dataset_version.py bump-gen first")
 
     jobs = []
     for spec in detector_specs:
@@ -124,9 +120,7 @@ def job_seed(kind: str, gen: str, job: SimJob, energy_gev: float | None) -> int:
     return zlib.crc32(key.encode()) & 0x7FFFFFFF
 
 
-def build_cmd(
-    executable: Path, job: SimJob, events_per_file: int, energy_gev: float | None
-) -> list[str]:
+def build_cmd(executable: Path, job: SimJob, events_per_file: int, energy_gev: float | None) -> list[str]:
     """minicalosim executables take positional `[configName] nEvents [energy_GeV]`."""
     cmd = [str(executable)]
     if job.config:
@@ -147,10 +141,7 @@ def run_job(
     gen: str,
     tmp_root: Path,
 ) -> JobResult:
-    workdir = (
-        tmp_root
-        / f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
-    )
+    workdir = tmp_root / f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
     workdir.mkdir(parents=True)
 
     cmd = build_cmd(executable, job, events_per_file, energy_gev)
@@ -174,20 +165,12 @@ def run_job(
             job,
             False,
             None,
-            f"expected exactly one .root output in {workdir}, found {len(produced)}: "
-            f"{[p.name for p in produced]}",
+            f"expected exactly one .root output in {workdir}, found {len(produced)}: {[p.name for p in produced]}",
             result.stdout,
             result.stderr,
         )
 
-    dest = (
-        dataset_root
-        / "raw"
-        / kind
-        / gen
-        / job.detector
-        / f"shard-{job.shard_index:03d}.root"
-    )
+    dest = dataset_root / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
     if dest.exists():
         return JobResult(
             job,
@@ -279,14 +262,7 @@ def run_make_root(
     print(f"executable: {executable}")
     for job in planned_jobs:
         cmd = build_cmd(executable, job, events_per_file, energy_gev)
-        dest = (
-            dataset_root_path
-            / "raw"
-            / kind
-            / gen
-            / job.detector
-            / f"shard-{job.shard_index:03d}.root"
-        )
+        dest = dataset_root_path / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
         seed = job_seed(kind, gen, job, energy_gev)
         print(f"  MINICALOSIM_SEED={seed} {' '.join(cmd)}  ->  {dest}")
 
diff --git a/scripts/dwarf.py b/scripts/dwarf.py
index 5e7cff5..674e252 100644
--- a/scripts/dwarf.py
+++ b/scripts/dwarf.py
@@ -80,8 +80,7 @@ def convert(
         typer.Option(
             "--output",
             "-o",
-            help="Output Parquet file (default: <input>.parquet). Only valid "
-            "with a single input file and --jobs 1.",
+            help="Output Parquet file (default: <input>.parquet). Only valid with a single input file and --jobs 1.",
         ),
     ] = None,
     batch_size: Annotated[
@@ -91,9 +90,7 @@ def convert(
             help="Uproot read batch size, e.g. '100 MB' or '500000' (rows)",
         ),
     ] = "100 MB",
-    tree: Annotated[
-        str, typer.Option("--tree", help="Tree name inside the ROOT file")
-    ] = "Steps",
+    tree: Annotated[str, typer.Option("--tree", help="Tree name inside the ROOT file")] = "Steps",
     compression: Annotated[
         Compression, typer.Option("--compression", help="Parquet compression codec")
     ] = Compression.snappy,
@@ -129,15 +126,11 @@ def convert(
         raise typer.Exit(1)
     _warn_if_exceeds_shared_quota(jobs, "--jobs")
 
-    compression_value = (
-        "uncompressed" if compression is Compression.none else compression.value
-    )
+    compression_value = "uncompressed" if compression is Compression.none else compression.value
 
     if jobs == 1:
         if output is not None and len(root_files) > 1:
-            typer.echo(
-                "error: --output can only be used with a single input file", err=True
-            )
+            typer.echo("error: --output can only be used with a single input file", err=True)
             raise typer.Exit(1)
         total_orphaned = 0
         for root_file in root_files:
@@ -150,10 +143,7 @@ def convert(
             )
             total_orphaned += n_orphaned
         if total_orphaned:
-            typer.echo(
-                f"\n{total_orphaned} orphaned child track(s) dropped across "
-                f"{len(root_files)} file(s)."
-            )
+            typer.echo(f"\n{total_orphaned} orphaned child track(s) dropped across {len(root_files)} file(s).")
         return
 
     if output is not None:
@@ -176,9 +166,7 @@ def convert(
 
 @app.command()
 def migrate(
-    root: Annotated[
-        Path, typer.Argument(help="Dataset root to migrate in place")
-    ] = _DATASET_ROOT_DEFAULT,
+    root: Annotated[Path, typer.Argument(help="Dataset root to migrate in place")] = _DATASET_ROOT_DEFAULT,
     execute: Annotated[
         bool,
         typer.Option(
@@ -190,8 +178,7 @@ def migrate(
         bool,
         typer.Option(
             "--copy",
-            help="Copy instead of move, leaving the originals in place "
-            "(e.g. if another process is still reading them)",
+            help="Copy instead of move, leaving the originals in place (e.g. if another process is still reading them)",
         ),
     ] = False,
 ) -> None:
@@ -202,12 +189,8 @@ def migrate(
 @app.command("bump-gen")
 def bump_gen(
     reason: Annotated[str, typer.Option("--reason", help="Why this gen exists")],
-    kind: Annotated[
-        str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
-    ] = "steps",
-    by: Annotated[
-        Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
-    ] = None,
+    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
+    by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
     date: Annotated[
         Optional[str],
         typer.Option("--date", help="Override date (default: today, ISO)"),
@@ -220,12 +203,8 @@ def bump_gen(
             help="Target gen tag (default: one past the current highest)",
         ),
     ] = None,
-    execute: Annotated[
-        bool, typer.Option("--execute", help="Apply (default: dry run)")
-    ] = False,
-    root: Annotated[
-        Path, typer.Option("--root", help="Dataset root")
-    ] = _DATASET_ROOT_DEFAULT,
+    execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
+    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
     """Cut a new raw generation."""
     run_bump_gen(
@@ -243,12 +222,8 @@ def bump_gen(
 def bump_schema(
     gen: Annotated[str, typer.Option("--gen", help="Existing gen tag, e.g. gen1")],
     reason: Annotated[str, typer.Option("--reason", help="Why this schema exists")],
-    kind: Annotated[
-        str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
-    ] = "steps",
-    by: Annotated[
-        Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
-    ] = None,
+    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
+    by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
     date: Annotated[
         Optional[str],
         typer.Option("--date", help="Override date (default: today, ISO)"),
@@ -261,12 +236,8 @@ def bump_schema(
             help="Target schema tag (default: one past the current highest)",
         ),
     ] = None,
-    execute: Annotated[
-        bool, typer.Option("--execute", help="Apply (default: dry run)")
-    ] = False,
-    root: Annotated[
-        Path, typer.Option("--root", help="Dataset root")
-    ] = _DATASET_ROOT_DEFAULT,
+    execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
+    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
     """Cut a new schema within a gen."""
     run_bump_schema(
@@ -283,9 +254,7 @@ def bump_schema(
 
 @app.command()
 def status(
-    root: Annotated[
-        Path, typer.Option("--root", help="Dataset root")
-    ] = _DATASET_ROOT_DEFAULT,
+    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
     """List existing gens/schemas per kind."""
     run_status(str(root))
@@ -293,9 +262,7 @@ def status(
 
 @app.command("update-manifest")
 def update_manifest(
-    manifests: Annotated[
-        list[Path], typer.Argument(help="One or more .manifest files to update")
-    ],
+    manifests: Annotated[list[Path], typer.Argument(help="One or more .manifest files to update")],
     schema: Annotated[
         Optional[str],
         typer.Option(
@@ -306,9 +273,7 @@ def update_manifest(
     ] = None,
     gen: Annotated[
         Optional[str],
-        typer.Option(
-            "--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"
-        ),
+        typer.Option("--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"),
     ] = None,
     execute: Annotated[
         bool,
@@ -316,9 +281,7 @@ def update_manifest(
     ] = False,
 ) -> None:
     """Repoint manifest(s) to a new gen and/or schema, verifying all target files exist."""
-    run_update_manifest(
-        [str(m) for m in manifests], schema=schema, execute=execute, gen=gen
-    )
+    run_update_manifest([str(m) for m in manifests], schema=schema, execute=execute, gen=gen)
 
 
 @app.command("create-manifest")
@@ -333,22 +296,15 @@ def create_manifest(
         typer.Option(
             "--pool",
             metavar="DETECTOR",
-            help="Detector name; combined with --type and --root to form "
-            "<root>/pools/<detector>/<type>.manifest",
+            help="Detector name; combined with --type and --root to form <root>/pools/<detector>/<type>.manifest",
         ),
     ] = None,
     type_: Annotated[
         Optional[PoolType],
-        typer.Option(
-            "--type", help="Pool type — full, holdout, or dev (required with --pool)"
-        ),
+        typer.Option("--type", help="Pool type — full, holdout, or dev (required with --pool)"),
     ] = None,
-    root: Annotated[
-        Path, typer.Option("--root", help="Dataset root (used with --pool)")
-    ] = _DATASET_ROOT_DEFAULT,
-    execute: Annotated[
-        bool, typer.Option("--execute", help="Write the manifest (default: dry run)")
-    ] = False,
+    root: Annotated[Path, typer.Option("--root", help="Dataset root (used with --pool)")] = _DATASET_ROOT_DEFAULT,
+    execute: Annotated[bool, typer.Option("--execute", help="Write the manifest (default: dry run)")] = False,
     force: Annotated[
         bool,
         typer.Option("--force", help="Overwrite the manifest if it already exists"),
@@ -368,9 +324,7 @@ def create_manifest(
 
 @app.command("make-root")
 def make_root(
-    executable: Annotated[
-        Path, typer.Option("--executable", help="Built minicalosim run_* executable")
-    ],
+    executable: Annotated[Path, typer.Option("--executable", help="Built minicalosim run_* executable")],
     detector: Annotated[
         list[str],
         typer.Option(
@@ -382,15 +336,9 @@ def make_root(
             "Repeatable.",
         ),
     ],
-    num_files: Annotated[
-        int, typer.Option("--num-files", help="New shards to create per detector")
-    ],
-    events_per_file: Annotated[
-        int, typer.Option("--events-per-file", help="nEvents passed to the executable")
-    ],
-    gen: Annotated[
-        str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")
-    ],
+    num_files: Annotated[int, typer.Option("--num-files", help="New shards to create per detector")],
+    events_per_file: Annotated[int, typer.Option("--events-per-file", help="nEvents passed to the executable")],
+    gen: Annotated[str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")],
     energy_gev: Annotated[
         float | None,
         typer.Option(
@@ -401,20 +349,12 @@ def make_root(
             "to name the dataset accordingly.",
         ),
     ] = None,
-    kind: Annotated[
-        str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
-    ] = "steps",
-    dataset_root: Annotated[
-        Path, typer.Option("--dataset-root", help="Dataset root")
-    ] = _DATASET_ROOT_DEFAULT,
-    jobs: Annotated[
-        int, typer.Option("--jobs", "-j", help="Parallel simulation runs (default: 4)")
-    ] = 4,
+    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
+    dataset_root: Annotated[Path, typer.Option("--dataset-root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
+    jobs: Annotated[int, typer.Option("--jobs", "-j", help="Parallel simulation runs (default: 4)")] = 4,
     execute: Annotated[
         bool,
-        typer.Option(
-            "--execute", help="Actually run jobs (default: dry run / print plan)"
-        ),
+        typer.Option("--execute", help="Actually run jobs (default: dry run / print plan)"),
     ] = False,
 ) -> None:
     """Generate new ROOT shards via a minicalosim executable."""
@@ -441,9 +381,7 @@ class OracleMethod(str, Enum):
 
 @app.command("build-geometry-oracle")
 def build_geometry_oracle(
-    data: Annotated[
-        Path, typer.Argument(help="Steps parquet file or directory of steps files")
-    ],
+    data: Annotated[Path, typer.Argument(help="Steps parquet file or directory of steps files")],
     out: Annotated[Path, typer.Option("--out", "-o", help="Output oracle .pkl path")],
     method: Annotated[
         OracleMethod,
@@ -456,9 +394,7 @@ def build_geometry_oracle(
             ),
         ),
     ] = OracleMethod.slab,
-    k: Annotated[
-        int, typer.Option("--k", help="Neighbours for the knn classifier")
-    ] = 1,
+    k: Annotated[int, typer.Option("--k", help="Neighbours for the knn classifier")] = 1,
     subsample: Annotated[
         int,
         typer.Option("--subsample", help="Max reference points sampled from the data"),
@@ -504,9 +440,7 @@ def build_geometry_oracle(
 def warm_cache(
     data: Annotated[
         Path,
-        typer.Argument(
-            help="Parquet file, directory, or .manifest — same as `giant train`'s"
-        ),
+        typer.Argument(help="Parquet file, directory, or .manifest — same as `giant train`'s"),
     ],
     val_fraction: Annotated[
         float,
@@ -518,16 +452,13 @@ def warm_cache(
     ] = 0.1,
     seed: Annotated[
         int,
-        typer.Option(
-            "--seed", "-s", help="Must match the `giant train` run(s) to warm for"
-        ),
+        typer.Option("--seed", "-s", help="Must match the `giant train` run(s) to warm for"),
     ] = 0,
     particle_conditioning: Annotated[
         Conditioning,
         typer.Option(
             "--particle-conditioning",
-            help="Must match the `giant train` run(s)' conditioning.particle.type "
-            "to warm for",
+            help="Must match the `giant train` run(s)' conditioning.particle.type to warm for",
         ),
     ] = Conditioning.physical,
     material_conditioning: Annotated[
@@ -543,21 +474,14 @@ def warm_cache(
         bool,
         typer.Option(
             "--router/--no-router",
-            help="Warm the process vocabulary too (only takes effect with "
-            "--router-type process)",
+            help="Warm the process vocabulary too (only takes effect with --router-type process)",
         ),
     ] = False,
-    router_type: Annotated[
-        str, typer.Option("--router-type", help="Router implementation name")
-    ] = "energy",
-    n_experts: Annotated[
-        int, typer.Option("--n-experts", help="Number of routed experts")
-    ] = 4,
+    router_type: Annotated[str, typer.Option("--router-type", help="Router implementation name")] = "energy",
+    n_experts: Annotated[int, typer.Option("--n-experts", help="Number of routed experts")] = 4,
     rebuild: Annotated[
         bool,
-        typer.Option(
-            "--rebuild", help="Ignore any existing sidecar and recompute every section"
-        ),
+        typer.Option("--rebuild", help="Ignore any existing sidecar and recompute every section"),
     ] = False,
 ) -> None:
     """Precompute `giant train`'s setup-stage sidecar for `data` ahead of time.
diff --git a/scripts/geometry_oracle.py b/scripts/geometry_oracle.py
index e5a36fb..762d288 100644
--- a/scripts/geometry_oracle.py
+++ b/scripts/geometry_oracle.py
@@ -38,9 +38,7 @@ def run_build_geometry_oracle(
         n_bins=n_bins,
     )
 
-    print(
-        f"method: {method}  reference points: {oracle.metadata['n_reference_points']:,}"
-    )
+    print(f"method: {method}  reference points: {oracle.metadata['n_reference_points']:,}")
     print("classes (material, layer_id):")
     for material, layer_id in oracle.classes:
         print(f"  {material:<12} layer_id={layer_id}")
diff --git a/scripts/hparam_scan.py b/scripts/hparam_scan.py
index 00efd42..0c64f6c 100644
--- a/scripts/hparam_scan.py
+++ b/scripts/hparam_scan.py
@@ -154,9 +154,7 @@ def run_hparam_scan(
         wall_time_s = time.monotonic() - start
 
         if metrics_path.exists():
-            epochs_completed, final_val_loss, best_val_loss = final_metrics(
-                metrics_path
-            )
+            epochs_completed, final_val_loss, best_val_loss = final_metrics(metrics_path)
             append_summary(
                 summary_path,
                 {
@@ -171,11 +169,6 @@ def run_hparam_scan(
                     "wall_time_s": round(wall_time_s, 1),
                 },
             )
-            print(
-                f"[{i}/{len(runs)}] {name} — val_loss {final_val_loss:.4f} "
-                f"({wall_time_s:.1f}s)"
-            )
+            print(f"[{i}/{len(runs)}] {name} — val_loss {final_val_loss:.4f} ({wall_time_s:.1f}s)")
         else:
-            print(
-                f"[{i}/{len(runs)}] {name} — no metrics.csv produced, check train.log"
-            )
+            print(f"[{i}/{len(runs)}] {name} — no metrics.csv produced, check train.log")
diff --git a/scripts/migrate_geant_steps.py b/scripts/migrate_geant_steps.py
index 27e96c9..447e0dc 100644
--- a/scripts/migrate_geant_steps.py
+++ b/scripts/migrate_geant_steps.py
@@ -50,12 +50,8 @@ PREDICTED_RE = re.compile(
     r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)"
     r"_predicted(?P<local>_local)?\.parquet$"
 )
-SHARD_RE = re.compile(
-    r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)\.(?P<ext>root|parquet)$"
-)
-LEGACY_PREDICTED_RE = re.compile(
-    r"^pbwo4_10000events_hits_predicted(?P<local>_local)?\.parquet$"
-)
+SHARD_RE = re.compile(r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)\.(?P<ext>root|parquet)$")
+LEGACY_PREDICTED_RE = re.compile(r"^pbwo4_10000events_hits_predicted(?P<local>_local)?\.parquet$")
 LEGACY_RE = re.compile(r"^pbwo4_10000events_hits\.(?P<ext>root|parquet)$")
 
 
@@ -110,24 +106,9 @@ def plan_moves(src_root: Path) -> tuple[list[tuple[Path, Path]], list[Path]]:
         if m:
             detector, shard, ext = m["detector"], int(m["shard"]), m["ext"]
             if ext == "root":
-                dst = (
-                    src_root
-                    / "raw"
-                    / "steps"
-                    / GEN
-                    / detector
-                    / f"shard-{shard:03d}.root"
-                )
+                dst = src_root / "raw" / "steps" / GEN / detector / f"shard-{shard:03d}.root"
             else:
-                dst = (
-                    src_root
-                    / "processed"
-                    / "steps"
-                    / GEN
-                    / SCHEMA
-                    / detector
-                    / f"shard-{shard:03d}.parquet"
-                )
+                dst = src_root / "processed" / "steps" / GEN / SCHEMA / detector / f"shard-{shard:03d}.parquet"
             moves.append((path, dst))
             continue
 
@@ -135,19 +116,9 @@ def plan_moves(src_root: Path) -> tuple[list[tuple[Path, Path]], list[Path]]:
         if m:
             ext = m["ext"]
             if ext == "root":
-                dst = (
-                    src_root / "raw" / "hits" / LEGACY_GEN / "pbwo4" / "shard-000.root"
-                )
+                dst = src_root / "raw" / "hits" / LEGACY_GEN / "pbwo4" / "shard-000.root"
             else:
-                dst = (
-                    src_root
-                    / "processed"
-                    / "hits"
-                    / LEGACY_GEN
-                    / LEGACY_SCHEMA
-                    / "pbwo4"
-                    / "shard-000.parquet"
-                )
+                dst = src_root / "processed" / "hits" / LEGACY_GEN / LEGACY_SCHEMA / "pbwo4" / "shard-000.parquet"
             moves.append((path, dst))
             continue
 
@@ -165,20 +136,9 @@ def plan_manifests(src_root: Path) -> dict[Path, list[str]]:
         for pool, shards in rules.items():
             manifest_path = manifest_dir / f"{pool}{MANIFEST_SUFFIX}"
             for shard in shards:
-                dst = (
-                    src_root
-                    / "processed"
-                    / "steps"
-                    / GEN
-                    / SCHEMA
-                    / detector
-                    / f"shard-{shard:03d}.parquet"
-                )
+                dst = src_root / "processed" / "steps" / GEN / SCHEMA / detector / f"shard-{shard:03d}.parquet"
                 manifests[manifest_path].append((shard, dst))
-    return {
-        k: [os.path.relpath(dst, start=k.parent) for _, dst in sorted(v)]
-        for k, v in manifests.items()
-    }
+    return {k: [os.path.relpath(dst, start=k.parent) for _, dst in sorted(v)] for k, v in manifests.items()}
 
 
 def run_migration(root: str, execute: bool, copy: bool) -> None:
diff --git a/scripts/profile_analysis_costs.py b/scripts/profile_analysis_costs.py
index de97a2b..60dda89 100644
--- a/scripts/profile_analysis_costs.py
+++ b/scripts/profile_analysis_costs.py
@@ -94,9 +94,7 @@ def _make_rollout(n: int, n_events: int, seed: int) -> pl.DataFrame:
             "post_dx": post_dir[:, 0],
             "post_dy": post_dir[:, 1],
             "post_dz": post_dir[:, 2],
-            "edep": np.where(
-                is_synthetic, np.where(reasons == "escaped", 0.0, pre_E), edep
-            ),
+            "edep": np.where(is_synthetic, np.where(reasons == "escaped", 0.0, pre_E), edep),
             "step_length": np.where(is_synthetic, 0.0, step_length),
             "material": rng.choice(_MATERIALS, size=n),
             "layer_id": rng.integers(0, 30, size=n),
@@ -163,9 +161,7 @@ def _make_reference(n: int, n_events: int, seed: int) -> pl.DataFrame:
     )
 
 
-def _time(
-    spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path
-) -> float:
+def _time(spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path) -> float:
     t0 = time.perf_counter()
     compute_reduced(
         spec_id,
diff --git a/scripts/steps_to_parquet_parallel.py b/scripts/steps_to_parquet_parallel.py
index b2ee41e..1817c36 100644
--- a/scripts/steps_to_parquet_parallel.py
+++ b/scripts/steps_to_parquet_parallel.py
@@ -49,9 +49,7 @@ def latest_schema_tag(processed_gen_dir: Path) -> str | None:
     return best_tag
 
 
-def resolve_destination(
-    root_file: Path, dataset_root: Path, schema_override: str | None
-) -> Path:
+def resolve_destination(root_file: Path, dataset_root: Path, schema_override: str | None) -> Path:
     """Map raw/<kind>/<gen>/<detector>/<file>.root (relative to *dataset_root*)
     to processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet.
 
@@ -66,12 +64,7 @@ def resolve_destination(
         raise DestinationError(f"{root_file} is not under dataset root {dataset_root}")
 
     parts = rel.parts
-    if (
-        len(parts) != 5
-        or parts[0] != "raw"
-        or not GEN_RE.match(parts[2])
-        or not parts[4].endswith(".root")
-    ):
+    if len(parts) != 5 or parts[0] != "raw" or not GEN_RE.match(parts[2]) or not parts[4].endswith(".root"):
         raise DestinationError(
             f"{root_file} does not match raw/<kind>/<gen>/<detector>/<file>.root "
             f"under {dataset_root} (got relative path: {rel})"
@@ -216,13 +209,7 @@ def run_parallel_job(
             print(f"  {root_file}", file=sys.stderr)
         raise SystemExit(1)
 
-    total_orphaned = sum(
-        int(m.group(1))
-        for _, _, stdout, _ in results
-        for m in _ORPHAN_RE.finditer(stdout)
-    )
+    total_orphaned = sum(int(m.group(1)) for _, _, stdout, _ in results for m in _ORPHAN_RE.finditer(stdout))
     if total_orphaned:
-        print(
-            f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s)."
-        )
+        print(f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s).")
     print(f"\nAll {len(results)} conversion(s) completed.")
diff --git a/tests/legacy/network_v02_snapshot.py b/tests/legacy/network_v02_snapshot.py
index ea9364c..a58992d 100644
--- a/tests/legacy/network_v02_snapshot.py
+++ b/tests/legacy/network_v02_snapshot.py
@@ -37,11 +37,7 @@ class SinusoidalEmbedding(nn.Module):
         super().__init__()
         assert dim % 2 == 0, "dim must be even"
         half = dim // 2
-        freqs = torch.exp(
-            -math.log(10000)
-            * torch.arange(half, dtype=torch.float32)
-            / max(half - 1, 1)
-        )
+        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32) / max(half - 1, 1))
         self.register_buffer("freqs", freqs)
 
     def forward(self, t: torch.Tensor) -> torch.Tensor:
@@ -107,9 +103,7 @@ class ConditionEncoder(nn.Module):
             pdg_e = self.pdg_emb(cond_cat[:, 0])
             mat_e = self.mat_emb(cond_cat[:, 1])
         else:
-            particle_phys = cond_cont[
-                :, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM
-            ]
+            particle_phys = cond_cont[:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM]
             material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :]
             pdg_e = self.particle_mlp(particle_phys)
             mat_e = self.material_mlp(material_phys)
@@ -168,12 +162,7 @@ class DenoisingMLP(nn.Module):
         )
         merged_cond_dim = time_dim + cond_out_dim
         self.input_proj = nn.Linear(x_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_proj = nn.Linear(hidden_dim, x_dim)
         # Predicts n_sec as classification over {0, 1, ..., k_max}.
         # Applied to the condition encoding (not the diffused latent).
@@ -286,12 +275,7 @@ class SecondaryDecoder(nn.Module):
         )
         merged_cond_dim = time_dim + cond_out_dim
         self.input_proj = nn.Linear(sec_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_proj = nn.Linear(hidden_dim, sec_dim)
 
     def forward(
@@ -345,12 +329,7 @@ class WGANGenerator(nn.Module):
             conditioning=conditioning,
         )
         self.input_proj = nn.Linear(noise_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_proj = nn.Linear(hidden_dim, x_dim)
         self.n_sec_head = nn.Sequential(
             nn.Linear(cond_out_dim, hidden_dim // 2),
@@ -410,12 +389,7 @@ class Critic(nn.Module):
             conditioning=conditioning,
         )
         self.input_proj = nn.Linear(x_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_norm = nn.LayerNorm(hidden_dim)
         self.out_proj = nn.Linear(hidden_dim, 1)
 
@@ -466,12 +440,7 @@ class WGANSecondaryGenerator(nn.Module):
             conditioning=conditioning,
         )
         self.input_proj = nn.Linear(noise_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_proj = nn.Linear(hidden_dim, sec_dim)
 
     def forward(
@@ -515,12 +484,7 @@ class SecondaryCritic(nn.Module):
             conditioning=conditioning,
         )
         self.input_proj = nn.Linear(sec_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_norm = nn.LayerNorm(hidden_dim)
         self.out_proj = nn.Linear(hidden_dim, 1)
 
@@ -550,9 +514,7 @@ class Router(nn.Module):
     def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         raise NotImplementedError
 
-    def combine_weights(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> torch.Tensor:
+    def combine_weights(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         probs = self.gate(cond_cont, cond_cat)
         if not (self.gumbel and self.training):
             return probs
@@ -562,26 +524,18 @@ class Router(nn.Module):
     def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         return self.gate(cond_cont, cond_cat).argmax(dim=-1)
 
-    def balance_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> torch.Tensor:
+    def balance_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         importance = self.gate(cond_cont, cond_cat).sum(dim=0)  # (n_experts,)
         return (importance.std() / (importance.mean() + 1e-8)) ** 2
 
-    def classify_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
-    ) -> torch.Tensor:
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
         return torch.zeros((), device=cond_cont.device)
 
-    def entropy_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> torch.Tensor:
+    def entropy_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         norm_entropy, _ = self.gate_stats(cond_cont, cond_cat)
         return norm_entropy
 
-    def gate_stats(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    def gate_stats(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         gate = self.gate(cond_cont, cond_cat)  # (B, n_experts)
         row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1)  # (B,)
         norm_entropy = row_entropy.mean() / math.log(self.n_experts)
@@ -602,9 +556,7 @@ def register_router(name: str):
 
 def build_router(name: str, n_experts: int, **kwargs) -> Router:
     if name not in ROUTER_REGISTRY:
-        raise ValueError(
-            f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}"
-        )
+        raise ValueError(f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}")
     cls = ROUTER_REGISTRY[name]
     accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
     filtered = {k: v for k, v in kwargs.items() if k in accepted}
@@ -644,8 +596,7 @@ class EnergyRouter(Router):
         if learn_width or learn_temperature:
             if not (width_min_ratio < 1.0 < width_max_ratio):
                 raise ValueError(
-                    f"width_min_ratio ({width_min_ratio}) and width_max_ratio "
-                    f"({width_max_ratio}) must bracket 1.0"
+                    f"width_min_ratio ({width_min_ratio}) and width_max_ratio ({width_max_ratio}) must bracket 1.0"
                 )
             self._width_lo = width_min_ratio * temperature
             self._width_hi = width_max_ratio * temperature
@@ -658,10 +609,7 @@ class EnergyRouter(Router):
             centers = torch.linspace(-2.0, 2.0, n_experts)
         else:
             if len(centers_init) != n_experts:
-                raise ValueError(
-                    f"centers_init has {len(centers_init)} values, "
-                    f"expected n_experts={n_experts}"
-                )
+                raise ValueError(f"centers_init has {len(centers_init)} values, expected n_experts={n_experts}")
             centers = torch.tensor(list(centers_init), dtype=torch.float32)
         if learn_centers:
             self.centers = nn.Parameter(centers)
@@ -702,9 +650,7 @@ class PdgRouter(Router):
 
     def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         e = self.pdg_emb(cond_cat[:, 0])  # (B, emb_dim)
-        d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(
-            -1
-        )  # (B, n_experts)
+        d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(-1)  # (B, n_experts)
         return torch.softmax(-d2 / self.temperature, dim=-1)
 
 
@@ -736,9 +682,7 @@ class ProcessRouter(Router):
     def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1)
 
-    def classify_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
-    ) -> torch.Tensor:
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
         return F.cross_entropy(self.logits(cond_cont, cond_cat), labels)
 
 
@@ -756,14 +700,10 @@ class ComposedRouter(Router):
         joint = self.routers[0].gate(cond_cont, cond_cat)  # (B, n_0)
         for router in self.routers[1:]:
             g = router.gate(cond_cont, cond_cat)  # (B, n_i)
-            joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(
-                1
-            )  # (B, prod so far)
+            joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(1)  # (B, prod so far)
         return joint
 
-    def classify_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
-    ) -> torch.Tensor:
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
         total = torch.zeros((), device=cond_cont.device)
         for router in self.routers:
             total = total + router.classify_loss(cond_cont, cond_cat, labels)
@@ -796,12 +736,7 @@ class ExpertTrunk(nn.Module):
     ) -> None:
         super().__init__()
         self.input_proj = nn.Linear(in_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_proj = nn.Linear(hidden_dim, in_dim)
 
     def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
@@ -891,9 +826,7 @@ class RoutedDenoisingMLP(nn.Module):
         t_emb = self.time_emb(t)
         c_emb = self.cond_enc(cond_cont, cond_cat)
         cond = torch.cat([t_emb, c_emb], dim=-1)
-        return _route_forward(
-            self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
-        )
+        return _route_forward(self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training)
 
     def predict_n_sec(
         self,
@@ -957,9 +890,7 @@ class RoutedSecondaryDecoder(nn.Module):
         t_emb = self.time_emb(t)
         c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
         cond = torch.cat([t_emb, c_emb], dim=-1)
-        return _route_forward(
-            self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
-        )
+        return _route_forward(self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training)
 
 
 _STAGE1_MODEL_KEYS = {
@@ -1006,9 +937,7 @@ def _parse_composed_axes(router_cfg: dict) -> list[dict]:
 _VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process")
 
 
-def _check_router_conditioning_compat(
-    router_types: list[str], conditioning: str
-) -> None:
+def _check_router_conditioning_compat(router_types: list[str], conditioning: str) -> None:
     bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES))
     if bad and conditioning == "physical":
         raise ValueError(
@@ -1019,9 +948,7 @@ def _check_router_conditioning_compat(
         )
 
 
-def _build_router_from_cfg(
-    router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding"
-) -> Router:
+def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding") -> Router:
     shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
     if router_cfg["type"] == "composed":
         axes = _parse_composed_axes(router_cfg)
@@ -1030,9 +957,7 @@ def _build_router_from_cfg(
         router.gumbel = bool(router_cfg.get("gumbel", False))
         return router
     _check_router_conditioning_compat([router_cfg["type"]], conditioning)
-    router_kwargs = {
-        k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")
-    }
+    router_kwargs = {k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")}
     router_kwargs.setdefault("pdg_vocab", pdg_vocab)
     router_kwargs.setdefault("mat_vocab", mat_vocab)
     router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
@@ -1042,15 +967,9 @@ def _build_router_from_cfg(
 
 def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
     if model_config.get("mode") == "wgan":
-        stage1 = WGANGenerator(
-            **{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS}
-        )
+        stage1 = WGANGenerator(**{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS})
         sec_decoder = WGANSecondaryGenerator(
-            **{
-                k: v
-                for k, v in model_config.items()
-                if k in _WGAN_SEC_GENERATOR_MODEL_KEYS
-            }
+            **{k: v for k, v in model_config.items() if k in _WGAN_SEC_GENERATOR_MODEL_KEYS}
         )
         return stage1, sec_decoder
 
@@ -1061,44 +980,30 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
         shared = dict(
             pdg_vocab=pdg_vocab,
             mat_vocab=mat_vocab,
-            expert_hidden_dim=model_config.get("expert_hidden_dim")
-            or model_config.get("hidden_dim", 128),
-            expert_n_blocks=model_config.get("expert_n_blocks")
-            or model_config.get("n_blocks", 3),
+            expert_hidden_dim=model_config.get("expert_hidden_dim") or model_config.get("hidden_dim", 128),
+            expert_n_blocks=model_config.get("expert_n_blocks") or model_config.get("n_blocks", 3),
             emb_dim=model_config.get("emb_dim", EMB_DIM),
             dropout=model_config.get("dropout", 0.1),
             conditioning=model_config.get("conditioning", "embedding"),
         )
         conditioning = shared["conditioning"]
         stage1 = RoutedDenoisingMLP(
-            router=_build_router_from_cfg(
-                router_cfg, pdg_vocab, mat_vocab, conditioning
-            ),
+            router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, conditioning),
             k_max=model_config.get("k_max", K_MAX),
             **shared,
         )
         sec_decoder = RoutedSecondaryDecoder(
-            router=_build_router_from_cfg(
-                router_cfg, pdg_vocab, mat_vocab, conditioning
-            ),
+            router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, conditioning),
             **shared,
         )
         return stage1, sec_decoder
 
-    stage1 = DenoisingMLP(
-        **{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS}
-    )
-    sec_decoder = SecondaryDecoder(
-        **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
-    )
+    stage1 = DenoisingMLP(**{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS})
+    sec_decoder = SecondaryDecoder(**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS})
     return stage1, sec_decoder
 
 
 def build_critics(model_config: dict) -> tuple[nn.Module, nn.Module]:
-    critic = Critic(
-        **{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS}
-    )
-    sec_critic = SecondaryCritic(
-        **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
-    )
+    critic = Critic(**{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS})
+    sec_critic = SecondaryCritic(**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS})
     return critic, sec_critic
diff --git a/tests/test_bump_dataset_version.py b/tests/test_bump_dataset_version.py
index 3f169a5..83f7dd9 100644
--- a/tests/test_bump_dataset_version.py
+++ b/tests/test_bump_dataset_version.py
@@ -22,9 +22,7 @@ def test_git_user_name_returns_none_on_timeout(monkeypatch):
 
 
 def test_bump_gen_starts_at_gen1_when_none_exist(tmp_path):
-    dirs, log_line = plan_bump_gen(
-        tmp_path, "steps", "first generation", None, "2026-01-01"
-    )
+    dirs, log_line = plan_bump_gen(tmp_path, "steps", "first generation", None, "2026-01-01")
     assert dirs == [
         tmp_path / "raw" / "steps" / "gen1",
         tmp_path / "processed" / "steps" / "gen1" / "schema1",
@@ -56,9 +54,7 @@ def test_bump_gen_kinds_are_independent(tmp_path):
 
 def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path):
     (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
-    dirs, log_line = plan_bump_schema(
-        tmp_path, "steps", "gen1", "added e_sec column", None, "2026-01-01"
-    )
+    dirs, log_line = plan_bump_schema(tmp_path, "steps", "gen1", "added e_sec column", None, "2026-01-01")
     assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema1"]
     assert "`gen1`/`schema1`" in log_line
 
@@ -66,18 +62,14 @@ def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path):
 def test_bump_schema_increments_within_its_gen(tmp_path):
     (tmp_path / "processed" / "steps" / "gen1" / "schema1").mkdir(parents=True)
     (tmp_path / "processed" / "steps" / "gen1" / "schema2").mkdir(parents=True)
-    dirs, _ = plan_bump_schema(
-        tmp_path, "steps", "gen1", "next schema", None, "2026-01-01"
-    )
+    dirs, _ = plan_bump_schema(tmp_path, "steps", "gen1", "next schema", None, "2026-01-01")
     assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema3"]
 
 
 def test_bump_schema_does_not_see_other_gens_schemas(tmp_path):
     (tmp_path / "processed" / "steps" / "gen1" / "schema5").mkdir(parents=True)
     (tmp_path / "raw" / "steps" / "gen2").mkdir(parents=True)
-    dirs, _ = plan_bump_schema(
-        tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01"
-    )
+    dirs, _ = plan_bump_schema(tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01")
     assert dirs == [tmp_path / "processed" / "steps" / "gen2" / "schema1"]
 
 
@@ -91,9 +83,7 @@ def test_bump_schema_rejects_nonexistent_gen(tmp_path):
 
 def test_bump_gen_to_specific_tag(tmp_path):
     (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
-    dirs, log_line = plan_bump_gen(
-        tmp_path, "steps", "jump to gen5", None, "2026-01-01", target="gen5"
-    )
+    dirs, log_line = plan_bump_gen(tmp_path, "steps", "jump to gen5", None, "2026-01-01", target="gen5")
     assert dirs[0] == tmp_path / "raw" / "steps" / "gen5"
     assert "`gen5`" in log_line
 
@@ -125,9 +115,7 @@ def test_bump_schema_to_specific_tag(tmp_path):
 def test_bump_schema_rejects_invalid_to_tag(tmp_path):
     (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
     try:
-        plan_bump_schema(
-            tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3"
-        )
+        plan_bump_schema(tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3")
         assert False, "expected SystemExit"
     except SystemExit:
         pass
@@ -164,15 +152,7 @@ def _make_parquet(path):
 
 
 def test_update_manifest_bumps_to_specified_schema(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -194,15 +174,7 @@ def test_update_manifest_auto_detects_highest_schema(tmp_path):
     for schema in ("schema1", "schema2", "schema3"):
         d = tmp_path / "processed" / "steps" / "gen1" / schema / "pbwo4"
         d.mkdir(parents=True)
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema3"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
     parquet.touch()
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -231,15 +203,7 @@ def test_update_manifest_reports_missing_targets(tmp_path):
 
 
 def test_update_manifest_skips_already_at_target(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -254,15 +218,7 @@ def test_update_manifest_skips_already_at_target(tmp_path):
 
 
 def test_update_manifest_preserves_comments_and_blanks(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -278,15 +234,7 @@ def test_update_manifest_preserves_comments_and_blanks(tmp_path):
 
 
 def test_update_manifest_bumps_gen(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen2"
-        / "schema1"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen2" / "schema1" / "pbwo4" / "shard-000.parquet"
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -303,15 +251,7 @@ def test_update_manifest_bumps_gen(tmp_path):
 
 
 def test_update_manifest_bumps_gen_and_schema(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen2"
-        / "schema3"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen2" / "schema3" / "pbwo4" / "shard-000.parquet"
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -328,15 +268,7 @@ def test_update_manifest_bumps_gen_and_schema(tmp_path):
 
 
 def test_apply_update_manifest_writes_file(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -358,24 +290,8 @@ def test_apply_update_manifest_writes_file(tmp_path):
 
 
 def test_create_manifest_writes_relative_paths(tmp_path):
-    pq1 = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
-    pq2 = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-001.parquet"
-    )
+    pq1 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
+    pq2 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-001.parquet"
     _make_parquet(pq1)
     _make_parquet(pq2)
 
@@ -419,9 +335,7 @@ def test_run_create_manifest_refuses_to_overwrite_existing_output(tmp_path):
     output.write_text("original contents\n")
 
     try:
-        bump_dataset_version.run_create_manifest(
-            [str(pq)], execute=True, output=str(output)
-        )
+        bump_dataset_version.run_create_manifest([str(pq)], execute=True, output=str(output))
         assert False, "expected SystemExit"
     except SystemExit:
         pass
@@ -435,9 +349,7 @@ def test_run_create_manifest_force_overwrites_existing_output(tmp_path):
     output.parent.mkdir(parents=True)
     output.write_text("original contents\n")
 
-    bump_dataset_version.run_create_manifest(
-        [str(pq)], execute=True, output=str(output), force=True
-    )
+    bump_dataset_version.run_create_manifest([str(pq)], execute=True, output=str(output), force=True)
     assert output.read_text() != "original contents\n"
 
 
diff --git a/tests/test_catalog.py b/tests/test_catalog.py
index c8a47f0..92ddb24 100644
--- a/tests/test_catalog.py
+++ b/tests/test_catalog.py
@@ -13,9 +13,7 @@ from tests.test_analysis_reduce import _reference_frame, _rollout_frame
 
 def _build_ctx() -> Context:
     r, t = _rollout_frame(), _reference_frame()
-    return build_context(
-        r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
-    )
+    return build_context(r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000)
 
 
 @pytest.fixture(scope="module")
@@ -142,10 +140,7 @@ def test_chunked_matches_unchunked(ctx: Context, spec_id: str):
 
     # 4 chunks over only 2 distinct event_ids also exercises empty chunks.
     n_chunks = 4 if spec.chunkable else 1
-    parts = [
-        spec.compute_partial(Bundle.open(r, t, ctx, chunk=(k, n_chunks)))
-        for k in range(n_chunks)
-    ]
+    parts = [spec.compute_partial(Bundle.open(r, t, ctx, chunk=(k, n_chunks))) for k in range(n_chunks)]
     chunked = spec.finalize(parts, ctx)
 
     assert chunked.id == unchunked.id
diff --git a/tests/test_cli_new_run.py b/tests/test_cli_new_run.py
index c6ef14f..a4f9450 100644
--- a/tests/test_cli_new_run.py
+++ b/tests/test_cli_new_run.py
@@ -101,9 +101,7 @@ def test_force_guard_refuses_to_clobber_existing_checkpoints(tmp_path: Path):
     assert "already has last.pt" in result.output
     assert not (out_dir / "config.toml").exists()
 
-    result = runner.invoke(
-        app, ["new-run", "--out", str(out_dir), "--mode", "ddpm", "--force"]
-    )
+    result = runner.invoke(app, ["new-run", "--out", str(out_dir), "--mode", "ddpm", "--force"])
     assert result.exit_code == 0, result.output
     assert (out_dir / "config.toml").exists()
 
diff --git a/tests/test_cli_predict.py b/tests/test_cli_predict.py
index a6f74fb..1c18ae8 100644
--- a/tests/test_cli_predict.py
+++ b/tests/test_cli_predict.py
@@ -116,9 +116,7 @@ def test_ref_yaml_includes_comment_when_provided(tmp_path):
     dataset = tmp_path / "full.manifest"
     pred_uuid = str(uuid.uuid4())
 
-    ref_path = _write_prediction_ref(
-        checkpoint, pred_uuid, out, dataset, comment="baseline sweep run 3"
-    )
+    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset, comment="baseline sweep run 3")
     data = yaml.safe_load(ref_path.read_text())
 
     assert data["comment"] == "baseline sweep run 3"
@@ -133,9 +131,7 @@ def test_ref_timestamp_is_iso_format(tmp_path):
     checkpoint.touch()
 
     pred_uuid = str(uuid.uuid4())
-    ref_path = _write_prediction_ref(
-        checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d"
-    )
+    ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
     data = yaml.safe_load(ref_path.read_text())
 
     # Must parse without error and be timezone-aware (UTC).
@@ -150,9 +146,7 @@ def test_ref_checkpoint_path_is_absolute(tmp_path):
     checkpoint.touch()
 
     pred_uuid = str(uuid.uuid4())
-    ref_path = _write_prediction_ref(
-        checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d"
-    )
+    ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
     data = yaml.safe_load(ref_path.read_text())
 
     assert data["checkpoint"].startswith("/")
diff --git a/tests/test_condor.py b/tests/test_condor.py
index 4027a50..02912ed 100644
--- a/tests/test_condor.py
+++ b/tests/test_condor.py
@@ -62,9 +62,7 @@ def _fake_venv(repo_dir: Path) -> None:
     giant.chmod(0o755)
 
 
-def _prep(
-    rollout_yaml: Path, run_dir: str | Path | None = None, chunks: int = 1
-) -> Path:
+def _prep(rollout_yaml: Path, run_dir: str | Path | None = None, chunks: int = 1) -> Path:
     """``prep`` with small test-sized context bins/sampling."""
     return prep(
         rollout_yaml,
@@ -224,16 +222,12 @@ def test_write_submit_description(tmp_path: Path):
     assert "--chunk" in body and "--run-dir" in body
 
 
-def test_write_submit_requires_synced_venv(
-    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
-):
+def test_write_submit_requires_synced_venv(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
     run_dir = _prep(_write_inputs(tmp_path))
     cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path)
     # No `giant` next to the (fake) active interpreter, so this falls through
     # to repo_dir/.venv/bin/giant, which _write_inputs/_prep also didn't create.
-    monkeypatch.setattr(
-        sys, "executable", str(tmp_path / "not-a-venv" / "bin" / "python")
-    )
+    monkeypatch.setattr(sys, "executable", str(tmp_path / "not-a-venv" / "bin" / "python"))
     with pytest.raises(FileNotFoundError, match="uv sync"):
         write_submit(cfg)
 
@@ -241,9 +235,7 @@ def test_write_submit_requires_synced_venv(
 def test_write_submit_remote_flag(tmp_path: Path):
     run_dir = _prep(_write_inputs(tmp_path))
     _fake_venv(tmp_path)
-    cfg = SubmitConfig(
-        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, remote=True
-    )
+    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, remote=True)
     txt = write_submit(cfg).read_text()
     assert "+RemoteJob = True" in txt
     assert "ProvidesETPResources" not in txt
@@ -253,9 +245,7 @@ def test_write_submit_chunks_respect_chunkable(tmp_path: Path):
     assert get_spec("router_gating").chunkable is False
     run_dir = _prep(_write_inputs(tmp_path), chunks=4)
     _fake_venv(tmp_path)
-    cfg = SubmitConfig(
-        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4
-    )
+    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4)
     write_submit(cfg)
     jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
     counts: dict[str, int] = {}
@@ -272,9 +262,7 @@ def test_write_submit_rejects_n_chunks_mismatch_with_run_meta(tmp_path: Path):
     _job_walltimes instead of a clear error here."""
     run_dir = _prep(_write_inputs(tmp_path), chunks=2)
     _fake_venv(tmp_path)
-    cfg = SubmitConfig(
-        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4
-    )
+    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4)
     with pytest.raises(ValueError, match="n_chunks"):
         write_submit(cfg)
 
@@ -297,16 +285,9 @@ def test_write_submit_walltime_grows_with_chunk_rows(tmp_path: Path):
     run_dir = _prep(_write_inputs(tmp_path), chunks=2)
     meta = RunMeta.load(run_dir / "run_meta.json")
     _fake_venv(tmp_path)
-    cfg = SubmitConfig(
-        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=2
-    )
+    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=2)
     write_submit(cfg)
-    jobs = {
-        (i, int(k)): int(w)
-        for i, k, w in (
-            line.split(",") for line in (run_dir / "jobs.txt").read_text().split()
-        )
-    }
+    jobs = {(i, int(k)): int(w) for i, k, w in (line.split(",") for line in (run_dir / "jobs.txt").read_text().split())}
     for chunk in range(2):
         expected = estimate_runtime_s("marginal_edep", meta.rows_per_chunk[chunk])
         assert jobs[("marginal_edep", chunk)] == expected
diff --git a/tests/test_config.py b/tests/test_config.py
index fb7b5c2..6aea480 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -112,9 +112,7 @@ def test_migrate_config_lambda_nsec_and_lambda_s2():
 
 
 def test_migrate_config_wgan_knobs_map_to_both_stages():
-    new = gconfig.migrate_config(
-        {"train": {"n_critic": 3, "gp_weight": 5.0, "critic_lr": 1e-4}}
-    )
+    new = gconfig.migrate_config({"train": {"n_critic": 3, "gp_weight": 5.0, "critic_lr": 1e-4}})
     for stage in ("stage1_model", "stage2_model"):
         assert new[stage]["wgan"]["n_critic"] == 3
         assert new[stage]["wgan"]["gp_weight"] == 5.0
@@ -122,9 +120,7 @@ def test_migrate_config_wgan_knobs_map_to_both_stages():
 
 
 def test_migrate_config_model_hidden_dim_n_blocks_dropout_map_to_both_stages():
-    new = gconfig.migrate_config(
-        {"model": {"hidden_dim": 128, "n_blocks": 4, "dropout": 0.2}}
-    )
+    new = gconfig.migrate_config({"model": {"hidden_dim": 128, "n_blocks": 4, "dropout": 0.2}})
     for stage in ("stage1_model", "stage2_model"):
         assert new[stage]["hidden_dim"] == 128
         assert new[stage]["n_res_blocks"] == 4
@@ -132,9 +128,7 @@ def test_migrate_config_model_hidden_dim_n_blocks_dropout_map_to_both_stages():
 
 
 def test_migrate_config_emb_dim_and_conditioning_map_to_both_axes():
-    new = gconfig.migrate_config(
-        {"model": {"emb_dim": 32, "conditioning": "embedding"}}
-    )
+    new = gconfig.migrate_config({"model": {"emb_dim": 32, "conditioning": "embedding"}})
     for axis in ("particle", "material"):
         assert new["conditioning"][axis]["emb_dim"] == 32
         assert new["conditioning"][axis]["type"] == "embedding"
@@ -181,11 +175,7 @@ def test_migrate_config_router_copied_to_both_stages_with_tie_to_stage1_false():
 
 
 def test_migrate_config_router_nonzero_expert_dims_raises():
-    cfg = {
-        "model": {
-            "router": {"enabled": True, "expert_hidden_dim": 128, "expert_n_blocks": 0}
-        }
-    }
+    cfg = {"model": {"router": {"enabled": True, "expert_hidden_dim": 128, "expert_n_blocks": 0}}}
     try:
         gconfig.migrate_config(cfg)
         assert False, "expected ValueError"
@@ -267,9 +257,7 @@ def test_merge_cli_overrides_nested_override_keeps_siblings():
     assert cfg["stage1_model"]["hidden_dim"] == 256  # untouched sibling section
 
 
-def test_merge_cli_overrides_file_then_explicit_override_precedence(
-    tmp_path, monkeypatch
-):
+def test_merge_cli_overrides_file_then_explicit_override_precedence(tmp_path, monkeypatch):
     monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123")
     path = tmp_path / "config.toml"
     _write_toml(
@@ -317,9 +305,7 @@ def test_merge_cli_overrides_warns_on_git_hash_mismatch(tmp_path, monkeypatch, c
     assert "current999" in captured.err
 
 
-def test_merge_cli_overrides_no_warning_on_matching_git_hash(
-    tmp_path, monkeypatch, capsys
-):
+def test_merge_cli_overrides_no_warning_on_matching_git_hash(tmp_path, monkeypatch, capsys):
     monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
     path = tmp_path / "config.toml"
     _write_toml(path, git_hash="same123", extra="[train]\nepochs = 5\n")
@@ -328,9 +314,7 @@ def test_merge_cli_overrides_no_warning_on_matching_git_hash(
     assert capsys.readouterr().err == ""
 
 
-def test_merge_cli_overrides_no_warning_when_git_hash_unknown(
-    tmp_path, monkeypatch, capsys
-):
+def test_merge_cli_overrides_no_warning_when_git_hash_unknown(tmp_path, monkeypatch, capsys):
     monkeypatch.setattr(gconfig, "git_hash", lambda: "unknown")
     path = tmp_path / "config.toml"
     _write_toml(path, git_hash="abc123", extra="[train]\nepochs = 5\n")
@@ -339,9 +323,7 @@ def test_merge_cli_overrides_no_warning_when_git_hash_unknown(
     assert capsys.readouterr().err == ""
 
 
-def test_merge_cli_overrides_no_warning_when_meta_section_absent(
-    tmp_path, monkeypatch, capsys
-):
+def test_merge_cli_overrides_no_warning_when_meta_section_absent(tmp_path, monkeypatch, capsys):
     monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
     path = tmp_path / "config.toml"
     path.write_text("[train]\nepochs = 5\n")
@@ -351,12 +333,8 @@ def test_merge_cli_overrides_no_warning_when_meta_section_absent(
 
 
 def test_merge_cli_overrides_real_default_toml_fixture(monkeypatch):
-    monkeypatch.setattr(
-        gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243"
-    )
-    cfg = gconfig.merge_cli_overrides(
-        gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / "default.toml", {}
-    )
+    monkeypatch.setattr(gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243")
+    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / "default.toml", {})
     assert cfg["stage1_model"]["generator"] == "flow"
     assert cfg["stage1_model"]["hidden_dim"] == 256
     assert cfg["stage2_model"]["hidden_dim"] == 256
@@ -431,10 +409,7 @@ def _cfg_with(**dotted_overrides):
 
 
 def test_default_out_dir_name_all_defaults_is_just_the_timestamp():
-    assert (
-        gconfig.default_out_dir_name(gconfig.DEFAULT_CONFIG, now=_NOW)
-        == "20260729_1430"
-    )
+    assert gconfig.default_out_dir_name(gconfig.DEFAULT_CONFIG, now=_NOW) == "20260729_1430"
 
 
 def test_default_out_dir_name_stage1_generator_shown_bare_no_prefix():
@@ -499,9 +474,7 @@ def test_default_out_dir_name_router_gumbel_shown_when_enabled():
             "stage1_model.router.gumbel": True,
         }
     )
-    assert (
-        gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8_s1gum"
-    )
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8_s1gum"
 
 
 def test_default_out_dir_name_overflow_caps_and_hashes_remainder():
@@ -524,9 +497,7 @@ def test_default_out_dir_name_overflow_caps_and_hashes_remainder():
     # First 6 by priority: stage1_generator, stage2_generator, stage2_decoder,
     # stage2_history, particle_type_target, stage1_router — stage2_router
     # overflows into the hash suffix.
-    assert name.startswith(
-        "20260729_1430_wgan_s2-flow_dec-one_shot_hist-attention_pt-physical_s1r-energy8_+"
-    )
+    assert name.startswith("20260729_1430_wgan_s2-flow_dec-one_shot_hist-attention_pt-physical_s1r-energy8_+")
 
 
 def test_default_out_dir_name_overflow_hash_is_deterministic_and_value_sensitive():
@@ -758,15 +729,11 @@ def test_validate_config_ar_checks_skipped_under_one_shot():
 # ---------------------------------------------------------------------------
 
 
-def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(
-    tmp_path, monkeypatch, capsys
-):
+def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(tmp_path, monkeypatch, capsys):
     monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
     ckpt_path = tmp_path / "best.pt"
     ckpt_path.write_bytes(b"")
-    _write_toml(
-        tmp_path / "config.toml", git_hash="old111", extra="[train]\nepochs = 5\n"
-    )
+    _write_toml(tmp_path / "config.toml", git_hash="old111", extra="[train]\nepochs = 5\n")
 
     gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
 
@@ -776,9 +743,7 @@ def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(
     assert "current999" in captured.err
 
 
-def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(
-    tmp_path, monkeypatch, capsys
-):
+def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(tmp_path, monkeypatch, capsys):
     monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
     ckpt_path = tmp_path / "best.pt"
     ckpt_path.write_bytes(b"")
@@ -787,15 +752,11 @@ def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(
     assert capsys.readouterr().err == ""
 
 
-def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(
-    tmp_path, monkeypatch, capsys
-):
+def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(tmp_path, monkeypatch, capsys):
     monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
     ckpt_path = tmp_path / "best.pt"
     ckpt_path.write_bytes(b"")
-    _write_toml(
-        tmp_path / "config.toml", git_hash="same123", extra="[train]\nepochs = 5\n"
-    )
+    _write_toml(tmp_path / "config.toml", git_hash="same123", extra="[train]\nepochs = 5\n")
 
     gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
     assert capsys.readouterr().err == ""
diff --git a/tests/test_create_root_files.py b/tests/test_create_root_files.py
index f07453c..686e198 100644
--- a/tests/test_create_root_files.py
+++ b/tests/test_create_root_files.py
@@ -16,9 +16,7 @@ SimJob = create_root_files.SimJob
 PlanError = create_root_files.PlanError
 
 
-def _write_fake_executable(
-    path: Path, *, output_count: int = 1, exit_code: int = 0, sleep: float = 0.0
-) -> Path:
+def _write_fake_executable(path: Path, *, output_count: int = 1, exit_code: int = 0, sleep: float = 0.0) -> Path:
     """Stand-in for run_pbwo4/run_sampling: writes *output_count* .root files
     into its own cwd (so callers can verify each job gets an isolated workdir
     and that the workdir ends up holding *only* the .root output, matching
@@ -89,17 +87,13 @@ def test_next_shard_index_continues_past_existing(tmp_path):
 
 def test_plan_jobs_rejects_missing_gen(tmp_path):
     with pytest.raises(PlanError):
-        plan_jobs(
-            ["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1"
-        )
+        plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1")
 
 
 def test_plan_jobs_rejects_malformed_gen(tmp_path):
     (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
     with pytest.raises(PlanError):
-        plan_jobs(
-            ["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen"
-        )
+        plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen")
 
 
 def test_plan_jobs_continues_from_existing_shards(tmp_path):
@@ -108,9 +102,7 @@ def test_plan_jobs_continues_from_existing_shards(tmp_path):
     (gen_dir / "pbwo4" / "shard-000.root").touch()
     (gen_dir / "pbwo4" / "shard-001.root").touch()
 
-    jobs = plan_jobs(
-        ["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1"
-    )
+    jobs = plan_jobs(["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1")
 
     assert [j.shard_index for j in jobs] == [2, 3, 4]
     assert all(j.detector == "pbwo4" and j.config is None for j in jobs)
@@ -320,9 +312,7 @@ def test_run_all_caps_concurrency(tmp_path):
     assert {d.name for d in dests} == {f"shard-{i:03d}.root" for i in range(6)}
 
     intervals = [json.loads(d.read_text()) for d in dests]
-    events = sorted(
-        [(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals]
-    )
+    events = sorted([(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals])
     concurrent = 0
     peak = 0
     for _, delta in events:
diff --git a/tests/test_dwarf.py b/tests/test_dwarf.py
index 767828e..2d9aeae 100644
--- a/tests/test_dwarf.py
+++ b/tests/test_dwarf.py
@@ -51,9 +51,7 @@ def test_convert_rejects_output_with_multiple_files(tmp_path):
 def test_convert_rejects_output_with_parallel_jobs(tmp_path):
     root_file = tmp_path / "shard.root"
     root_file.touch()
-    result = runner.invoke(
-        app, ["convert", str(root_file), "--output", "out.parquet", "--jobs", "2"]
-    )
+    result = runner.invoke(app, ["convert", str(root_file), "--output", "out.parquet", "--jobs", "2"])
     assert result.exit_code != 0
     assert "--output cannot be combined with --jobs > 1" in result.output
 
diff --git a/tests/test_geometry.py b/tests/test_geometry.py
index 9e9aace..aa67110 100644
--- a/tests/test_geometry.py
+++ b/tests/test_geometry.py
@@ -44,9 +44,7 @@ def test_iter_point_batches_without_post_columns_yields_pre_only(tmp_path):
     (pos, mat, lay) = next(g._iter_point_batches(path))
 
     assert pos.shape == (5, 3)
-    np.testing.assert_allclose(
-        pos, df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32)
-    )
+    np.testing.assert_allclose(pos, df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32))
     assert list(mat) == ["G4_AIR"] * 5
     np.testing.assert_array_equal(lay, np.arange(5))
 
@@ -63,12 +61,8 @@ def test_iter_point_batches_with_post_columns_doubles_and_concatenates_points(
     # Every step contributes both its pre_pos and post_pos, sharing the
     # step's material/layer_id label — so batches double in length.
     assert pos.shape == (10, 3)
-    np.testing.assert_allclose(
-        pos[:5], df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32)
-    )
-    np.testing.assert_allclose(
-        pos[5:], df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32)
-    )
+    np.testing.assert_allclose(pos[:5], df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32))
+    np.testing.assert_allclose(pos[5:], df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32))
     assert list(mat) == ["G4_AIR"] * 10
     np.testing.assert_array_equal(lay, np.concatenate([np.arange(5), np.arange(5)]))
 
@@ -208,9 +202,7 @@ def test_slab_classes_discovered():
 
 def test_slab_query_labels_by_depth():
     orc = _build_slab()
-    pos = np.array(
-        [[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0]]
-    )  # layer 0, gap, layer 1
+    pos = np.array([[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0]])  # layer 0, gap, layer 1
     material, layer_id, escaped = orc.query(pos)
     assert list(material) == ["G4_PbWO4", "G4_AIR", "G4_W"]
     assert list(layer_id) == [0, -1, 1]
@@ -239,9 +231,7 @@ def test_slab_save_load_roundtrip(tmp_path):
     orc.save(p)
     loaded = g.GeometryOracle.load(p)
 
-    pos = np.array(
-        [[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0], [0.0, 0.0, 1e5]]
-    )
+    pos = np.array([[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0], [0.0, 0.0, 1e5]])
     m0, l0, e0 = orc.query(pos)
     m1, l1, e1 = loaded.query(pos)
     assert (m0 == m1).all() and (l0 == l1).all() and (e0 == e1).all()
diff --git a/tests/test_loader.py b/tests/test_loader.py
index ad31f7a..5d57656 100644
--- a/tests/test_loader.py
+++ b/tests/test_loader.py
@@ -315,9 +315,7 @@ def test_build_index_maps_from_files_ordering_independent_of_file_order(tmp_path
 
 def test_build_index_maps_from_files_numeric_sort_for_nuclear_codes(tmp_path):
     path = tmp_path / "a.parquet"
-    pd.DataFrame({"pdg": [22, 1000060120, 11], "material": ["X", "X", "X"]}).to_parquet(
-        path
-    )
+    pd.DataFrame({"pdg": [22, 1000060120, 11], "material": ["X", "X", "X"]}).to_parquet(path)
 
     pdg_map, _ = build_index_maps_from_files([path])
     assert list(pdg_map.keys()) == [11, 22, 1000060120]
@@ -398,9 +396,7 @@ def test_load_event_ids_applies_offset(tmp_path):
     path = tmp_path / "a.parquet"
     pd.DataFrame({"event_id": [0, 1, 2]}).to_parquet(path)
     offset = event_id_offset(1)
-    np.testing.assert_array_equal(
-        load_event_ids(path, offset=offset), [offset, offset + 1, offset + 2]
-    )
+    np.testing.assert_array_equal(load_event_ids(path, offset=offset), [offset, offset + 1, offset + 2])
 
 
 def test_load_event_ids_raises_when_event_id_reaches_stride(tmp_path):
diff --git a/tests/test_materials.py b/tests/test_materials.py
index c0d0c8e..233bf93 100644
--- a/tests/test_materials.py
+++ b/tests/test_materials.py
@@ -23,11 +23,7 @@ def test_get_material_properties_unfilled_entry_raises():
 
 
 def test_get_material_properties_returns_filled_entry_from_injected_table():
-    table = {
-        "G4_Pb": MaterialProperties(
-            z_eff=82.0, a_eff=207.2, density=11.35, x0=0.5612, lambda_int=17.59
-        )
-    }
+    table = {"G4_Pb": MaterialProperties(z_eff=82.0, a_eff=207.2, density=11.35, x0=0.5612, lambda_int=17.59)}
     props = get_material_properties("G4_Pb", table)
     assert props.z_eff == 82.0
     assert props.a_eff == 207.2
diff --git a/tests/test_migration_v02_v03.py b/tests/test_migration_v02_v03.py
index 05829a3..35f4db5 100644
--- a/tests/test_migration_v02_v03.py
+++ b/tests/test_migration_v02_v03.py
@@ -56,9 +56,7 @@ def _random_batch(seed: int):
 
 def _assert_bit_identical(a: torch.Tensor, b: torch.Tensor, label: str) -> None:
     assert a.shape == b.shape, f"{label}: shape mismatch {a.shape} vs {b.shape}"
-    assert torch.equal(a, b), (
-        f"{label}: outputs diverged, max abs diff = {(a - b).abs().max().item()}"
-    )
+    assert torch.equal(a, b), f"{label}: outputs diverged, max abs diff = {(a - b).abs().max().item()}"
 
 
 def _run_migration_check(mode: str, conditioning: str) -> None:
@@ -137,9 +135,7 @@ def _run_migration_check(mode: str, conditioning: str) -> None:
     assert new_stage1.n_sec_head is not None
     assert new_stage2.n_sec_head is None
 
-    remapped1, remapped2 = net.migrate_legacy_state_dict(
-        old_stage1.state_dict(), old_stage2.state_dict()
-    )
+    remapped1, remapped2 = net.migrate_legacy_state_dict(old_stage1.state_dict(), old_stage2.state_dict())
     missing1, unexpected1 = new_stage1.load_state_dict(remapped1, strict=True)
     missing2, unexpected2 = new_stage2.load_state_dict(remapped2, strict=True)
     assert not missing1 and not unexpected1
@@ -159,9 +155,7 @@ def _run_migration_check(mode: str, conditioning: str) -> None:
             new_out2 = new_stage2(x2, cond_cont, cond_cat, new_out1, t=t)
 
     _assert_bit_identical(old_out1, new_out1, f"stage1 output ({mode}, {conditioning})")
-    _assert_bit_identical(
-        old_n_sec, new_n_sec, f"n_sec logits ({mode}, {conditioning})"
-    )
+    _assert_bit_identical(old_n_sec, new_n_sec, f"n_sec logits ({mode}, {conditioning})")
     _assert_bit_identical(old_out2, new_out2, f"stage2 output ({mode}, {conditioning})")
 
 
diff --git a/tests/test_network.py b/tests/test_network.py
index bbf0552..5deab01 100644
--- a/tests/test_network.py
+++ b/tests/test_network.py
@@ -85,9 +85,7 @@ def test_stage1_model_gradients_flow():
 def test_stage1_model_no_n_sec_head_by_default():
     """Fresh v0.3.0 construction (no n_sec_head_k_max) has no n_sec head —
     it moves to stage 2."""
-    model = Stage1Model(
-        pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG
-    )
+    model = Stage1Model(pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG)
     assert model.n_sec_head is None
 
 
@@ -121,29 +119,18 @@ def test_stage2_type_dim_onehot_and_embedding_are_emb_dim():
 
 def test_stage2_trunk_sec_dim_physical_matches_v02_sec_dim():
     k_max = 15
-    assert (
-        stage2_trunk_sec_dim({"target": "physical"}, "flow", k_max, emb_dim=16)
-        == k_max * SEC_SLOT_DIM
-    )
-    assert (
-        stage2_trunk_sec_dim({"target": "physical"}, "wgan", k_max, emb_dim=16)
-        == k_max * SEC_SLOT_DIM
-    )
+    assert stage2_trunk_sec_dim({"target": "physical"}, "flow", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
+    assert stage2_trunk_sec_dim({"target": "physical"}, "wgan", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
 
 
 def test_stage2_trunk_sec_dim_onehot_wgan_folds_type_in():
     k_max = 15
-    assert stage2_trunk_sec_dim(
-        {"target": "onehot"}, "wgan", k_max, emb_dim=16
-    ) == k_max * (CONT_SLOT_DIM + 16)
+    assert stage2_trunk_sec_dim({"target": "onehot"}, "wgan", k_max, emb_dim=16) == k_max * (CONT_SLOT_DIM + 16)
 
 
 def test_stage2_trunk_sec_dim_onehot_flow_excludes_type():
     k_max = 15
-    assert (
-        stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16)
-        == k_max * CONT_SLOT_DIM
-    )
+    assert stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16) == k_max * CONT_SLOT_DIM
 
 
 # --- ConditionEncoder onehot mode -------------------------------------------
@@ -440,16 +427,12 @@ def test_stage2_autoregressive_history_invalid_raises():
 @pytest.mark.parametrize("history", ["markov", "attention"])
 def test_stage2_autoregressive_forward_shape(target, generator, history):
     B, K, emb_dim = 4, 5, 6
-    model = _build_stage2_ar(
-        target, generator, emb_dim=emb_dim, k_max=K, history=history
-    )
+    model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K, history=history)
     cond_cont = torch.randn(B, COND_DIM)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
     type_dim = stage2_type_dim({"target": target}, emb_dim)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
-        B, K, CONT_SLOT_DIM + type_dim
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
     token_dim = stage2_trunk_sec_dim({"target": target}, generator, 1, emb_dim)
     if generator == "wgan":
         x_t = torch.randn(B, K, model.noise_dim)
@@ -488,9 +471,7 @@ def test_stage2_autoregressive_predict_type_shape():
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
     type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
-        B, K, CONT_SLOT_DIM + type_dim
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
     out = model.predict_type(
         cond_cont,
         cond_cat,
@@ -511,9 +492,7 @@ def test_stage2_autoregressive_predict_type_raises_when_no_type_head(target, gen
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
     type_dim = stage2_type_dim({"target": target}, emb_dim)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
-        B, K, CONT_SLOT_DIM + type_dim
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
     with pytest.raises(RuntimeError):
         model.predict_type(
             cond_cont,
@@ -533,9 +512,7 @@ def test_stage2_autoregressive_gradients_flow_wgan_onehot():
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
     type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
-        B, K, CONT_SLOT_DIM + type_dim
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
     z = torch.randn(B, K, model.noise_dim)
     gen_out = model(
         z,
@@ -560,9 +537,7 @@ def test_stage2_autoregressive_gradients_flow_onehot():
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
     type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
-        B, K, CONT_SLOT_DIM + type_dim
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
     token_dim = stage2_trunk_sec_dim({"target": "onehot"}, "flow", 1, emb_dim)
     x_t = torch.randn(B, K, token_dim)
     t = torch.rand(B, K)
@@ -601,17 +576,13 @@ def test_stage2_autoregressive_history_step_matches_parallel_history_encoder():
     itself rather than `AttentionHistory` in isolation
     (`test_attention_history_step_matches_forward` covers that lower layer)."""
     B, K, emb_dim = 3, 6, 6
-    model = _build_stage2_ar(
-        "physical", "wgan", emb_dim=emb_dim, k_max=K, history="attention"
-    )
+    model = _build_stage2_ar("physical", "wgan", emb_dim=emb_dim, k_max=K, history="attention")
     model.eval()
     type_dim = stage2_type_dim({"target": "physical"}, emb_dim)
     hist_in_dim = CONT_SLOT_DIM + type_dim
     own_feat = torch.randn(B, K, hist_in_dim)  # token i's own raw feature
     has_prev_full = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
-    history_feat = torch.cat(
-        [torch.zeros_like(own_feat[:, :1]), own_feat[:, :-1]], dim=1
-    )
+    history_feat = torch.cat([torch.zeros_like(own_feat[:, :1]), own_feat[:, :-1]], dim=1)
 
     with torch.no_grad():
         expected = model.history_encoder(history_feat, has_prev_full)
diff --git a/tests/test_particles.py b/tests/test_particles.py
index bcfbfdf..c80e724 100644
--- a/tests/test_particles.py
+++ b/tests/test_particles.py
@@ -49,9 +49,7 @@ def test_ground_state_nucleus_resolved_via_particle_package():
     """He-4 (Z=2, A=4) is a common nuclide in `particle`'s ground-state table."""
     mass, charge = particle_mass_charge(1000020040)
     assert charge == pytest.approx(2.0)
-    assert mass == pytest.approx(
-        4 * 931.494, rel=0.05
-    )  # near A*amu, binding-energy-corrected
+    assert mass == pytest.approx(4 * 931.494, rel=0.05)  # near A*amu, binding-energy-corrected
 
 
 def test_nuclear_isomer_falls_back_to_z_a_decode():
@@ -174,9 +172,7 @@ def test_decode_topn_class_other_drop_returns_zero_sentinel():
 def test_decode_topn_class_other_sample_stays_within_members():
     topn_map, n_classes = _topn_fixture()
     rng = np.random.default_rng(0)
-    out = decode_topn_class(
-        np.full(50, 3), topn_map, n_classes, other_policy="sample", rng=rng
-    )
+    out = decode_topn_class(np.full(50, 3), topn_map, n_classes, other_policy="sample", rng=rng)
     assert set(out.tolist()) <= {2212, 2112}
 
 
diff --git a/tests/test_phase2.py b/tests/test_phase2.py
index afb12bd..7a8efec 100644
--- a/tests/test_phase2.py
+++ b/tests/test_phase2.py
@@ -57,9 +57,7 @@ def _sec_decoder(pdg=3, mat=2, conditioning="embedding"):
 
 def _cond(B=8, pdg=3, mat=2):
     cond_cont = torch.randn(B, COND_DIM)
-    cond_cat = torch.stack(
-        [torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
-    )
+    cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
     return cond_cont, cond_cat
 
 
@@ -154,9 +152,7 @@ def test_flow_matching_loss_secondary_scalar():
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
     sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
-    loss = flow_matching_loss_secondary(
-        decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
-    )
+    loss = flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask)
     assert loss.shape == ()
     assert loss.item() >= 0.0
 
@@ -169,9 +165,7 @@ def test_flow_matching_loss_secondary_mask_zeros_padding():
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
     sec_mask = torch.zeros(B, K_MAX, dtype=torch.bool)
-    loss = flow_matching_loss_secondary(
-        decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
-    )
+    loss = flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask)
     assert loss.item() == pytest.approx(0.0, abs=1e-6)
 
 
@@ -182,9 +176,7 @@ def test_flow_matching_loss_secondary_has_grad():
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
     sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
-    flow_matching_loss_secondary(
-        decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
-    ).backward()
+    flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask).backward()
     assert any(p.grad is not None for p in decoder.parameters())
 
 
@@ -220,9 +212,7 @@ def test_flow_matching_loss_secondary_ar_scalar():
     x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(
-        B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     sec_mask = torch.ones(B, K, dtype=torch.bool)
     loss = flow_matching_loss_secondary_ar(
         decoder,
@@ -246,9 +236,7 @@ def test_flow_matching_loss_secondary_ar_mask_zeros_padding():
     x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(
-        B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     sec_mask = torch.zeros(B, K, dtype=torch.bool)
     loss = flow_matching_loss_secondary_ar(
         decoder,
@@ -271,9 +259,7 @@ def test_flow_matching_loss_secondary_ar_has_grad():
     x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(
-        B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     sec_mask = torch.ones(B, K, dtype=torch.bool)
     flow_matching_loss_secondary_ar(
         decoder,
@@ -299,9 +285,7 @@ def test_sample_secondaries_shapes():
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
-    sec_cont, sec_phys, sec_valid = sample_secondaries(
-        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=3
-    )
+    sec_cont, sec_phys, sec_valid = sample_secondaries(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=3)
     assert sec_cont.shape == (B, K_MAX, 4)
     assert sec_phys.shape == (B, K_MAX, PARTICLE_PHYS_DIM)
     assert sec_valid.shape == (B, K_MAX)
@@ -314,9 +298,7 @@ def test_sample_secondaries_valid_mask_matches_n_sec():
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.tensor([0, 1, 3, K_MAX])
-    _, _, sec_valid = sample_secondaries(
-        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
-    )
+    _, _, sec_valid = sample_secondaries(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
     for i, n in enumerate(n_sec_pred.tolist()):
         assert sec_valid[i, :n].all()
         assert not sec_valid[i, n:].any()
@@ -350,9 +332,7 @@ def test_encode_secondaries_energy_conservation():
     pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
     pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
 
-    sec_cont = encode_secondaries(
-        sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
-    )
+    sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list)
     assert sec_cont.shape == (N, K_MAX, 6)
     assert np.isfinite(sec_cont).all()
 
@@ -399,9 +379,7 @@ def test_encode_secondaries_stick_logits_match_naive_reference():
             logit = np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP)
             expected[row, i] = logit
 
-    np.testing.assert_allclose(
-        stick_logits[sec_valid], expected[sec_valid], rtol=1e-4, atol=1e-4
-    )
+    np.testing.assert_allclose(stick_logits[sec_valid], expected[sec_valid], rtol=1e-4, atol=1e-4)
 
 
 def test_encode_secondaries_direction_encoding():
@@ -513,9 +491,7 @@ def test_encode_secondaries_physical_columns_match_ground_truth_pdg():
     sec_valid[0, 0] = True
     pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
 
-    sec_cont = encode_secondaries(
-        sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
-    )
+    sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list)
     mass, charge = particle_mass_charge(11)
     assert sec_cont[0, 0, 4] == pytest.approx(log_transform(np.array([mass]))[0])
     assert sec_cont[0, 0, 5] == pytest.approx(charge)
@@ -549,9 +525,7 @@ def test_decode_secondaries_valid_slots_sum_to_e_sec():
     e_sec = rng.uniform(0.0, 50.0, size=N).astype(np.float32)
     pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
 
-    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
-        sec_cont, n_sec, e_sec, pre_dir
-    )
+    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
 
     valid_sum = (sec_E * sec_valid).sum(axis=1)
     has_secondaries = n_sec > 0
@@ -574,9 +548,7 @@ def test_decode_secondaries_zero_n_sec_has_zero_energy():
     e_sec = rng.uniform(1.0, 10.0, size=N).astype(np.float32)
     pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
 
-    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
-        sec_cont, n_sec, e_sec, pre_dir
-    )
+    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
 
     assert not sec_valid.any()
     np.testing.assert_allclose(sec_E, 0.0)
@@ -596,9 +568,7 @@ def test_decode_secondaries_degenerate_row_falls_back_to_even_split():
     e_sec = np.array([0.0, 4.0, 9.0, 30.0], dtype=np.float32)
     pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
 
-    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
-        sec_cont, n_sec, e_sec, pre_dir
-    )
+    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
 
     for i, k in enumerate(n_sec):
         if k == 0:
@@ -622,12 +592,8 @@ def test_decode_secondaries_rescale_preserves_relative_shares():
     n_sec = np.array([4])
     pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
 
-    sec_E_small, _, _, _, sec_valid = decode_secondaries(
-        sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir
-    )
-    sec_E_large, _, _, _, _ = decode_secondaries(
-        sec_cont, n_sec, np.array([50.0], dtype=np.float32), pre_dir
-    )
+    sec_E_small, _, _, _, sec_valid = decode_secondaries(sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir)
+    sec_E_large, _, _, _, _ = decode_secondaries(sec_cont, n_sec, np.array([50.0], dtype=np.float32), pre_dir)
 
     ratio_small = sec_E_small[0, :4] / sec_E_small[0, 0]
     ratio_large = sec_E_large[0, :4] / sec_E_large[0, 0]
@@ -649,20 +615,14 @@ def test_decode_secondaries_mass_charge_round_trip_with_normalizer():
     sec_valid[0, 0] = True
     pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
 
-    sec_cont = encode_secondaries(
-        sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
-    )
+    sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list)
     norm = Normalizer()
     norm.mean = np.array([-2.0, 0.5], dtype=np.float32)
     norm.std = np.array([3.0, 1.5], dtype=np.float32)
     sec_cont_normed = sec_cont.copy()
-    sec_cont_normed[:, :, 4:6] = norm.transform(
-        sec_cont[:, :, 4:6].reshape(-1, 2)
-    ).reshape(N, K_MAX, 2)
+    sec_cont_normed[:, :, 4:6] = norm.transform(sec_cont[:, :, 4:6].reshape(-1, 2)).reshape(N, K_MAX, 2)
 
     n_sec = np.array([1])
-    _, _, sec_mass, sec_charge, _ = decode_secondaries(
-        sec_cont_normed, n_sec, e_sec, pre_dir, sec_phys_normalizer=norm
-    )
+    _, _, sec_mass, sec_charge, _ = decode_secondaries(sec_cont_normed, n_sec, e_sec, pre_dir, sec_phys_normalizer=norm)
     assert sec_mass[0, 0] == pytest.approx(938.27208943, abs=1e-2)
     assert sec_charge[0, 0] == pytest.approx(1.0, abs=1e-4)
diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
index c0848c9..9c25e67 100644
--- a/tests/test_pipeline.py
+++ b/tests/test_pipeline.py
@@ -48,9 +48,7 @@ def _make_synthetic_steps(path, n_events=20, seed=0):
             pre_dir = np.array([0.0, 0.0, 1.0])
             post_dir = _unit(rng.normal(size=3))
             post_pos = pre_pos + step_length * pre_dir
-            sec_energies = (
-                list(rng.dirichlet(np.ones(n_sec)) * e_sec) if n_sec > 0 else []
-            )
+            sec_energies = list(rng.dirichlet(np.ones(n_sec)) * e_sec) if n_sec > 0 else []
             sec_pdgs = [pdgs[(row_idx + j) % 2] for j in range(n_sec)]
             sec_dirs = [_unit(rng.normal(size=3)) for _ in range(n_sec)]
             rows.append(
@@ -211,9 +209,7 @@ def test_run_train_job_no_topn_map_for_physical_target(tmp_path, data):
 
 
 @pytest.mark.filterwarnings("ignore::DeprecationWarning:multiprocessing.popen_fork")
-def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(
-    tmp_path, data, monkeypatch
-):
+def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(tmp_path, data, monkeypatch):
     # num_workers>0 makes DataLoader actually fork worker subprocesses
     # (unlike every other test here, which runs with num_workers=0) — pytest
     # itself is multi-threaded, hence Python's fork-safety warning below.
@@ -223,9 +219,7 @@ def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(
 
 
 @pytest.mark.filterwarnings("ignore::DeprecationWarning:multiprocessing.popen_fork")
-def test_run_train_job_no_warning_when_num_workers_within_shared_quota(
-    tmp_path, data, monkeypatch
-):
+def test_run_train_job_no_warning_when_num_workers_within_shared_quota(tmp_path, data, monkeypatch):
     monkeypatch.setattr("giant.pipeline.os.cpu_count", lambda: 8)  # quota = 2
     echo = _run(data, tmp_path / "out", num_workers=2)
     assert not any("exceeds" in m for m in echo)
@@ -299,12 +293,8 @@ def test_run_train_job_mixed_particle_material_conditioning_end_to_end(tmp_path,
     from giant.constants import COND_DIM_BASE, PARTICLE_PHYS_DIM
 
     assert cond_norm.mean is not None and cond_norm.std is not None
-    np.testing.assert_allclose(
-        cond_norm.mean[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 0.0
-    )
-    np.testing.assert_allclose(
-        cond_norm.std[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 1.0
-    )
+    np.testing.assert_allclose(cond_norm.mean[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 0.0)
+    np.testing.assert_allclose(cond_norm.std[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 1.0)
     material_std = cond_norm.std[COND_DIM_BASE + PARTICLE_PHYS_DIM :]
     assert np.all(material_std > 0) and not np.allclose(material_std, 1.0)
 
@@ -342,12 +332,8 @@ def test_run_train_job_matches_uncached_output(tmp_path, data):
     cached = torch.load(tmp_path / "cached2" / "last.pt", weights_only=False)
 
     for key in ("cond", "target", "sec_phys"):
-        np.testing.assert_allclose(
-            uncached["normalizer"][key]["mean"], cached["normalizer"][key]["mean"]
-        )
-        np.testing.assert_allclose(
-            uncached["normalizer"][key]["std"], cached["normalizer"][key]["std"]
-        )
+        np.testing.assert_allclose(uncached["normalizer"][key]["mean"], cached["normalizer"][key]["mean"])
+        np.testing.assert_allclose(uncached["normalizer"][key]["std"], cached["normalizer"][key]["std"])
     assert uncached["pdg_map"] == cached["pdg_map"]
     assert uncached["mat_map"] == cached["mat_map"]
 
@@ -376,9 +362,7 @@ def test_seed_energy_router_falls_back_to_default_and_warns_when_no_samples():
     router_cfg = {"enabled": True, "type": "energy", "n_experts": 4}
     cond_norm = _fitted_cond_norm()
     echoed = []
-    _seed_energy_router(
-        router_cfg, cond_norm, np.empty(0), energy_idx=3, echo=echoed.append
-    )
+    _seed_energy_router(router_cfg, cond_norm, np.empty(0), energy_idx=3, echo=echoed.append)
     assert "centers_init" not in router_cfg
     assert len(echoed) == 1
     assert "falls back to default centers" in echoed[0]
@@ -393,9 +377,7 @@ def test_seed_energy_router_seeds_centers_from_data_quantiles():
     # units as the conditioning column being normalized against.
     energy_quantiles = np.linspace(1.0, 10.0, 33).astype(np.float32)
     echoed = []
-    _seed_energy_router(
-        router_cfg, cond_norm, energy_quantiles, energy_idx, echoed.append
-    )
+    _seed_energy_router(router_cfg, cond_norm, energy_quantiles, energy_idx, echoed.append)
 
     assert "centers_init" in router_cfg
     centers = np.asarray(router_cfg["centers_init"], dtype=np.float32)
diff --git a/tests/test_render.py b/tests/test_render.py
index f1ec9d5..064cb32 100644
--- a/tests/test_render.py
+++ b/tests/test_render.py
@@ -73,10 +73,7 @@ def test_render_router_diagnostics_and_edge_cases(tmp_path: Path):
             "x",
             {
                 "edges": [0, 1, 2],
-                "groups": {
-                    lbl: {"rollout": [1, 2], "reference": [2, 1]}
-                    for lbl in ("a", "b", "c", "d")
-                },
+                "groups": {lbl: {"rollout": [1, 2], "reference": [2, 1]} for lbl in ("a", "b", "c", "d")},
                 "log_y": True,
             },
         ),
@@ -126,9 +123,7 @@ def test_render_all_run_gallery_invokes_subprocess(tmp_path: Path, monkeypatch):
     for r in reduced:
         r.save(tmp_path / "reduced" / f"{r.id}.json")
     try:
-        render_mod.render_all(
-            tmp_path / "reduced", tmp_path / "plots", run_gallery=True
-        )
+        render_mod.render_all(tmp_path / "reduced", tmp_path / "plots", run_gallery=True)
     except RuntimeError as e:
         pytest.skip(f"LaTeX rendering unavailable: {e}")
 
@@ -145,9 +140,7 @@ def test_render_run_glues_condor_run_meta_into_render_all(tmp_path: Path, monkey
     (run_dir / "reduced").mkdir(parents=True)
 
     merge_calls = []
-    monkeypatch.setattr(
-        condor_mod, "merge_all", lambda rd: merge_calls.append(Path(rd))
-    )
+    monkeypatch.setattr(condor_mod, "merge_all", lambda rd: merge_calls.append(Path(rd)))
     meta = condor_mod.RunMeta(
         rollout="rollout.parquet",
         reference="reference.parquet",
@@ -157,9 +150,9 @@ def test_render_run_glues_condor_run_meta_into_render_all(tmp_path: Path, monkey
     )
     monkeypatch.setattr(condor_mod.RunMeta, "load", classmethod(lambda cls, p: meta))
 
-    Reduced(
-        "s", "species", "single_hist", "Single", "x", {"edges": [0, 1], "rollout": [1]}
-    ).save(run_dir / "reduced" / "s.json")
+    Reduced("s", "species", "single_hist", "Single", "x", {"edges": [0, 1], "rollout": [1]}).save(
+        run_dir / "reduced" / "s.json"
+    )
 
     try:
         pdfs = render_mod.render_run(run_dir)
@@ -363,9 +356,7 @@ def test_figure_params_old_shape_wgan_reports_noise_dim_not_steps():
 
 
 def test_plot_metadata_includes_note_and_run_meta_parameters():
-    r = Reduced(
-        "u", "router", "unavailable", "Unavailable", "x", {"note": "no router data"}
-    )
+    r = Reduced("u", "router", "unavailable", "Unavailable", "x", {"note": "no router data"})
     meta = render_mod._plot_metadata(r, {"title": "run-1", "checkpoint": "ckpt.pt"})
     assert meta["note"] == "no router data"
     assert meta["parameters"] == {"checkpoint": "ckpt.pt"}
diff --git a/tests/test_rollout.py b/tests/test_rollout.py
index da8d624..1cdb799 100644
--- a/tests/test_rollout.py
+++ b/tests/test_rollout.py
@@ -121,12 +121,8 @@ def fake_material_props(monkeypatch):
     import giant.materials as gm
 
     fake = {
-        "G4_AIR": gm.MaterialProperties(
-            z_eff=7.3, a_eff=14.4, density=1.2e-3, x0=3.0e4, lambda_int=7.0e5
-        ),
-        "G4_PbWO4": gm.MaterialProperties(
-            z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7
-        ),
+        "G4_AIR": gm.MaterialProperties(z_eff=7.3, a_eff=14.4, density=1.2e-3, x0=3.0e4, lambda_int=7.0e5),
+        "G4_PbWO4": gm.MaterialProperties(z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7),
     }
     monkeypatch.setattr(gm, "MATERIAL_PROPERTIES", fake)
     return fake
@@ -496,9 +492,7 @@ def test_resolve_n_sec_raises_when_neither_stage_owns_head(fake_material_props):
 
 @pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
 @pytest.mark.parametrize("generator2", ["flow", "wgan"])
-def test_rollout_physical_target_decoder_generator_matrix(
-    fake_material_props, decoder, generator2
-):
+def test_rollout_physical_target_decoder_generator_matrix(fake_material_props, decoder, generator2):
     """Every (decoder, stage2 generator) combination under
     particle_type.target="physical" must run to completion and conserve
     energy."""
@@ -542,9 +536,7 @@ def test_rollout_onehot_target_end_to_end(fake_material_props, decoder):
     assert len(rec["event_id"]) > 0
     # Every spawned secondary's nominal pdg must be one decode_topn_class can
     # actually produce (the topn map's known classes + its "other" members).
-    possible = set(PDG_TOPN_MAP.class_map.keys()) | set(
-        PDG_TOPN_MAP.other_members.keys()
-    )
+    possible = set(PDG_TOPN_MAP.class_map.keys()) | set(PDG_TOPN_MAP.other_members.keys())
     secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
     assert secondary_pdgs <= possible
 
@@ -589,9 +581,7 @@ def _onehot_conditioning_models():
     return s1.eval(), s2.eval()
 
 
-def _run_onehot_conditioning(
-    pdg_topn_map=COND_PDG_TOPN_MAP, mat_topn_map=COND_MAT_TOPN_MAP
-):
+def _run_onehot_conditioning(pdg_topn_map=COND_PDG_TOPN_MAP, mat_topn_map=COND_MAT_TOPN_MAP):
     s1, s2 = _onehot_conditioning_models()
     cond, tgt, sec_phys = _norms()
     return rollout(
@@ -643,9 +633,7 @@ def test_rollout_embedding_target_end_to_end(decoder):
     conditioning's own particle embedding table, so every resolved PDG must
     be a real member of the dense training vocab (pdg_map) — unlike
     "onehot", there is no "other" bucket to fall outside of."""
-    s1, s2 = _models_v3(
-        conditioning="embedding", decoder=decoder, target="embedding", emb_dim=4
-    )
+    s1, s2 = _models_v3(conditioning="embedding", decoder=decoder, target="embedding", emb_dim=4)
     rec = _run_v3(s1, s2, conditioning="embedding")
     assert len(rec["event_id"]) > 0
     secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
diff --git a/tests/test_router.py b/tests/test_router.py
index 01c0ded..d733218 100644
--- a/tests/test_router.py
+++ b/tests/test_router.py
@@ -25,9 +25,7 @@ MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
 
 def _cond(B=8, pdg=3, mat=2):
     cond_cont = torch.randn(B, COND_DIM)
-    cond_cat = torch.stack(
-        [torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
-    )
+    cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
     return cond_cont, cond_cat
 
 
@@ -78,9 +76,7 @@ def test_energy_router_gate_partition_of_unity():
 def test_energy_router_top1_matches_gate_argmax():
     router = EnergyRouter(n_experts=4)
     cond_cont, cond_cat = _cond(16)
-    assert torch.equal(
-        router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
-    )
+    assert torch.equal(router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1))
 
 
 def test_energy_router_hardens_as_temperature_shrinks():
@@ -128,12 +124,8 @@ def test_energy_router_centers_init_wrong_length_raises():
 
 
 def test_energy_router_centers_init_respects_learn_centers_flag():
-    learned = EnergyRouter(
-        n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=True
-    )
-    fixed = EnergyRouter(
-        n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=False
-    )
+    learned = EnergyRouter(n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=True)
+    fixed = EnergyRouter(n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=False)
     assert isinstance(learned.centers, torch.nn.Parameter)
     assert not isinstance(fixed.centers, torch.nn.Parameter)
 
@@ -160,9 +152,7 @@ def test_energy_router_learn_width_matches_fixed_temperature_at_init():
     per-expert width must reproduce the fixed-temperature gate exactly."""
     centers_init = [-1.0, 0.0, 0.5, 1.5]
     fixed = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init)
-    learned = EnergyRouter(
-        n_experts=4, temperature=0.3, centers_init=centers_init, learn_width=True
-    )
+    learned = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init, learn_width=True)
     cond_cont, cond_cat = _cond(16)
     torch.testing.assert_close(
         learned.gate(cond_cont, cond_cat),
@@ -195,21 +185,15 @@ def test_energy_router_learn_width_and_temperature_mutually_exclusive_raises():
         EnergyRouter(n_experts=4, learn_width=True, learn_temperature=True)
     except ValueError:
         return
-    raise AssertionError(
-        "expected ValueError for learn_width and learn_temperature both set"
-    )
+    raise AssertionError("expected ValueError for learn_width and learn_temperature both set")
 
 
 def test_energy_router_width_ratio_bounds_must_bracket_one_raises():
     try:
-        EnergyRouter(
-            n_experts=4, learn_width=True, width_min_ratio=1.0, width_max_ratio=2.0
-        )
+        EnergyRouter(n_experts=4, learn_width=True, width_min_ratio=1.0, width_max_ratio=2.0)
     except ValueError:
         return
-    raise AssertionError(
-        "expected ValueError for width_min_ratio/width_max_ratio not bracketing 1.0"
-    )
+    raise AssertionError("expected ValueError for width_min_ratio/width_max_ratio not bracketing 1.0")
 
 
 def test_energy_router_effective_width_stays_within_bounds():
@@ -246,9 +230,7 @@ def test_energy_router_learn_width_hardens_when_pushed_to_floor():
     """Pushing every expert's width toward the (tiny) floor should harden the
     gate to a one-hot at the nearest center, generalizing the fixed-
     temperature->0 hardening test to the per-expert path."""
-    router = EnergyRouter(
-        n_experts=4, learn_width=True, width_min_ratio=1e-4, width_max_ratio=10.0
-    )
+    router = EnergyRouter(n_experts=4, learn_width=True, width_min_ratio=1e-4, width_max_ratio=10.0)
     with torch.no_grad():
         router.raw_width.fill_(-1e6)
     cond_cont, cond_cat = _cond(16)
@@ -264,9 +246,7 @@ def test_energy_router_own_width_controls_own_coverage_independent_of_others():
     expert's own gate share, without needing to touch any other expert's
     width — the "each expert learns its own coverage independently" property
     this feature is meant to add."""
-    router = EnergyRouter(
-        n_experts=2, temperature=1.0, learn_width=True, centers_init=[0.0, 10.0]
-    )
+    router = EnergyRouter(n_experts=2, temperature=1.0, learn_width=True, centers_init=[0.0, 10.0])
     cond_cont, cond_cat = _cond(4)
     cond_cont[:, 3] = 3.0  # fixed energy, unequal distance to each center
 
@@ -280,9 +260,7 @@ def test_energy_router_own_width_controls_own_coverage_independent_of_others():
 
 
 def test_build_router_threads_learn_width_kwargs_through():
-    router = build_router(
-        "energy", 4, learn_width=True, width_min_ratio=0.2, width_max_ratio=8.0
-    )
+    router = build_router("energy", 4, learn_width=True, width_min_ratio=0.2, width_max_ratio=8.0)
     assert isinstance(router, EnergyRouter)
     assert router.learn_width is True
     assert isinstance(router.raw_width, torch.nn.Parameter)
@@ -419,9 +397,7 @@ def test_pdg_router_gate_partition_of_unity():
 def test_pdg_router_top1_matches_gate_argmax():
     router = PdgRouter(n_experts=4, pdg_vocab=3)
     cond_cont, cond_cat = _cond(16)
-    assert torch.equal(
-        router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
-    )
+    assert torch.equal(router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1))
 
 
 def test_pdg_router_hardens_as_temperature_shrinks():
@@ -586,9 +562,7 @@ def test_process_router_gate_partition_of_unity():
 def test_process_router_top1_matches_gate_argmax():
     router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
     cond_cont, cond_cat = _cond(16)
-    assert torch.equal(
-        router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
-    )
+    assert torch.equal(router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1))
 
 
 def test_process_router_balance_loss_is_nonnegative_scalar():
@@ -663,16 +637,12 @@ def test_build_models_routed_with_process_router():
 
 
 def test_composed_router_n_experts_is_product():
-    router = ComposedRouter(
-        [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
-    )
+    router = ComposedRouter([EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)])
     assert router.n_experts == 12
 
 
 def test_composed_router_gate_partition_of_unity():
-    router = ComposedRouter(
-        [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
-    )
+    router = ComposedRouter([EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)])
     cond_cont, cond_cat = _cond(16, pdg=5)
     g = router.gate(cond_cont, cond_cat)
     assert g.shape == (16, 12)
@@ -709,9 +679,7 @@ def test_composed_router_top1_factors_into_per_axis_argmax():
 
 
 def test_composed_router_supports_different_expert_counts_per_axis():
-    router = ComposedRouter(
-        [EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)]
-    )
+    router = ComposedRouter([EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)])
     assert router.n_experts == 10
     cond_cont, cond_cat = _cond(8, pdg=5)
     assert router.gate(cond_cont, cond_cat).shape == (8, 10)
@@ -719,9 +687,7 @@ def test_composed_router_supports_different_expert_counts_per_axis():
 
 def test_composed_router_classify_loss_sums_sub_router_losses():
     """energy/pdg both default to zero, so the composed loss should too."""
-    router = ComposedRouter(
-        [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
-    )
+    router = ComposedRouter([EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)])
     cond_cont, cond_cat = _cond(16, pdg=5)
     labels = torch.randint(0, 4, (16,))
     loss = router.classify_loss(cond_cont, cond_cat, labels)
@@ -964,9 +930,7 @@ def test_routed_stage1_eval_dispatch_matches_manual_grouping():
         idx = model.trunk.router.top1(cond_cont, cond_cat)
         manual = torch.zeros_like(x_t)
         for i in range(B):
-            manual[i] = model.trunk.experts[int(idx[i])](
-                x_t[i : i + 1], cond[i : i + 1]
-            )[0]
+            manual[i] = model.trunk.experts[int(idx[i])](x_t[i : i + 1], cond[i : i + 1])[0]
 
     torch.testing.assert_close(batched, manual, atol=1e-5, rtol=1e-4)
 
diff --git a/tests/test_sample.py b/tests/test_sample.py
index f6b12e7..7906069 100644
--- a/tests/test_sample.py
+++ b/tests/test_sample.py
@@ -30,9 +30,7 @@ def _particle_material_cfg(conditioning: str, emb_dim: int) -> tuple[dict, dict]
 
 def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]:
     cond_cont = torch.randn(B, COND_DIM)
-    cond_cat = torch.stack(
-        [torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
-    )
+    cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
     return cond_cont, cond_cat
 
 
@@ -43,12 +41,8 @@ def _conditioning_for(target: str) -> str:
     return "embedding" if target == "embedding" else "physical"
 
 
-def _stage2_oneshot(
-    target: str, generator: str, emb_dim: int = 6, pdg: int = 3, mat: int = 2
-) -> Stage2OneShot:
-    particle_cfg, material_cfg = _particle_material_cfg(
-        _conditioning_for(target), emb_dim
-    )
+def _stage2_oneshot(target: str, generator: str, emb_dim: int = 6, pdg: int = 3, mat: int = 2) -> Stage2OneShot:
+    particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
     particle_type_cfg = {"target": target}
     # build_models (giant/model/network.py) computes sec_dim this same way
     # before constructing Stage2OneShot — its own default (SEC_DIM, the
@@ -78,9 +72,7 @@ def _stage2_ar(
     k_max: int = 5,
     history: str = "markov",
 ) -> Stage2Autoregressive:
-    particle_cfg, material_cfg = _particle_material_cfg(
-        _conditioning_for(target), emb_dim
-    )
+    particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
     return Stage2Autoregressive(
         pdg_vocab=pdg,
         mat_vocab=mat,
@@ -163,9 +155,7 @@ def test_sample_secondaries_flow_shapes_by_target(target):
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
-    sec_cont, sec_type, sec_valid = sample_secondaries(
-        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
-    )
+    sec_cont, sec_type, sec_valid = sample_secondaries(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
     assert sec_cont.shape == (B, K_MAX, CONT_SLOT_DIM)
     assert sec_type.shape == (B, K_MAX, _expected_type_dim(target, emb_dim))
     assert sec_valid.shape == (B, K_MAX)
@@ -180,9 +170,7 @@ def test_sample_secondaries_wgan_shapes_by_target(target):
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
-    sec_cont, sec_type, sec_valid = sample_secondaries_wgan(
-        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
-    )
+    sec_cont, sec_type, sec_valid = sample_secondaries_wgan(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
     assert sec_cont.shape == (B, K_MAX, CONT_SLOT_DIM)
     assert sec_type.shape == (B, K_MAX, _expected_type_dim(target, emb_dim))
     assert sec_valid.shape == (B, K_MAX)
@@ -196,15 +184,11 @@ def test_sample_secondaries_wgan_shapes_by_target(target):
 @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
 def test_sample_secondaries_ar_shapes(target, generator, history):
     B, k_max, emb_dim = 4, 5, 6
-    decoder = _stage2_ar(
-        target, generator, emb_dim=emb_dim, k_max=k_max, history=history
-    )
+    decoder = _stage2_ar(target, generator, emb_dim=emb_dim, k_max=k_max, history=history)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.randint(0, k_max + 1, (B,))
-    sec_cont, sec_type, sec_valid = sample_secondaries_ar(
-        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
-    )
+    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
     assert sec_cont.shape == (B, k_max, CONT_SLOT_DIM)
     assert sec_type.shape == (B, k_max, _expected_type_dim(target, emb_dim))
     assert sec_valid.shape == (B, k_max)
@@ -220,9 +204,7 @@ def test_sample_secondaries_ar_valid_mask_matches_n_sec(target, generator):
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.tensor([0, 2, k_max])
-    _, _, sec_valid = sample_secondaries_ar(
-        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
-    )
+    _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
     for i, n in enumerate(n_sec_pred.tolist()):
         assert sec_valid[i, :n].all()
         assert not sec_valid[i, n:].any()
@@ -237,8 +219,6 @@ def test_sample_secondaries_ar_first_slot_has_no_history():
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.tensor([0, 1, 1])
-    sec_cont, sec_type, sec_valid = sample_secondaries_ar(
-        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
-    )
+    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
     assert sec_cont.shape == (B, 1, CONT_SLOT_DIM)
     assert sec_valid.tolist() == [[False], [True], [True]]
diff --git a/tests/test_setup_cache.py b/tests/test_setup_cache.py
index 70e8e40..584b5e8 100644
--- a/tests/test_setup_cache.py
+++ b/tests/test_setup_cache.py
@@ -14,9 +14,7 @@ from giant.data.transforms import Normalizer
 
 def _touch_parquet(path, n=1):
     path.parent.mkdir(parents=True, exist_ok=True)
-    pd.DataFrame(
-        {"pdg": [11] * n, "material": ["G4_AIR"] * n, "process": ["eIoni"] * n}
-    ).to_parquet(path)
+    pd.DataFrame({"pdg": [11] * n, "material": ["G4_AIR"] * n, "process": ["eIoni"] * n}).to_parquet(path)
     return path
 
 
@@ -29,9 +27,7 @@ def _normalizer(width=3):
 
 def _entry(n_train_steps=100, sample=None):
     sample = np.array([1.0, 2.0, 3.0], dtype=np.float32) if sample is None else sample
-    return NormalizerEntry(
-        _normalizer(), _normalizer(), _normalizer(2), n_train_steps, sample
-    )
+    return NormalizerEntry(_normalizer(), _normalizer(), _normalizer(2), n_train_steps, sample)
 
 
 # ── sidecar_path ─────────────────────────────────────────────────────────
@@ -39,9 +35,7 @@ def _entry(n_train_steps=100, sample=None):
 
 def test_sidecar_path_single_file(tmp_path):
     f = tmp_path / "shard.parquet"
-    assert (
-        setup_cache.sidecar_path(f) == tmp_path / "shard.parquet.giant_train_cache.json"
-    )
+    assert setup_cache.sidecar_path(f) == tmp_path / "shard.parquet.giant_train_cache.json"
 
 
 def test_sidecar_path_directory(tmp_path):
@@ -51,10 +45,7 @@ def test_sidecar_path_directory(tmp_path):
 
 def test_sidecar_path_manifest(tmp_path):
     m = tmp_path / "pools" / "full.manifest"
-    assert (
-        setup_cache.sidecar_path(m)
-        == tmp_path / "pools" / "full.manifest.giant_train_cache.json"
-    )
+    assert setup_cache.sidecar_path(m) == tmp_path / "pools" / "full.manifest.giant_train_cache.json"
 
 
 # ── fingerprint_files ────────────────────────────────────────────────────
diff --git a/tests/test_steps_to_parquet.py b/tests/test_steps_to_parquet.py
index 8e839c1..169bda8 100644
--- a/tests/test_steps_to_parquet.py
+++ b/tests/test_steps_to_parquet.py
@@ -24,18 +24,14 @@ def _frame() -> pl.DataFrame:
 def test_e_sec_sums_child_first_step_energy():
     out, n_orphaned = steps_to_parquet._add_secondary_attributes(_frame())
     assert n_orphaned == 0
-    e_sec = dict(
-        zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"])
-    )
+    e_sec = dict(zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"]))
     assert e_sec[(1, 0, 0)] == 15.0  # one child, first-step pre_E 15
     assert e_sec[(1, 0, 1)] == 50.0  # two children, 20 + 30
 
 
 def test_e_sec_zero_when_no_children():
     out, _ = steps_to_parquet._add_secondary_attributes(_frame())
-    childless = out.filter(
-        (pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1)
-    )
+    childless = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1))
     assert childless["e_sec"].item() == 0.0
 
 
@@ -69,9 +65,7 @@ def test_orphaned_child_track_is_dropped_not_nulled():
         }
     )
     out, n_orphaned = steps_to_parquet._add_secondary_attributes(df)
-    row = out.filter(
-        (pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0)
-    )
+    row = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0))
 
     assert n_orphaned == 1
     assert row["child_track_ids"].to_list() == [[2]]
diff --git a/tests/test_steps_to_parquet_parallel.py b/tests/test_steps_to_parquet_parallel.py
index 55e0cce..101d36f 100644
--- a/tests/test_steps_to_parquet_parallel.py
+++ b/tests/test_steps_to_parquet_parallel.py
@@ -124,31 +124,13 @@ def _make_dataset(tmp_path: Path, schemas: list[str] | None = None) -> Path:
 def test_resolve_destination_uses_latest_schema(tmp_path):
     root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3", "schema2"])
     dest = resolve_destination(root_file, tmp_path, schema_override=None)
-    assert (
-        dest
-        == tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema3"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
 
 
 def test_resolve_destination_schema_override_wins(tmp_path):
     root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3"])
     dest = resolve_destination(root_file, tmp_path, schema_override="schema9")
-    assert (
-        dest
-        == tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema9"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema9" / "pbwo4" / "shard-000.parquet"
 
 
 def test_resolve_destination_errors_without_any_schema(tmp_path):
diff --git a/tests/test_train.py b/tests/test_train.py
index ad14113..9a220fe 100644
--- a/tests/test_train.py
+++ b/tests/test_train.py
@@ -74,9 +74,7 @@ def test_wandb_run_config_includes_full_cfg_and_param_counts():
         "stage1_model": {"generator": "flow"},
         "stage2_model": {"generator": "wgan"},
     }
-    wcfg = _wandb_run_config(
-        cfg, model_config={"pdg_vocab": 3}, param_counts={"stage1": 100}
-    )
+    wcfg = _wandb_run_config(cfg, model_config={"pdg_vocab": 3}, param_counts={"stage1": 100})
     assert wcfg["train"] == {"lr": 3e-4}
     assert wcfg["stage1_model"] == {"generator": "flow"}
     assert wcfg["stage2_model"] == {"generator": "wgan"}
@@ -162,9 +160,7 @@ def test_type_repr_shapes_and_values(target):
     expected_width = PARTICLE_PHYS_DIM if target == "physical" else emb_dim
     assert repr_.shape == (B, K, expected_width)
     if target == "physical":
-        assert torch.equal(
-            repr_, sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM]
-        )
+        assert torch.equal(repr_, sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM])
     if target == "onehot":
         assert torch.all(repr_.sum(-1) == 1.0)
 
@@ -180,9 +176,7 @@ def test_type_repr_shapes_and_values(target):
         ("embedding", "wgan"),
     ],
 )
-def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(
-    target, generator
-):
+def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(target, generator):
     """Regression test tying the refactor together: _assemble_stage2_real is
     now defined as _assemble_stage2_ar_target(...).flatten(1)."""
     B, emb_dim = 4, 6
@@ -192,12 +186,8 @@ def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(
     if target == "embedding":
         cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim)
     particle_type_cfg = {"target": target}
-    flat = _assemble_stage2_real(
-        sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim
-    )
-    unflat = _assemble_stage2_ar_target(
-        sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim
-    )
+    flat = _assemble_stage2_real(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim)
+    unflat = _assemble_stage2_ar_target(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim)
     assert torch.equal(unflat.flatten(1), flat)
 
 
@@ -206,9 +196,7 @@ def test_assemble_stage2_ar_inputs_shapes_and_history_feat_width():
     sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM)
     sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX))
     cond_enc = torch.nn.Module()
-    out = _assemble_stage2_ar_inputs(
-        sec_cont, sec_type_idx, {"target": "physical"}, cond_enc, emb_dim
-    )
+    out = _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, {"target": "physical"}, cond_enc, emb_dim)
     assert out["history_feat"].shape == (B, K_MAX, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     assert out["has_prev"].shape == (B, K_MAX)
     assert out["remaining_frac"].shape == (B, K_MAX)
@@ -221,9 +209,7 @@ def test_relax_onehot_type_slice_grad_probe_populates_both_norms():
     B, k_max, cont_dim, type_dim = 4, K_MAX, CONT_SLOT_DIM, 6
     x_flat = torch.randn(B, k_max * (cont_dim + type_dim), requires_grad=True)
     grad_probe: dict[str, float] = {}
-    out = _relax_onehot_type_slice(
-        x_flat, k_max, cont_dim, type_dim, tau=0.5, grad_probe=grad_probe
-    )
+    out = _relax_onehot_type_slice(x_flat, k_max, cont_dim, type_dim, tau=0.5, grad_probe=grad_probe)
     out.sum().backward()
     assert grad_probe["cont"] >= 0.0
     assert grad_probe["type"] >= 0.0
@@ -324,9 +310,7 @@ def _fake_batches(n_batches, batch_size, seed=0):
         sec_cont = torch.randn(batch_size, K_MAX, SEC_SLOT_DIM, generator=g)
         proc_idx = torch.zeros(batch_size, dtype=torch.long)
         sec_type_idx = torch.zeros(batch_size, K_MAX, dtype=torch.long)
-        batches.append(
-            (cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)
-        )
+        batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
     return batches
 
 
@@ -409,17 +393,13 @@ def _run_train(cfg, out_dir, resume_path=None):
         ),
         (
             "stage2_onehot_target_wgan",
-            lambda cfg: cfg["stage2_model"].__setitem__(
-                "particle_type", {"target": "onehot", "lambda": 1.0}
-            ),
+            lambda cfg: cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
         ),
         (
             "stage2_onehot_target_flow",
             lambda cfg: (
                 cfg["stage2_model"].__setitem__("generator", "flow"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "onehot", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
             ),
         ),
         (
@@ -427,9 +407,7 @@ def _run_train(cfg, out_dir, resume_path=None):
             lambda cfg: (
                 cfg["conditioning"]["particle"].__setitem__("type", "embedding"),
                 cfg["conditioning"]["material"].__setitem__("type", "embedding"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "embedding", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
             ),
         ),
         (
@@ -438,18 +416,14 @@ def _run_train(cfg, out_dir, resume_path=None):
                 cfg["conditioning"]["particle"].__setitem__("type", "embedding"),
                 cfg["conditioning"]["material"].__setitem__("type", "embedding"),
                 cfg["stage2_model"].__setitem__("generator", "flow"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "embedding", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
             ),
         ),
         (
             "ar_wgan_onehot",
             lambda cfg: (
                 cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "onehot", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
             ),
         ),
         (
@@ -461,9 +435,7 @@ def _run_train(cfg, out_dir, resume_path=None):
             lambda cfg: (
                 cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
                 cfg["stage2_model"].__setitem__("generator", "flow"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "onehot", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
             ),
         ),
         (
@@ -473,9 +445,7 @@ def _run_train(cfg, out_dir, resume_path=None):
                 cfg["conditioning"]["material"].__setitem__("type", "embedding"),
                 cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
                 cfg["stage2_model"].__setitem__("generator", "flow"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "embedding", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
             ),
         ),
         (
@@ -491,9 +461,7 @@ def _run_train(cfg, out_dir, resume_path=None):
                 cfg["stage1_model"].__setitem__("generator", "wgan"),
                 cfg["stage2_model"].__setitem__("generator", "flow"),
                 cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "onehot", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
             ),
         ),
     ],
@@ -585,9 +553,7 @@ def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch():
         ema_decay=0.0,
         steps_per_epoch=4,
     )
-    trainer = WGANStageTrainer(
-        spec, models["stage1"], critics["stage1"], torch.device("cpu")
-    )
+    trainer = WGANStageTrainer(spec, models["stage1"], critics["stage1"], torch.device("cpu"))
     assert trainer.model.n_sec_head is None
     batch = _fake_batches(1, 8)[0]
     stats = trainer.step(batch, torch.device("cpu"), global_step=1)
@@ -595,9 +561,7 @@ def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch():
 
 
 def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
-    spec = StageSpec(
-        name="stage2", is_stage2=True, generator="ddpm", ddpm_n_steps=50, ema_decay=0.0
-    )
+    spec = StageSpec(name="stage2", is_stage2=True, generator="ddpm", ddpm_n_steps=50, ema_decay=0.0)
     with pytest.raises(NotImplementedError):
         FlowDDPMStageTrainer(spec, torch.nn.Linear(1, 1), torch.device("cpu"))
 
@@ -608,9 +572,7 @@ def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
 @pytest.mark.parametrize("teacher_forcing", ["always", "scheduled", "never"])
 @pytest.mark.parametrize("history", ["markov", "attention"])
 @pytest.mark.parametrize("stage2_generator", ["wgan", "flow"])
-def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(
-    teacher_forcing, history, stage2_generator
-):
+def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(teacher_forcing, history, stage2_generator):
     """v0.3.0 step 7: history='attention' and teacher_forcing in
     {'scheduled', 'never'} must actually train — a stage-2 AR trainer.step()
     must run and produce a finite loss, for every {history} x
@@ -629,9 +591,7 @@ def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(
     model_config = _model_config(cfg)
     models = build_models(model_config)
     critics = build_critics(model_config)
-    trainers = build_stage_trainers(
-        cfg, models, critics, torch.device("cpu"), total_train_batches=4
-    )
+    trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4)
     trainer = trainers["stage2"]
     batch = _fake_batches(1, 4)[0]
     stats = trainer.step(batch, torch.device("cpu"), global_step=1)
@@ -665,9 +625,7 @@ def test_train_end_to_end_ar_attention_history_scheduled_teacher_forcing(
         with open(out_dir / "metrics.csv", newline="") as f:
             rows = list(csv.DictReader(f))
         assert len(rows) == cfg["train"]["epochs"]
-        loss_col = (
-            "stage2/train/g_loss" if stage2_generator == "wgan" else "stage2/train/loss"
-        )
+        loss_col = "stage2/train/g_loss" if stage2_generator == "wgan" else "stage2/train/loss"
         assert all(math.isfinite(float(r[loss_col])) for r in rows)
 
 
diff --git a/tests/test_transforms.py b/tests/test_transforms.py
index 04ce58d..7f19326 100644
--- a/tests/test_transforms.py
+++ b/tests/test_transforms.py
@@ -162,9 +162,7 @@ def test_local_frame_rotation_normalizes_non_unit_pre_dir():
     post_dir = rng.standard_normal((N, 3)).astype(np.float32)
     post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)
 
-    pre_dir_scaled = pre_dir_unit * rng.uniform(0.9, 1.1, size=(N, 1)).astype(
-        np.float32
-    )
+    pre_dir_scaled = pre_dir_unit * rng.uniform(0.9, 1.1, size=(N, 1)).astype(np.float32)
     expected = local_frame_rotation(pre_dir_unit, post_dir)
     result = local_frame_rotation(pre_dir_scaled, post_dir)
     np.testing.assert_allclose(result, expected, atol=1e-4)
@@ -200,14 +198,10 @@ def test_reconstruct_post_pos_straight_line():
     step_length = rng.uniform(0.1, 5.0, size=N).astype(np.float32)
     post_pos = pre_pos + step_length[:, None] * pre_dir
 
-    travel_dir_local = local_frame_rotation(
-        pre_dir, travel_direction(pre_pos, post_pos)
-    )
+    travel_dir_local = local_frame_rotation(pre_dir, travel_direction(pre_pos, post_pos))
     np.testing.assert_allclose(travel_dir_local, np.tile([0, 0, 1], (N, 1)), atol=1e-4)
 
-    reconstructed = reconstruct_post_pos(
-        pre_pos, pre_dir, step_length, travel_dir_local
-    )
+    reconstructed = reconstruct_post_pos(pre_pos, pre_dir, step_length, travel_dir_local)
     np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)
 
 
@@ -221,12 +215,8 @@ def test_reconstruct_post_pos_general_roundtrip():
     post_pos = pre_pos + rng.standard_normal((N, 3)).astype(np.float32)
     step_length = np.linalg.norm(post_pos - pre_pos, axis=1).astype(np.float32)
 
-    travel_dir_local = local_frame_rotation(
-        pre_dir, travel_direction(pre_pos, post_pos)
-    )
-    reconstructed = reconstruct_post_pos(
-        pre_pos, pre_dir, step_length, travel_dir_local
-    )
+    travel_dir_local = local_frame_rotation(pre_dir, travel_direction(pre_pos, post_pos))
+    reconstructed = reconstruct_post_pos(pre_pos, pre_dir, step_length, travel_dir_local)
     np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)
 
 
@@ -356,9 +346,7 @@ def _step_data_no_sec_lists(n_sec: np.ndarray) -> dict:
 
 
 def test_build_features_proc_idx_zero_without_proc_map():
-    data = _minimal_step_data(
-        3, process=np.array(["compt", "phot", "eIoni"], dtype=object)
-    )
+    data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object))
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
     *_, proc_idx, _, _, _ = build_features(data, pdg_map, mat_map)
@@ -367,9 +355,7 @@ def test_build_features_proc_idx_zero_without_proc_map():
 
 
 def test_build_features_proc_idx_looks_up_proc_map():
-    data = _minimal_step_data(
-        3, process=np.array(["compt", "phot", "eIoni"], dtype=object)
-    )
+    data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object))
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
     proc_map = {"compt": 0, "phot": 1, "eIoni": 2}
 
@@ -396,9 +382,7 @@ def test_build_features_require_secondaries_ok_when_no_secondaries():
     data = _step_data_no_sec_lists(np.zeros(3, dtype=np.int32))
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
-    _, _, _, _, sec_cont, *_ = build_features(
-        data, pdg_map, mat_map, require_secondaries=True
-    )
+    _, _, _, _, sec_cont, *_ = build_features(data, pdg_map, mat_map, require_secondaries=True)
 
     assert not sec_cont.any()
 
@@ -413,11 +397,7 @@ def fake_material_props(monkeypatch):
     in by the user (see giant.materials.MaterialPropertiesNotFilledError)."""
     import giant.materials as gm
 
-    fake = {
-        "PbWO4": gm.MaterialProperties(
-            z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7
-        )
-    }
+    fake = {"PbWO4": gm.MaterialProperties(z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7)}
     monkeypatch.setattr(gm, "MATERIAL_PROPERTIES", fake)
     return fake
 
@@ -455,9 +435,7 @@ def test_build_features_physical_mode_shape_and_values(fake_material_props):
     assert cond_cont.shape[1] == COND_DIM
     mass, charge = particle_mass_charge(11)
     expected_log_mass = log_transform(np.array([mass]))[0]
-    np.testing.assert_allclose(
-        cond_cont[:, COND_DIM_BASE], expected_log_mass, atol=1e-5
-    )
+    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE], expected_log_mass, atol=1e-5)
     np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], charge)
     np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 2], 75.6)  # z_eff
     np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 3], 205.3)  # a_eff
@@ -500,9 +478,7 @@ def test_build_cond_features_mass_charge_override(fake_material_props):
         material_conditioning="physical",
     )
 
-    np.testing.assert_allclose(
-        cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0]))
-    )
+    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0])))
     np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], [2.0, -2.0])
 
 
@@ -714,9 +690,7 @@ def test_welford_accumulator_matches_naive_running_mean_reference():
     naive_M2 = np.zeros(F)
     naive_n = 0
     for chunk in chunks:
-        naive_mean, naive_M2, naive_n = naive_update(
-            naive_mean, naive_M2, naive_n, chunk
-        )
+        naive_mean, naive_M2, naive_n = naive_update(naive_mean, naive_M2, naive_n, chunk)
 
     acc = _WelfordAccumulator(F)
     for chunk in chunks:
diff --git a/tests/test_validate.py b/tests/test_validate.py
index 1e29517..ca220ae 100644
--- a/tests/test_validate.py
+++ b/tests/test_validate.py
@@ -57,9 +57,7 @@ def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int
         sec_cont = torch.randn(B, _K_MAX, SEC_SLOT_DIM)
         proc_idx = torch.zeros(B, dtype=torch.long)
         sec_type_idx = torch.randint(0, n_classes, (B, _K_MAX), dtype=torch.long)
-        batches.append(
-            (cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)
-        )
+        batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
     return batches
 
 
@@ -71,9 +69,7 @@ def test_validate_marginals_all_zero_secondaries_returns_nan_phys_kl(monkeypatch
     s1, s2 = _tiny_models()
     loader = _loader(n_sec_value=0)
 
-    def _fake_resolve_n_sec(
-        stage1_model, sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
-    ):
+    def _fake_resolve_n_sec(stage1_model, sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred):
         return torch.zeros(cond_cont.size(0), dtype=torch.long)
 
     monkeypatch.setattr("giant.validate.resolve_n_sec", _fake_resolve_n_sec)
diff --git a/tests/test_wgan.py b/tests/test_wgan.py
index 851527b..5b17aa8 100644
--- a/tests/test_wgan.py
+++ b/tests/test_wgan.py
@@ -152,9 +152,7 @@ def test_sample_secondaries_wgan_shape():
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.randint(0, K_MAX, (B,))
-    sec_cont, sec_phys, sec_valid = sample_secondaries_wgan(
-        model, cond_cont, cond_cat, stage1_out, n_sec_pred
-    )
+    sec_cont, sec_phys, sec_valid = sample_secondaries_wgan(model, cond_cont, cond_cat, stage1_out, n_sec_pred)
     assert sec_cont.shape == (B, K_MAX, 4)
     assert sec_phys.shape == (B, K_MAX, 2)
     assert sec_valid.shape == (B, K_MAX)
@@ -183,9 +181,7 @@ def test_gradient_penalty_masked():
     mask = _mask(B, n_sec)
     real = torch.randn(B, SEC_DIM) * mask
     fake = torch.randn(B, SEC_DIM) * mask
-    gp = gradient_penalty(
-        lambda x: sec_critic(x, cond_cont, cond_cat, stage1_out), real, fake, mask=mask
-    )
+    gp = gradient_penalty(lambda x: sec_critic(x, cond_cont, cond_cat, stage1_out), real, fake, mask=mask)
     assert gp.item() >= 0.0
 
 
@@ -195,9 +191,7 @@ def test_critic_loss_scalar_and_grad():
     cond_cont, cond_cat = _cond(B)
     real = torch.randn(B, X_DIM)
     fake = torch.randn(B, X_DIM)
-    loss = critic_loss(
-        lambda x: critic(x, cond_cont, cond_cat), real, fake.detach(), gp_weight=10.0
-    )
+    loss = critic_loss(lambda x: critic(x, cond_cont, cond_cat), real, fake.detach(), gp_weight=10.0)
     assert loss.shape == ()
     loss.backward()
     assert any(p.grad is not None for p in critic.parameters())