lars
55332db67a
CI / Lint (ruff check) (push) Successful in 31s
CI / Format (ruff format) (push) Successful in 32s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 26s
CI / Type check (ty) (push) Successful in 29s
CI / Format (ruff format) (pull_request) Successful in 33s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 35s
CI / Tests (pull_request) Successful in 3m47s
CI / Tests (push) Successful in 3m55s
Rejoins lines that only wrapped because they exceeded the old 88-char limit; ruff check and the full test suite (725 passed) are unaffected.
263 lines
10 KiB
Python
263 lines
10 KiB
Python
"""Streaming compute primitives — the reduce half of the analysis.
|
|
|
|
Everything here turns a (possibly larger-than-RAM) LazyFrame into a *compact*
|
|
numpy/DataFrame artifact in bounded memory, and never imports plotstyle so it can
|
|
run on an HTCondor worker. Efficiency rules (see the plan's "Histogram
|
|
efficiency" section):
|
|
|
|
* ``hist1d`` is a single streaming ``group_by([group, bin]).len()`` pass against
|
|
**fixed** edges (no min/max range pass) with a strict column projection — only
|
|
the columns the value/group expressions reference are read from the parquet.
|
|
* the multi-quantity reductions (``event_scalars``, profiles, ``species_share``,
|
|
``leakage_fraction``) each emit *all* their outputs from one ``group_by``.
|
|
* per-event -> per-row lookups (shower entry/axis) use ``replace_strict`` (a hash
|
|
map applied as an expression, bounded memory), never a streaming join.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import polars as pl
|
|
|
|
from giant.constants import TERM_ESCAPED
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 1-D histogram primitive
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _bin_expr(value: pl.Expr, lo: float, hi: float, nbins: int) -> pl.Expr:
|
|
"""Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins."""
|
|
return ((value - lo) / (hi - lo) * nbins).floor().cast(pl.Int32).clip(0, nbins - 1)
|
|
|
|
|
|
def hist1d(
|
|
lf: pl.LazyFrame,
|
|
value: pl.Expr,
|
|
edges: np.ndarray,
|
|
group: pl.Expr | None = None,
|
|
) -> dict[object, np.ndarray]:
|
|
"""Streaming histogram of ``value`` over fixed uniform ``edges``, by ``group``.
|
|
|
|
Returns ``{group_key: counts}`` (counts is an ``int64`` array of length
|
|
``len(edges)-1``). One hash pass; runtime is independent of group cardinality,
|
|
so every pdg/material/energy stratum falls out together. Only the tiny
|
|
``(n_groups x nbins)`` result is materialized.
|
|
"""
|
|
lo, hi, nbins = float(edges[0]), float(edges[-1]), len(edges) - 1
|
|
group = pl.lit(0, dtype=pl.Int64) if group is None else group
|
|
res = (
|
|
lf.select(group.alias("_g"), _bin_expr(value, lo, hi, nbins).alias("_b"))
|
|
.group_by("_g", "_b")
|
|
.agg(pl.len().alias("_n"))
|
|
.collect(engine="streaming")
|
|
)
|
|
out: dict[object, np.ndarray] = {}
|
|
for g, b, n in res.iter_rows():
|
|
out.setdefault(g, np.zeros(nbins, dtype=np.int64))[b] = n
|
|
return out
|
|
|
|
|
|
def sum_merge(dicts: list[dict[str, Any]]) -> dict[str, Any]:
|
|
"""Elementwise-sum a list of sum-mergeable count/total dicts (JSON-safe keys).
|
|
|
|
Used to merge chunked ``hist1d``/``species_share``-style partials, whose
|
|
values bin/group against edges or keys fixed by ``Context`` — a chunk's raw
|
|
count dict is exactly a partial sum, so merging is a plain elementwise sum
|
|
over the union of keys (a key absent from some chunk is all-zero there).
|
|
Values may be per-bin count lists or plain scalar totals; both round-trip
|
|
through ``np.asarray``/``.tolist()`` unchanged in shape.
|
|
"""
|
|
out: dict[str, np.ndarray] = {}
|
|
for d in dicts:
|
|
for k, v in d.items():
|
|
arr = np.asarray(v)
|
|
out[k] = arr.copy() if k not in out else out[k] + arr
|
|
return {k: v.tolist() for k, v in out.items()}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Per-event scalar observables (one bounded group_by pass)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def event_scalars(lf: pl.LazyFrame) -> pl.DataFrame:
|
|
"""One row per event: total/mean deposited energy, path length, step count.
|
|
|
|
Columns: ``event_id, total_edep, total_length, n_steps, mean_length,
|
|
incident_E`` (incident = ``max(pre_E)``, the primary). The caller chooses
|
|
whether ``lf`` includes the rollout's synthetic termination rows — pass the
|
|
full scan for energy totals (they carry the deposited remainder), physical
|
|
steps only for step-count / mean-length.
|
|
"""
|
|
return (
|
|
lf.group_by("event_id")
|
|
.agg(
|
|
pl.col("edep").sum().alias("total_edep"),
|
|
pl.col("step_length").sum().alias("total_length"),
|
|
pl.len().alias("n_steps"),
|
|
pl.col("pre_E").max().alias("incident_E"),
|
|
)
|
|
.with_columns((pl.col("total_length") / pl.col("n_steps")).alias("mean_length"))
|
|
.collect(engine="streaming")
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Shower shape: entry/axis + edep-weighted longitudinal / transverse profiles
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def entry_axis(lf: pl.LazyFrame) -> pl.DataFrame:
|
|
"""Per-event shower entry point and axis (from the highest-``pre_E`` step).
|
|
|
|
One bounded ``group_by``: the primary's ``pre_pos`` becomes the entry point
|
|
and its ``pre_dir`` the shower axis.
|
|
"""
|
|
return (
|
|
lf.group_by("event_id")
|
|
.agg(
|
|
pl.col("pre_x").get(pl.col("pre_E").arg_max()).alias("entry_x"),
|
|
pl.col("pre_y").get(pl.col("pre_E").arg_max()).alias("entry_y"),
|
|
pl.col("pre_z").get(pl.col("pre_E").arg_max()).alias("entry_z"),
|
|
pl.col("pre_dx").get(pl.col("pre_E").arg_max()).alias("axis_x"),
|
|
pl.col("pre_dy").get(pl.col("pre_E").arg_max()).alias("axis_y"),
|
|
pl.col("pre_dz").get(pl.col("pre_E").arg_max()).alias("axis_z"),
|
|
)
|
|
.collect(engine="streaming")
|
|
.sort("event_id")
|
|
)
|
|
|
|
|
|
_ENTRY_AXIS_COLS = ("entry_x", "entry_y", "entry_z", "axis_x", "axis_y", "axis_z")
|
|
|
|
|
|
def attach_entry_axis(lf: pl.LazyFrame, entry: pl.DataFrame) -> pl.LazyFrame:
|
|
"""Broadcast each event's entry/axis onto its rows via ``replace_strict``.
|
|
|
|
A hash map applied as an expression — streams in bounded memory, unlike a
|
|
join which would buffer the whole file-sized left side.
|
|
"""
|
|
ids = entry["event_id"].to_numpy()
|
|
return lf.with_columns(
|
|
pl.col("event_id").replace_strict(ids, entry[col].to_numpy(), return_dtype=pl.Float64).alias(col)
|
|
for col in _ENTRY_AXIS_COLS
|
|
)
|
|
|
|
|
|
def depth_expr() -> pl.Expr:
|
|
"""Signed distance of ``post_pos`` from the entry point along the shower axis."""
|
|
dx = pl.col("post_x") - pl.col("entry_x")
|
|
dy = pl.col("post_y") - pl.col("entry_y")
|
|
dz = pl.col("post_z") - pl.col("entry_z")
|
|
return dx * pl.col("axis_x") + dy * pl.col("axis_y") + dz * pl.col("axis_z")
|
|
|
|
|
|
def transverse_expr() -> pl.Expr:
|
|
"""Perpendicular distance of ``post_pos`` from the shower axis."""
|
|
dx = pl.col("post_x") - pl.col("entry_x")
|
|
dy = pl.col("post_y") - pl.col("entry_y")
|
|
dz = pl.col("post_z") - pl.col("entry_z")
|
|
depth = dx * pl.col("axis_x") + dy * pl.col("axis_y") + dz * pl.col("axis_z")
|
|
tx = dx - depth * pl.col("axis_x")
|
|
ty = dy - depth * pl.col("axis_y")
|
|
tz = dz - depth * pl.col("axis_z")
|
|
return (tx**2 + ty**2 + tz**2).sqrt()
|
|
|
|
|
|
def profile_partial(
|
|
lf: pl.LazyFrame,
|
|
coord: pl.Expr,
|
|
edges: np.ndarray,
|
|
weight: pl.Expr,
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
"""One chunk's per-event x bin ``weight``-sum matrix: ``(event_ids, matrix)``.
|
|
|
|
One streaming ``group_by(event_id, bin)`` sums ``weight`` per (event, bin).
|
|
A chunk's matrix rows are only the events present in that chunk, so chunks'
|
|
matrices stack cleanly with no cross-chunk lookup — this requires chunking
|
|
to be event-disjoint (every row of an event lands in one chunk).
|
|
``coord``/``weight`` require the entry/axis columns attached.
|
|
"""
|
|
lo, hi, nbins = float(edges[0]), float(edges[-1]), len(edges) - 1
|
|
grid = (
|
|
lf.select(
|
|
"event_id",
|
|
_bin_expr(coord, lo, hi, nbins).alias("_b"),
|
|
weight.alias("_w"),
|
|
)
|
|
.group_by("event_id", "_b")
|
|
.agg(pl.col("_w").sum().alias("_ws"))
|
|
.collect(engine="streaming")
|
|
)
|
|
ev = grid["event_id"].to_numpy()
|
|
uniq, inv = np.unique(ev, return_inverse=True)
|
|
mat = np.zeros((len(uniq), nbins), dtype=np.float64)
|
|
np.add.at(mat, (inv, grid["_b"].to_numpy()), grid["_ws"].to_numpy())
|
|
return uniq, mat
|
|
|
|
|
|
def profile_finalize(mats: list[np.ndarray]) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Collapse per-chunk per-event x bin matrices into the final mean/std profile.
|
|
|
|
Chunks are event-disjoint, so row-wise concatenation of their matrices
|
|
reconstructs the full per-event matrix; the mean/event-RMS collapse must
|
|
happen once over that full matrix — an average of per-chunk means/stds
|
|
would be wrong (chunks generally hold different numbers of events).
|
|
Returns ``(mean, std)``, each length ``nbins``.
|
|
"""
|
|
full = np.concatenate(mats, axis=0)
|
|
return full.mean(axis=0), full.std(axis=0)
|
|
|
|
|
|
def weighted_profile(
|
|
lf: pl.LazyFrame,
|
|
coord: pl.Expr,
|
|
edges: np.ndarray,
|
|
weight: pl.Expr,
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Single-pass profile of ``coord`` (mean +/- event-RMS band over events).
|
|
|
|
Convenience wrapper for the unchunked (whole-dataset) case; ``mean_std +
|
|
profile_partial`` is what a chunked compute/finalize split uses instead.
|
|
"""
|
|
_, mat = profile_partial(lf, coord, edges, weight)
|
|
return profile_finalize([mat])
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Species contribution and leakage
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def species_share(lf: pl.LazyFrame) -> pl.DataFrame:
|
|
"""Total deposited energy per PDG species (``pdg, total_edep``), one pass."""
|
|
return (
|
|
lf.group_by("pdg")
|
|
.agg(pl.col("edep").sum().alias("total_edep"))
|
|
.collect(engine="streaming")
|
|
.sort("total_edep", descending=True)
|
|
)
|
|
|
|
|
|
def leakage_fraction(lf: pl.LazyFrame) -> np.ndarray:
|
|
"""Per-event escaped-energy fraction (rollout only), one bounded pass.
|
|
|
|
Escaped rows carry the leaked energy in ``pre_E`` (``edep`` is 0 there); the
|
|
fraction is ``escaped / (deposited + escaped)`` per event.
|
|
"""
|
|
per_event = (
|
|
lf.group_by("event_id")
|
|
.agg(
|
|
pl.col("edep").sum().alias("deposited"),
|
|
pl.col("pre_E").filter(pl.col("termination_reason") == TERM_ESCAPED).sum().alias("escaped"),
|
|
)
|
|
.collect(engine="streaming")
|
|
)
|
|
deposited = per_event["deposited"].to_numpy()
|
|
escaped = per_event["escaped"].fill_null(0.0).to_numpy()
|
|
total = deposited + escaped
|
|
return np.where(total > 0, escaped / total, 0.0)
|