diff --git a/giant/analysis.py b/giant/analysis.py index 2e25aa6..d40f5e1 100644 --- a/giant/analysis.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 2bc8b6f..ef6ff16 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ convert = [ ] analysis = [ "matplotlib>=3.8", + "polars>=1.0", ] [project.scripts] diff --git a/tests/test_analysis.py b/tests/test_analysis.py index 608c3eb..16ea7ea 100644 --- a/tests/test_analysis.py +++ b/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, + ) diff --git a/uv.lock b/uv.lock index aa505b3..418cc41 100644 --- a/uv.lock +++ b/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" },