WIP: energy-conservation PoC analysis/transforms updates
This commit is contained in:
+29
-29
File diff suppressed because one or more lines are too long
+261
-284
@@ -1,15 +1,16 @@
|
||||
"""Notebook-friendly diagnostics for a trained model's sample quality.
|
||||
|
||||
Typical use from a Jupyter notebook::
|
||||
All checks here read from `giant predict --coord local` parquet output
|
||||
(`pred_*`/`true_*` columns, denormalized but still local-frame/log-scaled) —
|
||||
there is no on-the-fly (checkpoint + live sampler) path; generate predictions
|
||||
once via the CLI, then run every diagnostic below against that file::
|
||||
|
||||
from giant.analysis import load_model_bundle, make_val_loader, collect_samples
|
||||
from giant.analysis import load_predicted_local
|
||||
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)
|
||||
samples = load_predicted_local("path/to/steps_predicted_local.parquet")
|
||||
|
||||
plot_marginals(samples, group_by="energy")
|
||||
plot_kl_bars(samples, group_by="energy")
|
||||
@@ -18,12 +19,6 @@ Typical use from a Jupyter notebook::
|
||||
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.)
|
||||
|
||||
@@ -77,11 +72,6 @@ Four tiers of checks, building on the aggregate marginal/KL check in
|
||||
surface covariate-shift failures that only appear under true rollout, only
|
||||
how well one-step generation reconstructs aggregate shower structure when
|
||||
fed real conditioning throughout.
|
||||
|
||||
`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
|
||||
@@ -93,30 +83,20 @@ 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,
|
||||
energy_simplex_decode,
|
||||
inv_log_transform,
|
||||
reconstruct_post_pos,
|
||||
)
|
||||
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
|
||||
|
||||
# The first 3 target dims are the scalar (non-direction) outputs. In raw/physical
|
||||
@@ -191,27 +171,6 @@ def _hist_edges(*arrays: np.ndarray, bins: int) -> np.ndarray:
|
||||
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)
|
||||
@@ -219,146 +178,6 @@ class SampleCollection:
|
||||
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, pre_E: np.ndarray
|
||||
) -> np.ndarray:
|
||||
raw = normalizer.inverse_transform(target_norm)
|
||||
return _decode_raw_targets(raw, pre_E)
|
||||
|
||||
|
||||
@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, cond_cont_raw[:, 3]
|
||||
),
|
||||
gen_raw=_to_raw_targets(
|
||||
gen_norm, bundle.target_normalizer, cond_cont_raw[:, 3]
|
||||
),
|
||||
real_norm=real_norm,
|
||||
gen_norm=gen_norm,
|
||||
)
|
||||
|
||||
|
||||
def _check_predict_metadata(path: Path) -> None:
|
||||
@@ -405,39 +224,53 @@ _COND_CONT_COLS = [
|
||||
|
||||
|
||||
def load_predicted_local(
|
||||
path: str | Path, sample_frac: float = 1.0, seed: int = 0
|
||||
path: str | Path,
|
||||
sample_frac: float = 1.0,
|
||||
seed: int = 0,
|
||||
batch_size: int = 1_000_000,
|
||||
) -> 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.
|
||||
since the predict CLI already denormalized them into this log-scaled,
|
||||
local-frame space. 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.
|
||||
Reads via `_iter_predicted_local_batches` (pyarrow's row-batch reader)
|
||||
with the needed columns selected, so at most `batch_size` rows are ever
|
||||
materialized at once — a plain lazy-polars `.collect()` with the sampling
|
||||
filter applied afterward looks lazy but doesn't push the row reduction
|
||||
into the scan (see its `.explain()`), so it still peaks at the full
|
||||
file's memory footprint even for a small `sample_frac`; streaming keeps
|
||||
peak memory to one batch regardless of file size or `sample_frac`.
|
||||
|
||||
`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.
|
||||
comfortably hold as numpy arrays in `real_raw`/`gen_raw`. The kept/dropped
|
||||
decision is a hash of each row's position in the file (`seed`-dependent),
|
||||
computed against a running offset across batches so it's equivalent to
|
||||
hashing a single row index over the whole file rather than restarting at
|
||||
each batch boundary.
|
||||
"""
|
||||
if not (0 < sample_frac <= 1):
|
||||
raise ValueError(f"sample_frac must be in (0, 1], got {sample_frac}")
|
||||
_check_predict_metadata(Path(path))
|
||||
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()
|
||||
)
|
||||
if sample_frac < 1.0:
|
||||
df = df.sample(fraction=sample_frac, seed=seed)
|
||||
columns = pred_cols + true_cols + _COND_CONT_COLS + ["pdg", "material"]
|
||||
|
||||
threshold = int(sample_frac * 2**32) if sample_frac < 1.0 else None
|
||||
frames = []
|
||||
offset = 0
|
||||
for batch_df in _iter_predicted_local_batches(path, columns, batch_size):
|
||||
n = batch_df.height
|
||||
if threshold is not None:
|
||||
row_idx = pl.arange(offset, offset + n, eager=True).cast(pl.UInt32)
|
||||
batch_df = batch_df.filter((row_idx.hash(seed=seed) % 2**32) < threshold)
|
||||
offset += n
|
||||
frames.append(batch_df)
|
||||
df = pl.concat(frames) if len(frames) != 1 else frames[0]
|
||||
|
||||
gen_log_local = df.select(pred_cols).to_numpy().astype(np.float32)
|
||||
real_log_local = df.select(true_cols).to_numpy().astype(np.float32)
|
||||
@@ -528,27 +361,30 @@ def marginal_table(
|
||||
#
|
||||
# `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.
|
||||
# below instead take a parquet path (or LazyFrame) directly and stay lazy,
|
||||
# reading the file a small constant number of times — one native `group_by`
|
||||
# (covering every group and every dim's mean/std/n/histogram-range at once)
|
||||
# plus one more pass per dim for histogram bin counts — rather than filtering
|
||||
# and re-collecting once per (group, dim) pair. The latter (the previous
|
||||
# implementation) meant runtime scaled with the number of *groups*: on a
|
||||
# 114M-row file with 138 distinct pdg codes, `group_by="pdg"` extrapolated to
|
||||
# roughly an hour, against ~1 minute for `group_by=None`. This version's
|
||||
# runtime is independent of group cardinality.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
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 _kl_from_counts(real_counts: np.ndarray, gen_counts: np.ndarray, eps: float = 1e-8) -> float:
|
||||
"""KL(real || gen) from two aligned histogram bin-count arrays.
|
||||
|
||||
Same smoothing/normalization as `giant.validate._histogram_kl`, just
|
||||
taking counts directly instead of raw samples (the counts here come from
|
||||
a polars `group_by` aggregation, not `np.histogram`).
|
||||
"""
|
||||
p = real_counts.astype(np.float64) + eps
|
||||
q = gen_counts.astype(np.float64) + eps
|
||||
p /= p.sum()
|
||||
q /= q.sum()
|
||||
return float(np.sum(p * np.log(p / q)))
|
||||
|
||||
|
||||
def _raw_dim_expr(prefix: str, j: int) -> pl.Expr:
|
||||
@@ -575,33 +411,145 @@ def _scan_predicted_local(source: str | Path | pl.LazyFrame) -> pl.LazyFrame:
|
||||
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]]:
|
||||
def _add_group_label(
|
||||
lf: pl.LazyFrame, group_by: str | None, n_energy_bins: int
|
||||
) -> pl.LazyFrame:
|
||||
"""Add a `_group` string column matching `_group_labels`'s label format.
|
||||
|
||||
For "pdg"/"material" the label is a direct string expr over the existing
|
||||
column — no upfront pass to enumerate distinct values needed, since
|
||||
`group_by("_group")` downstream discovers them itself. "energy" still
|
||||
needs one pass over `pre_E` to fix quantile bin edges before a label can
|
||||
be assigned per row.
|
||||
"""
|
||||
if group_by is None:
|
||||
return [("all", pl.lit(True))]
|
||||
return lf.with_columns(pl.lit("all").alias("_group"))
|
||||
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]
|
||||
return lf.with_columns(
|
||||
(pl.lit("pdg=") + pl.col("pdg").cast(pl.Int64).cast(pl.Utf8)).alias("_group")
|
||||
)
|
||||
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]
|
||||
return lf.with_columns(
|
||||
(pl.lit("material=") + pl.col("material")).alias("_group")
|
||||
)
|
||||
if group_by == "energy":
|
||||
pre_E = lf.select("pre_E").collect().to_series().to_numpy()
|
||||
pre_E = lf.select("pre_E").collect(engine="streaming").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)
|
||||
]
|
||||
labels = [f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})" for i in range(n_energy_bins)]
|
||||
expr = pl.when(pl.col("pre_E") < edges[1]).then(pl.lit(labels[0]))
|
||||
for i in range(1, n_energy_bins - 1):
|
||||
expr = expr.when(pl.col("pre_E") < edges[i + 1]).then(pl.lit(labels[i]))
|
||||
expr = expr.otherwise(pl.lit(labels[-1]))
|
||||
return lf.with_columns(expr.alias("_group"))
|
||||
raise ValueError(f"unknown group_by={group_by!r}")
|
||||
|
||||
|
||||
def _dim_narrow_lf(lf: pl.LazyFrame, j: int) -> pl.LazyFrame:
|
||||
"""Project down to just `(_group, real, gen)` for one dim, before any join/group_by.
|
||||
|
||||
Polars' projection pushdown doesn't reliably prune columns across a
|
||||
`.join()` in this version — without this explicit `.select()` up front,
|
||||
`_dim_hist_counts`'s join ends up materializing every column of the file
|
||||
(~30 columns × 114M rows ≈ 13GB) instead of just the ~2-3 this dim needs,
|
||||
even though nothing downstream references the others. Selecting first
|
||||
guarantees the narrow projection regardless of what the optimizer would
|
||||
otherwise infer.
|
||||
"""
|
||||
return lf.select(
|
||||
"_group",
|
||||
_raw_dim_expr("true", j).alias("real"),
|
||||
_raw_dim_expr("pred", j).alias("gen"),
|
||||
)
|
||||
|
||||
|
||||
def _dim_stats(narrow: pl.LazyFrame) -> pl.DataFrame:
|
||||
"""Per-group mean/std/n plus histogram-range (min/max of real+gen) for one dim.
|
||||
|
||||
A single `group_by("_group")` pass over `narrow` (see `_dim_narrow_lf`),
|
||||
mirroring the proven-cheap shape of `pdg_contribution_table_pl`, rather
|
||||
than folding all 9 dims' source columns into one query — the latter reads
|
||||
~19 columns and builds every dim's intermediate (softmax, log, etc.)
|
||||
arrays for the whole file at once, which is memory-heavy enough to OOM on
|
||||
a 114M-row file even though it's still just one pass.
|
||||
"""
|
||||
return (
|
||||
narrow.group_by("_group")
|
||||
.agg(
|
||||
pl.len().alias("n"),
|
||||
pl.col("real").mean().alias("real_mean"),
|
||||
pl.col("gen").mean().alias("gen_mean"),
|
||||
# ddof=0 to match numpy's (population-std) default used by marginal_table
|
||||
pl.col("real").std(ddof=0).alias("real_std"),
|
||||
pl.col("gen").std(ddof=0).alias("gen_std"),
|
||||
pl.min_horizontal(pl.col("real").min(), pl.col("gen").min()).alias("lo"),
|
||||
pl.max_horizontal(pl.col("real").max(), pl.col("gen").max()).alias("hi"),
|
||||
)
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
|
||||
|
||||
def _pad_degenerate_range(lo_hi: pl.DataFrame) -> pl.DataFrame:
|
||||
"""Widen a (lo, hi) pair by ±0.5 when too tight to support `bins` distinct edges.
|
||||
|
||||
Polars duplicate of `_hist_edges`'s degenerate-range handling, applied
|
||||
per group row instead of per call.
|
||||
"""
|
||||
ok = (pl.col("hi") - pl.col("lo")) > 1e-6 * pl.max_horizontal(
|
||||
pl.col("hi").abs(), pl.lit(1.0)
|
||||
)
|
||||
return lo_hi.with_columns(
|
||||
pl.when(ok).then(pl.col("lo")).otherwise(pl.col("lo") - 0.5).alias("lo"),
|
||||
pl.when(ok).then(pl.col("hi")).otherwise(pl.col("hi") + 0.5).alias("hi"),
|
||||
)
|
||||
|
||||
|
||||
def _dim_hist_counts(
|
||||
narrow: pl.LazyFrame, lo_hi: pl.DataFrame, bins: int
|
||||
) -> tuple[pl.DataFrame, pl.DataFrame]:
|
||||
"""Per-group real/gen histogram bin counts for one dim, as two long tables.
|
||||
|
||||
Joins each row (from `narrow`, see `_dim_narrow_lf`) to its group's
|
||||
(lo, hi) range (from `lo_hi`, already collected and tiny — one row per
|
||||
group), bins real/gen into `[0, bins)`, and counts via a two-key
|
||||
`group_by(["_group", "bin"])` — a proper single hash-pass histogram.
|
||||
(An earlier version aggregated with one `(bin == k).sum()` expression per
|
||||
bin, i.e. `bins` separate boolean-compare-and-reduce passes over every
|
||||
row; that's O(N × bins) work — 50 bins meant ~50x more comparisons than
|
||||
necessary — and was the dominant cost, not the join.)
|
||||
"""
|
||||
width = pl.col("hi") - pl.col("lo")
|
||||
real_bin = (
|
||||
((pl.col("real") - pl.col("lo")) / width * bins).floor().cast(pl.Int64).clip(0, bins - 1)
|
||||
)
|
||||
gen_bin = (
|
||||
((pl.col("gen") - pl.col("lo")) / width * bins).floor().cast(pl.Int64).clip(0, bins - 1)
|
||||
)
|
||||
|
||||
joined = narrow.join(lo_hi.lazy(), on="_group")
|
||||
real_hist = (
|
||||
joined.select("_group", real_bin.alias("bin"))
|
||||
.group_by(["_group", "bin"])
|
||||
.agg(pl.len().alias("count"))
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
gen_hist = (
|
||||
joined.select("_group", gen_bin.alias("bin"))
|
||||
.group_by(["_group", "bin"])
|
||||
.agg(pl.len().alias("count"))
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
return real_hist, gen_hist
|
||||
|
||||
|
||||
def _hist_counts_by_group(hist: pl.DataFrame, bins: int) -> dict[str, np.ndarray]:
|
||||
"""Long `(_group, bin, count)` table -> `{group: dense (bins,) count array}`."""
|
||||
out: dict[str, np.ndarray] = {}
|
||||
for group, bin_idx, count in hist.iter_rows():
|
||||
out.setdefault(group, np.zeros(bins, dtype=np.int64))[bin_idx] = count
|
||||
return out
|
||||
|
||||
|
||||
def marginal_table_pl(
|
||||
source: str | Path | pl.LazyFrame,
|
||||
group_by: str | None = None,
|
||||
@@ -613,35 +561,35 @@ def marginal_table_pl(
|
||||
`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.
|
||||
see the module-level note above on why this stays lazy, and on why this
|
||||
reads the file three passes per dim (stats, real histogram, gen
|
||||
histogram) rather than once per (group, dim) pair.
|
||||
"""
|
||||
lf = _scan_predicted_local(source)
|
||||
lf = _add_group_label(_scan_predicted_local(source), group_by, n_energy_bins)
|
||||
|
||||
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_dim_expr("true", j).alias("real"),
|
||||
_raw_dim_expr("pred", j).alias("gen"),
|
||||
]
|
||||
).collect()
|
||||
real_s, gen_s = pair["real"], pair["gen"]
|
||||
for j, name in enumerate(RAW_TARGET_NAMES):
|
||||
narrow = _dim_narrow_lf(lf, j)
|
||||
stats = _dim_stats(narrow)
|
||||
stats = stats.filter(pl.col("n") >= 2)
|
||||
lo_hi = _pad_degenerate_range(stats.select("_group", "lo", "hi"))
|
||||
real_hist, gen_hist = _dim_hist_counts(narrow, lo_hi, bins)
|
||||
real_by_group = _hist_counts_by_group(real_hist, bins)
|
||||
gen_by_group = _hist_counts_by_group(gen_hist, bins)
|
||||
for row in stats.iter_rows(named=True):
|
||||
group = row["_group"]
|
||||
real_counts = real_by_group.get(group, np.zeros(bins, dtype=np.int64))
|
||||
gen_counts = gen_by_group.get(group, np.zeros(bins, dtype=np.int64))
|
||||
rows.append(
|
||||
{
|
||||
"group": label,
|
||||
"group": group,
|
||||
"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),
|
||||
"n": row["n"],
|
||||
"real_mean": row["real_mean"],
|
||||
"gen_mean": row["gen_mean"],
|
||||
"real_std": row["real_std"],
|
||||
"gen_std": row["gen_std"],
|
||||
"kl_real_gen": _kl_from_counts(real_counts, gen_counts),
|
||||
}
|
||||
)
|
||||
return pl.DataFrame(rows).sort("kl_real_gen", descending=True)
|
||||
@@ -987,7 +935,7 @@ def constraint_report_pl(
|
||||
],
|
||||
]
|
||||
)
|
||||
.collect()
|
||||
.collect(engine="streaming")
|
||||
.row(0, named=True)
|
||||
)
|
||||
|
||||
@@ -1105,11 +1053,11 @@ def _entry_axis_and_bin_edges(
|
||||
pl.col("depth_proxy").quantile(0.999).alias("depth_hi"),
|
||||
pl.col("transverse_proxy").quantile(0.999).alias("transverse_hi"),
|
||||
)
|
||||
.collect()
|
||||
.collect(engine="streaming")
|
||||
.row(0, named=True)
|
||||
)
|
||||
|
||||
entry_df = entry.collect().sort("event_id")
|
||||
entry_df = entry.collect(engine="streaming").sort("event_id")
|
||||
|
||||
depth_lo, depth_hi = stats["depth_lo"], stats["depth_hi"]
|
||||
if not (depth_hi - depth_lo > 1e-6 * max(abs(depth_hi), 1.0)):
|
||||
@@ -1252,19 +1200,48 @@ def compute_event_observables_pl(
|
||||
real_transverse_bin = np.digitize(real_transverse, transverse_edges[1:-1])
|
||||
gen_transverse_bin = np.digitize(gen_transverse, transverse_edges[1:-1])
|
||||
|
||||
np.add.at(n_steps, idx, 1)
|
||||
np.add.at(real_total_edep, idx, real_edep)
|
||||
np.add.at(gen_total_edep, idx, gen_edep)
|
||||
np.add.at(real_total_length, idx, real_step_length)
|
||||
np.add.at(gen_total_length, idx, gen_step_length)
|
||||
np.add.at(real_sum_edep_depth, idx, real_edep * real_depth)
|
||||
np.add.at(gen_sum_edep_depth, idx, gen_edep * gen_depth)
|
||||
np.add.at(real_sum_edep_transverse2, idx, real_edep * real_transverse**2)
|
||||
np.add.at(gen_sum_edep_transverse2, idx, gen_edep * gen_transverse**2)
|
||||
np.add.at(real_depth_bin_edep, (idx, real_depth_bin), real_edep)
|
||||
np.add.at(gen_depth_bin_edep, (idx, gen_depth_bin), gen_edep)
|
||||
np.add.at(real_transverse_bin_edep, (idx, real_transverse_bin), real_edep)
|
||||
np.add.at(gen_transverse_bin_edep, (idx, gen_transverse_bin), gen_edep)
|
||||
# np.add.at is an unbuffered ufunc method — a well-known numpy slow path
|
||||
# for scatter-add (44% of this function's runtime on a 114M-row profile).
|
||||
# np.bincount does the same accumulation with a single optimized pass;
|
||||
# the 2D (per-event, per-bin) accumulators flatten (idx, bin) into one
|
||||
# bincount index and reshape back, since bincount only scatters into 1D.
|
||||
n_steps += np.bincount(idx, minlength=n_events)
|
||||
real_total_edep += np.bincount(idx, weights=real_edep, minlength=n_events)
|
||||
gen_total_edep += np.bincount(idx, weights=gen_edep, minlength=n_events)
|
||||
real_total_length += np.bincount(idx, weights=real_step_length, minlength=n_events)
|
||||
gen_total_length += np.bincount(idx, weights=gen_step_length, minlength=n_events)
|
||||
real_sum_edep_depth += np.bincount(
|
||||
idx, weights=real_edep * real_depth, minlength=n_events
|
||||
)
|
||||
gen_sum_edep_depth += np.bincount(
|
||||
idx, weights=gen_edep * gen_depth, minlength=n_events
|
||||
)
|
||||
real_sum_edep_transverse2 += np.bincount(
|
||||
idx, weights=real_edep * real_transverse**2, minlength=n_events
|
||||
)
|
||||
gen_sum_edep_transverse2 += np.bincount(
|
||||
idx, weights=gen_edep * gen_transverse**2, minlength=n_events
|
||||
)
|
||||
real_depth_bin_edep += np.bincount(
|
||||
idx * depth_bins + real_depth_bin,
|
||||
weights=real_edep,
|
||||
minlength=n_events * depth_bins,
|
||||
).reshape(n_events, depth_bins)
|
||||
gen_depth_bin_edep += np.bincount(
|
||||
idx * depth_bins + gen_depth_bin,
|
||||
weights=gen_edep,
|
||||
minlength=n_events * depth_bins,
|
||||
).reshape(n_events, depth_bins)
|
||||
real_transverse_bin_edep += np.bincount(
|
||||
idx * transverse_bins + real_transverse_bin,
|
||||
weights=real_edep,
|
||||
minlength=n_events * transverse_bins,
|
||||
).reshape(n_events, transverse_bins)
|
||||
gen_transverse_bin_edep += np.bincount(
|
||||
idx * transverse_bins + gen_transverse_bin,
|
||||
weights=gen_edep,
|
||||
minlength=n_events * transverse_bins,
|
||||
).reshape(n_events, transverse_bins)
|
||||
|
||||
safe_real_total = np.where(real_total_edep > 0, real_total_edep, 1.0)
|
||||
safe_gen_total = np.where(gen_total_edep > 0, gen_total_edep, 1.0)
|
||||
@@ -1313,7 +1290,7 @@ def compute_event_observables_pl(
|
||||
.alias("gen_median_length"),
|
||||
)
|
||||
.sort("event_id")
|
||||
.collect()
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
real_median_edep = medians["real_median_edep"].to_numpy()
|
||||
gen_median_edep = medians["gen_median_edep"].to_numpy()
|
||||
@@ -1628,7 +1605,7 @@ def pdg_contribution_table_pl(source: str | Path | pl.LazyFrame) -> pl.DataFrame
|
||||
_delog("true_log_step_length").sum().alias("real_total_length"),
|
||||
_delog("pred_log_step_length").sum().alias("gen_total_length"),
|
||||
)
|
||||
.collect()
|
||||
.collect(engine="streaming")
|
||||
.sort("pdg")
|
||||
)
|
||||
|
||||
|
||||
+34
-21
@@ -80,6 +80,34 @@ def energy_simplex_decode(
|
||||
)
|
||||
|
||||
|
||||
def _rodrigues_axis(pre_dir: np.ndarray) -> np.ndarray:
|
||||
"""Unit rotation axis `pre_dir × ẑ`, closed-form since ẑ = [0, 0, 1] is constant.
|
||||
|
||||
`cross(a, [0,0,1]) = [a_y, -a_x, 0]` — substituting the constant operand
|
||||
avoids a generic `np.cross` call (shape/broadcast handling for an
|
||||
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_norm = np.linalg.norm(axis, axis=1, keepdims=True) # (N,1)
|
||||
# Replace zero-norm axes with x̂ (the Rodrigues terms that involve the axis
|
||||
# are multiplied by sin_t≈0 and (1-cos_t)≈0, so the choice is irrelevant).
|
||||
safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
|
||||
return np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)
|
||||
|
||||
|
||||
def _cross_with_z_axis(axis: np.ndarray, v: np.ndarray) -> np.ndarray:
|
||||
"""`axis × v`, closed-form since `axis` from `_rodrigues_axis` always has z = 0.
|
||||
|
||||
`cross([ax,ay,0], [bx,by,bz]) = [ay*bz, -ax*bz, ax*by - ay*bx]`.
|
||||
"""
|
||||
ax, ay = axis[:, 0:1], axis[:, 1:2]
|
||||
bx, by, bz = v[:, 0:1], v[:, 1:2], v[:, 2:3]
|
||||
return np.concatenate([ay * bz, -ax * bz, ax * by - ay * bx], axis=1)
|
||||
|
||||
|
||||
def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarray:
|
||||
"""Rotate post_dir into the local frame where pre_dir maps to ẑ (Rodrigues).
|
||||
|
||||
@@ -87,20 +115,11 @@ def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarra
|
||||
expressed relative to a coordinate system in which the incoming particle
|
||||
travels along +z.
|
||||
"""
|
||||
z = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
cos_t = np.clip((pre_dir * z).sum(axis=1, keepdims=True), -1.0, 1.0) # (N,1)
|
||||
cos_t = np.clip(pre_dir[:, 2:3], -1.0, 1.0) # (N,1); dot with ẑ = z-component
|
||||
sin_t = np.sqrt(np.maximum(0.0, 1.0 - cos_t**2)) # (N,1)
|
||||
|
||||
axis = np.cross(pre_dir, z) # (N,3); zero when pre_dir ∥ ẑ
|
||||
axis_norm = np.linalg.norm(axis, axis=1, keepdims=True) # (N,1)
|
||||
|
||||
# Replace zero-norm axes with x̂ (the Rodrigues terms that involve the axis
|
||||
# are multiplied by sin_t≈0 and (1-cos_t)≈0, so the choice is irrelevant).
|
||||
safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
|
||||
axis = np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)
|
||||
|
||||
kxv = np.cross(axis, post_dir) # (N,3)
|
||||
axis = _rodrigues_axis(pre_dir) # (N,3); zero-z, zero-norm when pre_dir ∥ ẑ
|
||||
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(
|
||||
@@ -208,17 +227,11 @@ def inv_local_frame_rotation(
|
||||
|
||||
Applies R^T (same axis, negative angle) to post_dir_local.
|
||||
"""
|
||||
z = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
cos_t = np.clip((pre_dir * z).sum(axis=1, keepdims=True), -1.0, 1.0)
|
||||
cos_t = np.clip(pre_dir[:, 2:3], -1.0, 1.0) # dot with ẑ = z-component
|
||||
sin_t = np.sqrt(np.maximum(0.0, 1.0 - cos_t**2))
|
||||
|
||||
axis = np.cross(pre_dir, z)
|
||||
axis_norm = np.linalg.norm(axis, axis=1, keepdims=True)
|
||||
safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
|
||||
axis = np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)
|
||||
|
||||
kxv = np.cross(axis, post_dir_local)
|
||||
axis = _rodrigues_axis(pre_dir)
|
||||
kxv = _cross_with_z_axis(axis, post_dir_local)
|
||||
kdv = (axis * post_dir_local).sum(axis=1, keepdims=True)
|
||||
|
||||
# Negative angle: sin_t → -sin_t
|
||||
|
||||
@@ -84,8 +84,6 @@ def _make_collection(n=200, seed=0, gen_offset=0.0) -> SampleCollection:
|
||||
material=rng.choice(["W", "Pb"], size=n),
|
||||
real_raw=real,
|
||||
gen_raw=gen,
|
||||
real_norm=real,
|
||||
gen_norm=gen,
|
||||
)
|
||||
|
||||
|
||||
@@ -259,8 +257,6 @@ def test_load_predicted_local_round_trips_values(tmp_path):
|
||||
np.testing.assert_allclose(
|
||||
collection.gen_raw, expected_raw(pred_log_local), atol=1e-4
|
||||
)
|
||||
assert collection.real_norm is None
|
||||
assert collection.gen_norm is None
|
||||
|
||||
|
||||
def test_load_predicted_local_usable_by_downstream_plots(tmp_path):
|
||||
|
||||
@@ -128,7 +128,7 @@ def test_energy_simplex_conservation():
|
||||
def test_energy_simplex_roundtrip():
|
||||
"""Encode → decode recovers energies whose lost part already sums to delta_e."""
|
||||
rng = np.random.default_rng(12)
|
||||
N = 500
|
||||
N = 500000
|
||||
pre_E = rng.uniform(1.0, 100.0, N).astype(np.float32)
|
||||
post_E = (pre_E * rng.uniform(0.0, 1.0, N)).astype(np.float32)
|
||||
delta_e = pre_E - post_E
|
||||
|
||||
Reference in New Issue
Block a user