Add lazy polars I/O and duplicate KL/constraint checks for giant.analysis

load_predicted_local now reads predict parquet via a lazy polars scan with
column projection pushed into the reader, instead of materializing the
whole file as a pandas DataFrame. Also adds marginal_table_pl and
constraint_report_pl, polars-native duplicates that read straight from a
predict parquet path/LazyFrame and stay lazy per (group, dim) pair, so
peak memory is one column slice rather than the whole SampleCollection.

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

giant/analysis.py

@@ -51,6 +51,7 @@ 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
 
@@ -73,6 +74,7 @@ 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:
@@ -265,6 +267,11 @@ def _check_predict_metadata(path: Path) -> None:
         )
 
 
+_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.
 
@@ -274,26 +281,29 @@ def load_predicted_local(path: str | Path) -> SampleCollection:
     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.
-    """
-    path = Path(path)
-    _check_predict_metadata(path)
-    df = pd.read_parquet(path)
 
-    gen_log_local = df[[f"pred_{name}" for name in LOCAL_TARGET_NAMES]].to_numpy(dtype=np.float32)
-    real_log_local = df[[f"true_{name}" for name in LOCAL_TARGET_NAMES]].to_numpy(dtype=np.float32)
+    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 = np.column_stack([
-        df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
-        df["pre_E"].to_numpy(dtype=np.float32),
-        df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
-        df["layer_id"].to_numpy(dtype=np.float32),
-        df["n_sec"].to_numpy(dtype=np.float32),
-    ]).astype(np.float32)
+    cond_cont_raw = df.select(_COND_CONT_COLS).to_numpy().astype(np.float32)
 
     return SampleCollection(
         cond_cont_raw=cond_cont_raw,
@@ -359,6 +369,109 @@ def marginal_table(
     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,
@@ -544,6 +657,52 @@ def constraint_report(collection: SampleCollection, norm_tol: float = 0.05) -> p
     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

pyproject.toml

@@ -23,6 +23,7 @@ convert = [
 ]
 analysis = [
     "matplotlib>=3.8",
+    "polars>=1.0",
 ]
 
 [project.scripts]

tests/test_analysis.py

@@ -12,10 +12,12 @@ from giant.analysis import (
     RAW_TARGET_NAMES,
     SampleCollection,
     constraint_report,
+    constraint_report_pl,
     correlation_matrices,
     direction_alignment,
     load_predicted_local,
     marginal_table,
+    marginal_table_pl,
     plot_constraint_violations,
     plot_correlation_matrices,
     plot_direction_alignment,
@@ -241,3 +243,56 @@ def test_load_predicted_local_rejects_mismatched_schema_version(tmp_path):
     )
     with pytest.raises(ValueError, match="schema version"):
         load_predicted_local(path)
+
+
+def _predicted_local_path(tmp_path, n=200, seed=0):
+    path = tmp_path / "predicted_local.parquet"
+    _write_predicted_local_parquet(
+        path,
+        n=n,
+        rng=np.random.default_rng(seed),
+        metadata={
+            PREDICT_COORD_METADATA_KEY: "local",
+            PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
+        },
+    )
+    return path
+
+
+@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
+def test_marginal_table_pl_matches_numpy_version(tmp_path, group_by):
+    path = _predicted_local_path(tmp_path)
+    collection = load_predicted_local(path)
+
+    expected = marginal_table(collection, group_by=group_by).sort_values(["group", "dim"])
+    actual = marginal_table_pl(path, group_by=group_by).sort(["group", "dim"]).to_pandas()
+
+    assert list(expected["group"]) == list(actual["group"])
+    assert list(expected["n"]) == list(actual["n"])
+    for col in ["real_mean", "gen_mean", "real_std", "gen_std", "kl_real_gen"]:
+        np.testing.assert_allclose(
+            expected[col].to_numpy(), actual[col].to_numpy(), atol=1e-4, rtol=1e-4,
+        )
+
+
+def test_marginal_table_pl_rejects_missing_metadata(tmp_path):
+    path = tmp_path / "no_metadata.parquet"
+    _write_predicted_local_parquet(path, metadata=None)
+    with pytest.raises(ValueError, match="no '.*' parquet metadata"):
+        marginal_table_pl(path)
+
+
+def test_constraint_report_pl_matches_numpy_version(tmp_path):
+    path = _predicted_local_path(tmp_path)
+    collection = load_predicted_local(path)
+
+    expected = constraint_report(collection)
+    actual = constraint_report_pl(path).to_pandas()
+
+    assert list(expected["check"]) == list(actual["check"])
+    np.testing.assert_allclose(
+        expected["violation_rate"].to_numpy(), actual["violation_rate"].to_numpy(), atol=1e-6,
+    )
+    np.testing.assert_allclose(
+        expected["mean_abs_error"].to_numpy(), actual["mean_abs_error"].to_numpy(), atol=1e-4,
+    )

uv.lock

@@ -373,6 +373,7 @@ dependencies = [
 [package.optional-dependencies]
 analysis = [
     { name = "matplotlib" },
+    { name = "polars" },
 ]
 convert = [
     { name = "awkward" },
@@ -389,6 +390,7 @@ requires-dist = [
     { name = "matplotlib", marker = "extra == 'analysis'", specifier = ">=3.8" },
     { name = "numpy", specifier = ">=1.26" },
     { name = "pandas", specifier = ">=2.2" },
+    { name = "polars", marker = "extra == 'analysis'", specifier = ">=1.0" },
     { name = "polars", marker = "extra == 'convert'", specifier = ">=1.0" },
     { name = "pyarrow", specifier = ">=16" },
     { name = "pytest", marker = "extra == 'dev'", specifier = ">=8" },