Add KL bar plots and sample_frac to load_predicted_local; ignore root parquet scratch files

Adds plot_kl_bars/plot_kl_bars_pl (numpy/polars variants) for ranking which
target dimension or pdg/material stratum drives KL regressions, with the
same kl*n group capping as plot_marginals. Switches existing histogram
plots to step-type/log-scale. Adds sample_frac to load_predicted_local for
subsampling large predict parquets.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

.gitignore

@@ -11,3 +11,6 @@ wheels/
 
 # Training results
 checkpoints/
+
+# Predicted/scratch parquet files dropped in the repo root
+/*.parquet

giant/analysis.py

@@ -3,14 +3,16 @@
 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
+    from giant.analysis import plot_marginals, plot_kl_bars, plot_correlation_matrices
+    from giant.analysis import plot_pairwise, plot_direction_alignment
+    from giant.analysis import 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_kl_bars(samples, group_by="energy")
     plot_correlation_matrices(samples)
     plot_pairwise(samples)
     plot_direction_alignment(samples)
@@ -303,7 +305,9 @@ _COND_CONT_COLS = [
 ]
 
 
-def load_predicted_local(path: str | Path) -> SampleCollection:
+def load_predicted_local(
+    path: str | Path, sample_frac: float = 1.0, seed: int = 0
+) -> SampleCollection:
     """Build a SampleCollection from a `giant predict --coord local` parquet file.
 
     Reads the `pred_*`/`true_*` columns directly — no checkpoint or model needed,
@@ -317,7 +321,15 @@ def load_predicted_local(path: str | Path) -> SampleCollection:
     `.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.
+
+    `sample_frac` (0 < sample_frac <= 1) randomly keeps only that fraction of
+    rows after the column projection — useful for files too large to
+    comfortably hold as numpy arrays in `real_raw`/`gen_raw`. Sampling happens
+    after `.collect()` since polars' row-level sampling isn't pushed down into
+    the lazy scan; `seed` makes the subsample reproducible.
     """
+    if not (0 < sample_frac <= 1):
+        raise ValueError(f"sample_frac must be in (0, 1], got {sample_frac}")
     pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]
     true_cols = [f"true_{name}" for name in LOCAL_TARGET_NAMES]
     df = (
@@ -325,6 +337,8 @@ def load_predicted_local(path: str | Path) -> SampleCollection:
         .select(pred_cols + true_cols + _COND_CONT_COLS + ["pdg", "material"])
         .collect()
     )
+    if sample_frac < 1.0:
+        df = df.sample(fraction=sample_frac, seed=seed)
 
     gen_log_local = df.select(pred_cols).to_numpy().astype(np.float32)
     real_log_local = df.select(true_cols).to_numpy().astype(np.float32)
@@ -542,8 +556,10 @@ def plot_marginals(
     """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.
+    "material", or "energy", one row per stratum, ranked by max(kl) * n
+    (capped at `max_groups`) so groups that are both badly wrong and common
+    in the dataset surface first, rather than rare groups with a noisy,
+    high-variance KL estimate from just one or two samples.
     """
     dims = dims or RAW_TARGET_NAMES
     dim_idx = [RAW_TARGET_NAMES.index(d) for d in dims]
@@ -553,10 +569,10 @@ def plot_marginals(
         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)
-        )
+        by_group = table.groupby("group").agg(kl_max=("kl_real_gen", "max"), n=("n", "first"))
+        worst_first = (by_group["kl_max"] * by_group["n"]).sort_values(ascending=False)
         order = {label: rank for rank, label in enumerate(worst_first.index)}
+        groups = [g for g in groups if g[0] in order]
         groups = sorted(groups, key=lambda g: order[g[0]])[:max_groups]
 
     n_rows, n_cols = len(groups), len(dims)
@@ -571,8 +587,9 @@ def plot_marginals(
         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")
+            ax.hist(real[:, j], bins=edges, density=True, histtype="step", label="real")
+            ax.hist(gen[:, j], bins=edges, density=True, histtype="step", label="generated")
+            ax.set_yscale("log")
             if row == 0:
                 ax.set_title(dims[col], fontsize=9)
             if col == 0:
@@ -583,6 +600,96 @@ def plot_marginals(
     return fig
 
 
+def _limit_groups_by_kl_n(table: pd.DataFrame, max_groups: int) -> pd.DataFrame:
+    """Keep only the `max_groups` groups with the largest max(kl) * n, like `plot_marginals`.
+
+    Ranks by how badly wrong *and* how common a stratum is, rather than by KL
+    alone, so a rare pdg/material with a noisy, high-variance KL estimate from
+    a handful of samples doesn't crowd out groups that actually matter.
+    """
+    by_group = table.groupby("group").agg(kl_max=("kl_real_gen", "max"), n=("n", "first"))
+    worst_first = (by_group["kl_max"] * by_group["n"]).sort_values(ascending=False)
+    keep = set(worst_first.index[:max_groups])
+    return table[table["group"].isin(keep)]
+
+
+def _plot_kl_bars(table: pd.DataFrame, figsize: tuple[float, float]):
+    """Shared bar-plot body for `plot_kl_bars`/`plot_kl_bars_pl`.
+
+    `table` is a `marginal_table`/`marginal_table_pl` result (already converted
+    to pandas in the polars case) with `group`/`dim`/`kl_real_gen` columns.
+    """
+    pivot = table.pivot(index="dim", columns="group", values="kl_real_gen")
+    pivot = pivot.reindex(RAW_TARGET_NAMES)
+    groups = sorted(pivot.columns)
+    pivot = pivot[groups]
+
+    n_dims, n_groups = len(RAW_TARGET_NAMES), len(groups)
+    x = np.arange(n_dims)
+    width = 0.8 / n_groups
+
+    fig, ax = plt.subplots(figsize=figsize)
+    for i, group in enumerate(groups):
+        offset = (i - (n_groups - 1) / 2) * width
+        ax.bar(x + offset, pivot[group].to_numpy(), width=width, label=group)
+    ax.set_xticks(x)
+    ax.set_xticklabels(RAW_TARGET_NAMES, rotation=45, ha="right")
+    ax.set_ylabel("KL(real || gen)")
+    if n_groups > 1:
+        ax.legend(fontsize=7)
+    fig.tight_layout()
+    return fig
+
+
+def plot_kl_bars(
+    collection: SampleCollection,
+    group_by: str | None = None,
+    n_energy_bins: int = 4,
+    bins: int = 50,
+    max_groups: int = 6,
+    figsize: tuple[float, float] = (8, 4),
+):
+    """Bar plot of KL(real||gen) per target dimension, optionally stratified.
+
+    One bar cluster per dimension; with `group_by` set, one bar per stratum
+    within each cluster, so which dimension/stratum combination drives a KL
+    regression is visible at a glance rather than buried in `marginal_table`'s
+    sorted rows. With "pdg" or "material" — open-ended vocabularies that can
+    run to many distinct values — groups are capped at `max_groups`, ranked by
+    max(kl) * n as in `plot_marginals`; "energy" is already bounded by
+    `n_energy_bins` and isn't capped. Built on top of `marginal_table`; see
+    `plot_kl_bars_pl` for the polars-native, parquet-direct equivalent.
+    """
+    table = marginal_table(
+        collection, group_by=group_by, n_energy_bins=n_energy_bins, bins=bins
+    )
+    if group_by in ("pdg", "material"):
+        table = _limit_groups_by_kl_n(table, max_groups)
+    return _plot_kl_bars(table, figsize=figsize)
+
+
+def plot_kl_bars_pl(
+    source: str | Path | pl.LazyFrame,
+    group_by: str | None = None,
+    n_energy_bins: int = 4,
+    bins: int = 50,
+    max_groups: int = 6,
+    figsize: tuple[float, float] = (8, 4),
+):
+    """Polars duplicate of `plot_kl_bars`, reading straight from a predict parquet.
+
+    Built on top of `marginal_table_pl`; see that function for the `source`
+    argument and why it stays lazy until the final per-(group, dim) collect.
+    See `plot_kl_bars` for the `max_groups` capping behavior.
+    """
+    table = marginal_table_pl(
+        source, group_by=group_by, n_energy_bins=n_energy_bins, bins=bins
+    ).to_pandas()
+    if group_by in ("pdg", "material"):
+        table = _limit_groups_by_kl_n(table, max_groups)
+    return _plot_kl_bars(table, figsize=figsize)
+
+
 # ---------------------------------------------------------------------------
 # Tier 2: joint structure
 # ---------------------------------------------------------------------------
@@ -679,8 +786,9 @@ 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.hist(real_cos, bins=edges, density=True, histtype="step", label="real")
+    ax.hist(gen_cos, bins=edges, density=True, histtype="step", label="generated")
+    ax.set_yscale("log")
     ax.set_xlabel("cos(angle) between post_dir and travel_dir")
     ax.legend()
     fig.tight_layout()
@@ -805,12 +913,14 @@ def plot_constraint_violations(collection: SampleCollection):
         (axes[0], post_norm, "||post_dir||"),
         (axes[1], travel_norm, "||travel_dir||"),
     ]:
-        ax.hist(norm, bins=_hist_edges(norm, bins=50))
+        ax.hist(norm, bins=_hist_edges(norm, bins=50), histtype="step")
+        ax.set_yscale("log")
         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.hist(gen[:, k], bins=_hist_edges(gen[:, k], bins=50), histtype="step")
+        ax.set_yscale("log")
         ax.axvline(0.0, color="k", linestyle="--", linewidth=1)
         ax.set_title(f"generated {name}")
     fig.tight_layout()

tests/test_analysis.py

@@ -20,6 +20,8 @@ from giant.analysis import (
     plot_constraint_violations,
     plot_correlation_matrices,
     plot_direction_alignment,
+    plot_kl_bars,
+    plot_kl_bars_pl,
     plot_marginals,
     plot_pairwise,
 )
@@ -155,6 +157,24 @@ def test_plot_constraint_violations_runs_without_error():
     assert fig is not None
 
 
+def test_plot_kl_bars_runs_without_error():
+    fig = plot_kl_bars(_make_collection())
+    assert fig is not None
+
+
+def test_plot_kl_bars_grouped_runs_without_error():
+    fig = plot_kl_bars(_make_collection(), group_by="material")
+    assert fig is not None
+
+
+def test_plot_kl_bars_caps_groups_by_pdg():
+    collection = _make_collection(n=600)
+    collection.pdg = np.arange(600) % 8  # 8 distinct pdg values, > max_groups
+    fig = plot_kl_bars(collection, group_by="pdg", max_groups=3)
+    ax = fig.axes[0]
+    assert len({line.get_label() for line in ax.containers}) <= 3
+
+
 def _write_predicted_local_parquet(path, n=50, metadata=None, rng=None):
     """Mimic `giant predict --coord local`'s output schema for the loader tests."""
     rng = rng or np.random.default_rng(0)
@@ -300,6 +320,13 @@ def test_marginal_table_pl_matches_numpy_version(tmp_path, group_by):
         )
 
 
+@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
+def test_plot_kl_bars_pl_runs_without_error(tmp_path, group_by):
+    path = _predicted_local_path(tmp_path)
+    fig = plot_kl_bars_pl(path, group_by=group_by)
+    assert fig is not None
+
+
 def test_marginal_table_pl_rejects_missing_metadata(tmp_path):
     path = tmp_path / "no_metadata.parquet"
     _write_predicted_local_parquet(path, metadata=None)