giant/giant/analysis.py

"""Notebook-friendly diagnostics for a trained model's sample quality.

Typical use from a Jupyter notebook::

    from giant.analysis import load_model_bundle, make_val_loader, collect_samples
    from giant.analysis import plot_marginals, plot_correlation_matrices, plot_pairwise
    from giant.analysis import plot_direction_alignment, plot_constraint_violations

    bundle = load_model_bundle("runs/my_run/best.pt")
    val_loader = make_val_loader(bundle, "path/to/steps.parquet")
    samples = collect_samples(bundle, val_loader)

    plot_marginals(samples, group_by="energy")
    plot_correlation_matrices(samples)
    plot_pairwise(samples)
    plot_direction_alignment(samples)
    plot_constraint_violations(samples)

If predictions were already generated offline via `giant predict --coord local`,
skip the checkpoint/model entirely and load the parquet directly::

    from giant.analysis import load_predicted_local
    samples = load_predicted_local("path/to/steps_predicted_local.parquet")

(`--coord global` output isn't supported here — it has no ground-truth columns
to compare against.)

Three tiers of checks, building on the aggregate marginal/KL check in
`giant.validate.validate_marginals`:

1. stratified marginals  — per-dimension real-vs-generated comparison, sliced by
   pdg / material / energy so failures hidden by the aggregate don't go unnoticed.
2. joint structure        — correlation matrices, physically-coupled pairwise
   plots, and post/travel direction alignment, since marginals can match while
   the model decorrelates targets that are coupled by the underlying physics.
3. physical constraints   — unit-norm direction vectors and non-negative raw
   step_length/delta_e/edep, checked in denormalized physical units; nothing in
   the unconstrained MLP output enforces these, so violations are a pure
   generation artifact.

`collect_samples` takes a `steps` argument (forwarded to the flow ODE
integrator or, in ddim mode, the DDIM substep count) so a later
sampler-step-count ablation can sweep it by calling this function repeatedly
without any new plumbing.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import polars as pl
import torch
from torch.utils.data import DataLoader

import pyarrow.parquet as pq

from giant.config import warn_if_checkpoint_config_mismatch
from giant.constants import (
    LOCAL_TARGET_NAMES,
    PREDICT_COORD_METADATA_KEY,
    PREDICT_SCHEMA_VERSION,
    PREDICT_SCHEMA_VERSION_KEY,
)
from giant.data.dataset import train_val_split
from giant.data.loader import find_parquet_files, load_steps
from giant.data.transforms import Normalizer, build_features, inv_log_transform
from giant.model.network import DenoisingMLP
from giant.model.schedule import CosineSchedule
from giant.sample import sample_ddim, sample_ddpm, sample_flow
from giant.validate import _histogram_kl

_N_LOG_DIMS = 3  # log_step_length, log_delta_e, log_edep are the first 3 target dims
RAW_TARGET_NAMES = [n.removeprefix("log_") for n in LOCAL_TARGET_NAMES]
_LOG_EPS = (
    1e-8  # mirrors giant.data.transforms._EPS, duplicated for use in polars exprs
)


def _hist_edges(*arrays: np.ndarray, bins: int) -> np.ndarray:
    """Bin edges that don't blow up on near-constant data (e.g. a tight unit-norm cluster).

    Plain `np.linspace(lo, hi, bins+1)` raises when `hi - lo` is too small relative
    to float precision to support `bins` distinct edges, which is a real failure
    mode here (a well-trained model can push direction norms to within float32
    epsilon of 1.0), not just a test artifact.
    """
    lo = min(a.min() for a in arrays)
    hi = max(a.max() for a in arrays)
    if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
        lo, hi = lo - 0.5, hi + 0.5
    return np.linspace(lo, hi, bins + 1)


@dataclass
class ModelBundle:
    model: torch.nn.Module
    cond_normalizer: Normalizer
    target_normalizer: Normalizer
    pdg_map: dict[int, int]
    mat_map: dict[str, int]
    model_config: dict
    mode: str
    schedule: CosineSchedule | None
    device: torch.device

    @property
    def idx_to_pdg(self) -> dict[int, int]:
        return {v: k for k, v in self.pdg_map.items()}

    @property
    def idx_to_mat(self) -> dict[int, str]:
        return {v: k for k, v in self.mat_map.items()}


@dataclass
class SampleCollection:
    cond_cont_raw: np.ndarray  # (N, 9) denormalized conditioning (pre_E delogged)
    pdg: np.ndarray  # (N,) raw PDG codes
    material: np.ndarray  # (N,) raw material names
    real_raw: np.ndarray  # (N, 9) denormalized + delogged real targets
    gen_raw: np.ndarray  # (N, 9) denormalized + delogged generated targets
    # Normalized-space (model's native training space) targets — only available
    # when collected live with the normalizer (collect_samples). A predict
    # parquet has already been denormalized on disk with no normalizer
    # attached, so loaders built from one (e.g. load_predicted_local) leave
    # these as None rather than fabricate a value.
    real_norm: np.ndarray | None = None
    gen_norm: np.ndarray | None = None


def load_model_bundle(
    ckpt_path: str | Path,
    mode: str = "flow",
    device: torch.device | None = None,
) -> ModelBundle:
    """Reconstruct a trained model and its normalizers/vocab from a training checkpoint."""
    warn_if_checkpoint_config_mismatch(ckpt_path)
    device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)

    pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
    mat_map = dict(ckpt["mat_map"])

    model = DenoisingMLP(**ckpt["model_config"])
    model.load_state_dict(ckpt["model"])
    model.to(device).eval()

    cond_normalizer = Normalizer.from_dict(ckpt["normalizer"]["cond"])
    target_normalizer = Normalizer.from_dict(ckpt["normalizer"]["target"])
    schedule = CosineSchedule().to(device) if mode != "flow" else None

    return ModelBundle(
        model=model,
        cond_normalizer=cond_normalizer,
        target_normalizer=target_normalizer,
        pdg_map=pdg_map,
        mat_map=mat_map,
        model_config=ckpt["model_config"],
        mode=mode,
        schedule=schedule,
        device=device,
    )


def make_val_loader(
    bundle: ModelBundle,
    data: str | Path,
    val_fraction: float = 0.1,
    seed: int = 42,
    batch_size: int = 4096,
) -> DataLoader:
    """Build a DataLoader over the val split, normalized with the bundle's fitted stats.

    Loads the whole file into memory — fine for typical validation-set sizes; for
    very large datasets, build a StreamingStepsDataset directly (see giant.pipeline).
    """
    files = find_parquet_files(data)
    chunks = [load_steps(f) for f in files]
    full = {k: np.concatenate([c[k] for c in chunks]) for k in chunks[0]}

    cond_cont, cond_cat, target, _, _ = build_features(
        full,
        bundle.pdg_map,
        bundle.mat_map,
        cond_normalizer=bundle.cond_normalizer,
        target_normalizer=bundle.target_normalizer,
    )
    _, val_ds = train_val_split(
        full, cond_cont, cond_cat, target, val_fraction=val_fraction, seed=seed
    )
    return DataLoader(val_ds, batch_size=batch_size, shuffle=False)


def _to_raw_targets(target_norm: np.ndarray, normalizer: Normalizer) -> np.ndarray:
    raw = normalizer.inverse_transform(target_norm)
    raw[:, :_N_LOG_DIMS] = inv_log_transform(raw[:, :_N_LOG_DIMS])
    return raw


@torch.no_grad()
def collect_samples(
    bundle: ModelBundle,
    val_loader: DataLoader,
    n_batches: int | None = None,
    steps: int | None = None,
) -> SampleCollection:
    """Run the sampler over `val_loader`, pairing generations with real targets + conditioning.

    `steps` overrides the number of sampler steps (flow ODE steps or DDIM
    substeps); `None` keeps each sampler's own default. Unused in "ddpm"
    mode, which always runs the full schedule.
    """
    model = bundle.model
    device = bundle.device
    steps_kw = {} if steps is None else {"steps": steps}

    cond_list, real_list, gen_list = [], [], []
    for i, (cond_cont, cond_cat, x1) in enumerate(val_loader):
        if n_batches is not None and i >= n_batches:
            break
        cond_cont = cond_cont.to(device)
        cond_cat = cond_cat.to(device)

        if bundle.mode == "flow":
            gen = sample_flow(model, cond_cont, cond_cat, **steps_kw)
        elif bundle.mode == "ddpm":
            gen = sample_ddpm(model, cond_cont, cond_cat, bundle.schedule)
        else:
            gen = sample_ddim(model, cond_cont, cond_cat, bundle.schedule, **steps_kw)

        cond_full = torch.cat([cond_cont.cpu(), cond_cat.cpu().float()], dim=-1)
        cond_list.append(cond_full.numpy())
        real_list.append(x1.numpy())
        gen_list.append(gen.cpu().numpy())

    cond_all = np.concatenate(cond_list, axis=0)
    real_norm = np.concatenate(real_list, axis=0)
    gen_norm = np.concatenate(gen_list, axis=0)

    cond_cont_norm, cond_cat_arr = cond_all[:, :-2], cond_all[:, -2:]
    cond_cont_raw = bundle.cond_normalizer.inverse_transform(cond_cont_norm)
    cond_cont_raw[:, 3] = inv_log_transform(cond_cont_raw[:, 3])  # log(pre_E) -> pre_E

    idx_to_pdg, idx_to_mat = bundle.idx_to_pdg, bundle.idx_to_mat
    pdg = np.array([idx_to_pdg[i] for i in cond_cat_arr[:, 0].astype(np.int64)])
    material = np.array([idx_to_mat[i] for i in cond_cat_arr[:, 1].astype(np.int64)])

    return SampleCollection(
        cond_cont_raw=cond_cont_raw,
        pdg=pdg,
        material=material,
        real_raw=_to_raw_targets(real_norm, bundle.target_normalizer),
        gen_raw=_to_raw_targets(gen_norm, bundle.target_normalizer),
        real_norm=real_norm,
        gen_norm=gen_norm,
    )


def _check_predict_metadata(path: Path) -> None:
    """Verify a parquet file's giant-predict tag before trusting its column layout.

    Raises rather than warns: a wrong or missing tag means the column-layout
    assumptions below don't hold, so silently proceeding could mix up which
    columns are predictions vs. ground truth.
    """
    metadata = pq.read_schema(path).metadata or {}
    coord = metadata.get(PREDICT_COORD_METADATA_KEY.encode())
    if coord is None:
        raise ValueError(
            f"{path} has no '{PREDICT_COORD_METADATA_KEY}' parquet metadata — it wasn't "
            "written by `giant predict` (or predates schema tagging), so its column "
            "layout can't be verified"
        )
    if coord.decode() != "local":
        raise ValueError(
            f"{path} was written with --coord {coord.decode()!r}, not 'local' — "
            "load_predicted_local only supports coord=local predict output, which "
            "is the only mode that also writes ground-truth columns"
        )
    version = metadata.get(PREDICT_SCHEMA_VERSION_KEY.encode())
    if version is not None and version.decode() != PREDICT_SCHEMA_VERSION:
        raise ValueError(
            f"{path} has predict schema version {version.decode()!r}, but "
            f"giant.analysis expects {PREDICT_SCHEMA_VERSION!r} — column layout may "
            "have changed; update load_predicted_local to match"
        )


_COND_CONT_COLS = [
    "pre_x",
    "pre_y",
    "pre_z",
    "pre_E",
    "pre_dx",
    "pre_dy",
    "pre_dz",
    "layer_id",
    "n_sec",
]


def load_predicted_local(path: str | Path) -> SampleCollection:
    """Build a SampleCollection from a `giant predict --coord local` parquet file.

    Reads the `pred_*`/`true_*` columns directly — no checkpoint or model needed,
    since the predict CLI already denormalized them into the same log-scaled,
    local-frame space `collect_samples` produces internally before raw conversion.
    Requires the file to carry the `giant predict` metadata tag (see
    `_check_predict_metadata`); raises if it's missing or from --coord global,
    rather than guessing from column names.

    Reads via a lazy polars scan with the needed columns selected before
    `.collect()`, so column projection is pushed down into the parquet reader
    (e.g. `event_id` is never read) instead of materializing every column of
    the file as a pandas DataFrame first.
    """
    pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]
    true_cols = [f"true_{name}" for name in LOCAL_TARGET_NAMES]
    df = (
        _scan_predicted_local(path)
        .select(pred_cols + true_cols + _COND_CONT_COLS + ["pdg", "material"])
        .collect()
    )

    gen_log_local = df.select(pred_cols).to_numpy().astype(np.float32)
    real_log_local = df.select(true_cols).to_numpy().astype(np.float32)

    def to_raw(log_local: np.ndarray) -> np.ndarray:
        raw = log_local.copy()
        raw[:, :_N_LOG_DIMS] = inv_log_transform(raw[:, :_N_LOG_DIMS])
        return raw

    cond_cont_raw = df.select(_COND_CONT_COLS).to_numpy().astype(np.float32)

    return SampleCollection(
        cond_cont_raw=cond_cont_raw,
        pdg=df["pdg"].to_numpy(),
        material=df["material"].to_numpy(),
        real_raw=to_raw(real_log_local),
        gen_raw=to_raw(gen_log_local),
    )


# ---------------------------------------------------------------------------
# Tier 1: stratified marginals
# ---------------------------------------------------------------------------


def _group_labels(
    collection: SampleCollection,
    group_by: str | None,
    n_energy_bins: int,
) -> list[tuple[str, np.ndarray]]:
    n = len(collection.pdg)
    if group_by is None:
        return [("all", np.ones(n, dtype=bool))]
    if group_by == "pdg":
        return [(f"pdg={v}", collection.pdg == v) for v in np.unique(collection.pdg)]
    if group_by == "material":
        return [
            (f"material={v}", collection.material == v)
            for v in np.unique(collection.material)
        ]
    if group_by == "energy":
        pre_E = collection.cond_cont_raw[:, 3]
        edges = np.quantile(pre_E, np.linspace(0, 1, n_energy_bins + 1))
        edges[-1] += 1e-6
        bin_idx = np.digitize(pre_E, edges[1:-1])
        return [
            (f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})", bin_idx == i)
            for i in range(n_energy_bins)
        ]
    raise ValueError(f"unknown group_by={group_by!r}")


def marginal_table(
    collection: SampleCollection,
    group_by: str | None = None,
    n_energy_bins: int = 4,
    bins: int = 50,
) -> pd.DataFrame:
    """Per-dimension real-vs-generated summary stats + KL(real||gen), in raw units.

    `group_by`: None for an aggregate table, or one of "pdg", "material", "energy"
    to stratify rows by that conditioning variable. Sorted worst-KL first, so
    failure modes hidden by the aggregate surface at the top.
    """
    rows = []
    for label, mask in _group_labels(collection, group_by, n_energy_bins):
        if mask.sum() < 2:
            continue
        real, gen = collection.real_raw[mask], collection.gen_raw[mask]
        for j, name in enumerate(RAW_TARGET_NAMES):
            rows.append(
                {
                    "group": label,
                    "dim": name,
                    "n": int(mask.sum()),
                    "real_mean": real[:, j].mean(),
                    "gen_mean": gen[:, j].mean(),
                    "real_std": real[:, j].std(),
                    "gen_std": gen[:, j].std(),
                    "kl_real_gen": _histogram_kl(real[:, j], gen[:, j], bins=bins),
                }
            )
    return (
        pd.DataFrame(rows)
        .sort_values("kl_real_gen", ascending=False)
        .reset_index(drop=True)
    )


# ---------------------------------------------------------------------------
# Polars-native duplicates
#
# `marginal_table`/`constraint_report` above require a `SampleCollection`
# with the full real/gen arrays already materialized in numpy. The functions
# below instead take a parquet path (or LazyFrame) directly and stay lazy
# end to end: each (group, dim) pair is filtered, projected to just the two
# columns it needs, and collected on its own, so peak memory is one column
# pair rather than the whole file — useful when `load_predicted_local` itself
# would be too large to hold in memory at once.
# ---------------------------------------------------------------------------


def _histogram_kl_pl(
    p: pl.Series, q: pl.Series, bins: int = 50, eps: float = 1e-8
) -> float:
    """Polars duplicate of `giant.validate._histogram_kl`, binning via `Series.hist`."""
    lo, hi = min(p.min(), q.min()), max(p.max(), q.max())
    if hi <= lo:
        return 0.0
    edges = np.linspace(lo, hi, bins + 1).tolist()
    p_hist = p.hist(bins=edges)["count"].to_numpy().astype(np.float64) + eps
    q_hist = q.hist(bins=edges)["count"].to_numpy().astype(np.float64) + eps
    p_hist /= p_hist.sum()
    q_hist /= q_hist.sum()
    return float(np.sum(p_hist * np.log(p_hist / q_hist)))


def _raw_expr(col: str, j: int) -> pl.Expr:
    """`pred_*`/`true_*` columns are log-scaled for the first `_N_LOG_DIMS` dims."""
    e = pl.col(col)
    return (e.exp() - _LOG_EPS) if j < _N_LOG_DIMS else e


def _scan_predicted_local(source: str | Path | pl.LazyFrame) -> pl.LazyFrame:
    if isinstance(source, pl.LazyFrame):
        return source
    path = Path(source)
    _check_predict_metadata(path)
    return pl.scan_parquet(path)


def _group_filters_pl(
    lf: pl.LazyFrame,
    group_by: str | None,
    n_energy_bins: int,
) -> list[tuple[str, pl.Expr]]:
    if group_by is None:
        return [("all", pl.lit(True))]
    if group_by == "pdg":
        vals = lf.select("pdg").unique().collect()["pdg"].sort().to_list()
        return [(f"pdg={v}", pl.col("pdg") == v) for v in vals]
    if group_by == "material":
        vals = lf.select("material").unique().collect()["material"].sort().to_list()
        return [(f"material={v}", pl.col("material") == v) for v in vals]
    if group_by == "energy":
        pre_E = lf.select("pre_E").collect().to_series().to_numpy()
        edges = np.quantile(pre_E, np.linspace(0, 1, n_energy_bins + 1))
        edges[-1] += 1e-6
        return [
            (
                f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})",
                (pl.col("pre_E") >= edges[i]) & (pl.col("pre_E") < edges[i + 1]),
            )
            for i in range(n_energy_bins)
        ]
    raise ValueError(f"unknown group_by={group_by!r}")


def marginal_table_pl(
    source: str | Path | pl.LazyFrame,
    group_by: str | None = None,
    n_energy_bins: int = 4,
    bins: int = 50,
) -> pl.DataFrame:
    """Polars duplicate of `marginal_table`, reading straight from a predict parquet.

    `source` is a path to a `giant predict --coord local` parquet file, or an
    already-built LazyFrame with the same `pred_*`/`true_*`/`pdg`/`material`/
    `pre_E` columns (e.g. for testing). Never builds a `SampleCollection` —
    see the module-level note above on why this stays lazy.
    """
    lf = _scan_predicted_local(source)
    pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]
    true_cols = [f"true_{name}" for name in LOCAL_TARGET_NAMES]

    rows = []
    for label, cond in _group_filters_pl(lf, group_by, n_energy_bins):
        glf = lf.filter(cond)
        n = glf.select(pl.len()).collect().item()
        if n < 2:
            continue
        for j, name in enumerate(RAW_TARGET_NAMES):
            pair = glf.select(
                [
                    _raw_expr(true_cols[j], j).alias("real"),
                    _raw_expr(pred_cols[j], j).alias("gen"),
                ]
            ).collect()
            real_s, gen_s = pair["real"], pair["gen"]
            rows.append(
                {
                    "group": label,
                    "dim": name,
                    "n": n,
                    "real_mean": real_s.mean(),
                    "gen_mean": gen_s.mean(),
                    # ddof=0 to match numpy's (population-std) default used by marginal_table
                    "real_std": real_s.std(ddof=0),
                    "gen_std": gen_s.std(ddof=0),
                    "kl_real_gen": _histogram_kl_pl(real_s, gen_s, bins=bins),
                }
            )
    return pl.DataFrame(rows).sort("kl_real_gen", descending=True)


def plot_marginals(
    collection: SampleCollection,
    dims: list[str] | None = None,
    group_by: str | None = None,
    n_energy_bins: int = 4,
    bins: int = 50,
    max_groups: int = 6,
    figsize_per_axis: tuple[float, float] = (3.5, 2.8),
):
    """Overlaid real-vs-generated histograms: one row per group, one column per dim.

    Without `group_by`, a single row over the whole val set. With "pdg",
    "material", or "energy", one row per stratum, worst-KL groups first
    (capped at `max_groups`), so failures hidden by the aggregate are visible.
    """
    dims = dims or RAW_TARGET_NAMES
    dim_idx = [RAW_TARGET_NAMES.index(d) for d in dims]
    groups = _group_labels(collection, group_by, n_energy_bins)

    if group_by is not None:
        table = marginal_table(
            collection, group_by=group_by, n_energy_bins=n_energy_bins, bins=bins
        )
        worst_first = (
            table.groupby("group")["kl_real_gen"].max().sort_values(ascending=False)
        )
        order = {label: rank for rank, label in enumerate(worst_first.index)}
        groups = sorted(groups, key=lambda g: order[g[0]])[:max_groups]

    n_rows, n_cols = len(groups), len(dims)
    fig, axes = plt.subplots(
        n_rows,
        n_cols,
        squeeze=False,
        figsize=(figsize_per_axis[0] * n_cols, figsize_per_axis[1] * n_rows),
    )
    for row, (label, mask) in enumerate(groups):
        real, gen = collection.real_raw[mask], collection.gen_raw[mask]
        for col, j in enumerate(dim_idx):
            ax = axes[row][col]
            edges = _hist_edges(real[:, j], gen[:, j], bins=bins)
            ax.hist(real[:, j], bins=edges, density=True, alpha=0.5, label="real")
            ax.hist(gen[:, j], bins=edges, density=True, alpha=0.5, label="generated")
            if row == 0:
                ax.set_title(dims[col], fontsize=9)
            if col == 0:
                ax.set_ylabel(label, fontsize=8)
            if row == 0 and col == n_cols - 1:
                ax.legend(fontsize=7)
    fig.tight_layout()
    return fig


# ---------------------------------------------------------------------------
# Tier 2: joint structure
# ---------------------------------------------------------------------------


def correlation_matrices(collection: SampleCollection) -> tuple[np.ndarray, np.ndarray]:
    """Pearson correlation matrices of the raw targets, real vs generated."""
    return (
        np.corrcoef(collection.real_raw, rowvar=False),
        np.corrcoef(collection.gen_raw, rowvar=False),
    )


def plot_correlation_matrices(collection: SampleCollection):
    """Side-by-side real/generated correlation heatmaps, plus their difference."""
    real_corr, gen_corr = correlation_matrices(collection)
    diff = gen_corr - real_corr

    fig, axes = plt.subplots(1, 3, figsize=(13, 4))
    for ax, mat, title, cmap, vlim in [
        (axes[0], real_corr, "real", "coolwarm", (-1, 1)),
        (axes[1], gen_corr, "generated", "coolwarm", (-1, 1)),
        (axes[2], diff, "generated − real", "PuOr", (-0.5, 0.5)),
    ]:
        im = ax.imshow(mat, vmin=vlim[0], vmax=vlim[1], cmap=cmap)
        ax.set_xticks(range(len(RAW_TARGET_NAMES)))
        ax.set_xticklabels(RAW_TARGET_NAMES, rotation=90, fontsize=7)
        ax.set_yticks(range(len(RAW_TARGET_NAMES)))
        ax.set_yticklabels(RAW_TARGET_NAMES, fontsize=7)
        ax.set_title(title)
        fig.colorbar(im, ax=ax, fraction=0.046)
    fig.tight_layout()
    return fig


_DEFAULT_PAIRS = [
    ("step_length", "delta_e"),
    ("delta_e", "edep"),
    ("step_length", "edep"),
]


def plot_pairwise(
    collection: SampleCollection,
    pairs: list[tuple[str, str]] | None = None,
    n_sample: int = 3000,
    seed: int = 0,
):
    """Real-vs-generated scatter for physically coupled target pairs.

    Marginals matching doesn't imply the joint does — these pairs are coupled by
    the underlying physics (energy loss tracks distance, edep is part of
    delta_e), so a model that decorrelates them shows up here even with clean
    per-dimension marginals.
    """
    pairs = pairs or _DEFAULT_PAIRS
    rng = np.random.default_rng(seed)
    n = len(collection.pdg)
    idx = rng.choice(n, size=min(n_sample, n), replace=False)

    fig, axes = plt.subplots(2, len(pairs), squeeze=False, figsize=(4 * len(pairs), 7))
    for col, (a, b) in enumerate(pairs):
        ia, ib = RAW_TARGET_NAMES.index(a), RAW_TARGET_NAMES.index(b)
        for row, (data, title) in enumerate(
            [(collection.real_raw, "real"), (collection.gen_raw, "generated")]
        ):
            ax = axes[row][col]
            ax.scatter(data[idx, ia], data[idx, ib], s=3, alpha=0.3)
            ax.set_xlabel(a)
            ax.set_ylabel(b)
            if col == 0:
                ax.set_title(title, loc="left", fontsize=9)
    fig.tight_layout()
    return fig


def direction_alignment(collection: SampleCollection) -> tuple[np.ndarray, np.ndarray]:
    """cos(angle) between post_dir_local and travel_dir_local, real vs generated.

    These two unit vectors are coupled through the scattering physics, so their
    joint alignment is a check the per-dimension marginals can't see.
    """

    def cos_angle(raw: np.ndarray) -> np.ndarray:
        post, travel = raw[:, 3:6], raw[:, 6:9]
        return np.sum(post * travel, axis=1) / (
            np.linalg.norm(post, axis=1) * np.linalg.norm(travel, axis=1) + 1e-8
        )

    return cos_angle(collection.real_raw), cos_angle(collection.gen_raw)


def plot_direction_alignment(collection: SampleCollection, bins: int = 50):
    real_cos, gen_cos = direction_alignment(collection)
    fig, ax = plt.subplots(figsize=(5, 4))
    edges = np.linspace(-1, 1, bins + 1)
    ax.hist(real_cos, bins=edges, density=True, alpha=0.5, label="real")
    ax.hist(gen_cos, bins=edges, density=True, alpha=0.5, label="generated")
    ax.set_xlabel("cos(angle) between post_dir and travel_dir")
    ax.legend()
    fig.tight_layout()
    return fig


# ---------------------------------------------------------------------------
# Tier 3: physical constraints
# ---------------------------------------------------------------------------


def constraint_report(
    collection: SampleCollection, norm_tol: float = 0.05
) -> pd.DataFrame:
    """Rate of physical-constraint violations in the generated raw-space samples.

    The model is an unconstrained MLP, so nothing forces post_dir_local /
    travel_dir_local to stay unit-norm or step_length/delta_e/edep to stay
    non-negative — both hold by construction in the real data, so any
    violation rate here is purely a generation artifact.
    """
    gen = collection.gen_raw
    post_norm = np.linalg.norm(gen[:, 3:6], axis=1)
    travel_norm = np.linalg.norm(gen[:, 6:9], axis=1)

    rows = [
        {
            "check": "post_dir unit norm",
            "violation_rate": float(np.mean(np.abs(post_norm - 1) > norm_tol)),
            "mean_abs_error": float(np.mean(np.abs(post_norm - 1))),
        },
        {
            "check": "travel_dir unit norm",
            "violation_rate": float(np.mean(np.abs(travel_norm - 1) > norm_tol)),
            "mean_abs_error": float(np.mean(np.abs(travel_norm - 1))),
        },
    ]
    for j, name in enumerate(RAW_TARGET_NAMES[:_N_LOG_DIMS]):
        rows.append(
            {
                "check": f"{name} >= 0",
                "violation_rate": float(np.mean(gen[:, j] < 0)),
                "mean_abs_error": float(np.mean(np.clip(-gen[:, j], 0, None))),
            }
        )
    return pd.DataFrame(rows)


def constraint_report_pl(
    source: str | Path | pl.LazyFrame, norm_tol: float = 0.05
) -> pl.DataFrame:
    """Polars duplicate of `constraint_report`, reading straight from a predict parquet.

    Computed as a single set of lazy aggregation expressions over the
    `pred_*` columns — never materializes `gen_raw`. See `marginal_table_pl`
    for the `source` argument and why this stays lazy.
    """
    lf = _scan_predicted_local(source)
    pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]

    post_norm = sum(pl.col(pred_cols[k]) ** 2 for k in range(3, 6)).sqrt()
    travel_norm = sum(pl.col(pred_cols[k]) ** 2 for k in range(6, 9)).sqrt()
    raw_log_dims = [_raw_expr(pred_cols[j], j) for j in range(_N_LOG_DIMS)]

    agg = (
        lf.select(
            [
                ((post_norm - 1).abs() > norm_tol)
                .mean()
                .alias("post_dir_violation_rate"),
                (post_norm - 1).abs().mean().alias("post_dir_mean_abs_error"),
                ((travel_norm - 1).abs() > norm_tol)
                .mean()
                .alias("travel_dir_violation_rate"),
                (travel_norm - 1).abs().mean().alias("travel_dir_mean_abs_error"),
                *[
                    (raw < 0).mean().alias(f"{name}_violation_rate")
                    for raw, name in zip(raw_log_dims, RAW_TARGET_NAMES[:_N_LOG_DIMS])
                ],
                *[
                    raw.clip(upper_bound=0).abs().mean().alias(f"{name}_mean_abs_error")
                    for raw, name in zip(raw_log_dims, RAW_TARGET_NAMES[:_N_LOG_DIMS])
                ],
            ]
        )
        .collect()
        .row(0, named=True)
    )

    rows = [
        {
            "check": "post_dir unit norm",
            "violation_rate": agg["post_dir_violation_rate"],
            "mean_abs_error": agg["post_dir_mean_abs_error"],
        },
        {
            "check": "travel_dir unit norm",
            "violation_rate": agg["travel_dir_violation_rate"],
            "mean_abs_error": agg["travel_dir_mean_abs_error"],
        },
    ]
    for name in RAW_TARGET_NAMES[:_N_LOG_DIMS]:
        rows.append(
            {
                "check": f"{name} >= 0",
                "violation_rate": agg[f"{name}_violation_rate"],
                "mean_abs_error": agg[f"{name}_mean_abs_error"],
            }
        )
    return pl.DataFrame(rows)


def plot_constraint_violations(collection: SampleCollection):
    """Histograms backing `constraint_report`: direction norms and sign of the log-dims."""
    gen = collection.gen_raw
    post_norm = np.linalg.norm(gen[:, 3:6], axis=1)
    travel_norm = np.linalg.norm(gen[:, 6:9], axis=1)

    n_panels = 2 + _N_LOG_DIMS
    fig, axes = plt.subplots(1, n_panels, figsize=(4 * n_panels, 3.5))
    for ax, norm, title in [
        (axes[0], post_norm, "||post_dir||"),
        (axes[1], travel_norm, "||travel_dir||"),
    ]:
        ax.hist(norm, bins=_hist_edges(norm, bins=50))
        ax.axvline(1.0, color="k", linestyle="--", linewidth=1)
        ax.set_title(title)
    for k, name in enumerate(RAW_TARGET_NAMES[:_N_LOG_DIMS]):
        ax = axes[2 + k]
        ax.hist(gen[:, k], bins=_hist_edges(gen[:, k], bins=50))
        ax.axvline(0.0, color="k", linestyle="--", linewidth=1)
        ax.set_title(f"generated {name}")
    fig.tight_layout()
    return fig