giant/giant/analysis/reduce.py

"""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.

    Out-of-range values clamp into the edge bins, and the clamp deliberately
    happens in f64 *before* the integer cast: a rollout is free to emit a wildly
    out-of-range outlier (a step_length of 1e10 mm, say) or an inf, whose
    unclamped bin index overflows i32 and makes the cast fail outright. NaN has
    no edge to clamp to, so it becomes null and is dropped by the callers below
    — the same thing ``np.histogram`` does with it.
    """
    idx = ((value - lo) / (hi - lo) * nbins).floor().clip(0, nbins - 1)
    return pl.when(idx.is_nan()).then(None).otherwise(idx).cast(pl.Int32)


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"))
        .drop_nulls("_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"),
        )
        .drop_nulls("_b")
        .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)


def sec_count_by_event(lf_all: pl.LazyFrame, sec_lf: pl.LazyFrame) -> tuple[np.ndarray, np.ndarray]:
    """Per-event secondary count, zero-filled for events that produced none.

    Two bounded per-event ``group_by``s — the full event set (from ``lf_all``)
    and the secondary counts (from ``sec_lf``, see ``sources.secondaries``) —
    merged in Python via a dict. Both results are event-granularity (not
    per-row), so this stays in the same bounded-memory budget as
    ``event_scalars``; a plain ``group_by`` on ``sec_lf`` alone would silently
    drop zero-secondary events instead of zero-filling them.
    """
    ev = lf_all.select("event_id").unique().collect(engine="streaming")["event_id"].to_numpy()
    cnt_df = sec_lf.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")
    cnt = dict(zip(cnt_df["event_id"].to_list(), cnt_df["n"].to_list()))
    counts = np.array([cnt.get(int(e), 0) for e in ev], dtype=np.int64)
    return ev, counts