giant/giant/analysis/grouping.py

"""Grouping axes (overall / energy / pdg / material) and their labels.

Energy grouping is by the event's **incident (primary) energy** — the largest
``pre_E`` in the event — so every step of a shower lands in one bin, the physically
meaningful stratification for a calorimeter surrogate. The quantile bin *edges* are
sized once in ``prep`` (from a subsample) and shipped in ``shared.json``; a compute
job that needs them re-derives the small per-event ``event_id -> bin`` map itself
(one bounded streaming ``group_by`` over ``pre_E``), so ``shared.json`` stays tiny.

Pure/plotstyle-free so it can run on the compute workers.
"""

from __future__ import annotations

import numpy as np
import polars as pl

# Common electromagnetic/hadronic species; anything else falls back to its code.
PDG_NAMES: dict[int, str] = {
    11: "e-",
    -11: "e+",
    22: "gamma",
    2112: "n",
    2212: "p",
    -2212: "pbar",
    111: "pi0",
    211: "pi+",
    -211: "pi-",
    13: "mu-",
    -13: "mu+",
    321: "K+",
    -321: "K-",
    130: "K0L",
}


def pdg_label(code: int) -> str:
    """Human-readable species label for a PDG code (falls back to the code)."""
    code = int(code)
    if code in PDG_NAMES:
        return PDG_NAMES[code]
    if abs(code) > 1_000_000_000:
        return f"ion {code}"
    return str(code)


def material_label(name: str) -> str:
    """Display label for a Geant4 material, dropping the ``G4_`` prefix."""
    return name.removeprefix("G4_")


def energy_bin_edges(incident_E: np.ndarray, n_bins: int = 4) -> np.ndarray:
    """Equal-population (quantile) bin edges over per-event incident energies.

    Returns ``n_bins + 1`` monotonically non-decreasing edges. The top edge is
    nudged up so the largest value falls inside the last bin under a
    right-open convention. Degenerate (single-value) input widens by +/-0.5.
    """
    incident_E = np.asarray(incident_E, dtype=np.float64)
    edges = np.quantile(incident_E, np.linspace(0.0, 1.0, n_bins + 1))
    edges = np.unique(edges)
    if edges.size < 2:
        v = edges[0] if edges.size else 0.0
        edges = np.array([v - 0.5, v + 0.5])
    edges[-1] = np.nextafter(edges[-1], np.inf)
    return edges


def energy_bin_labels(edges: np.ndarray) -> list[str]:
    """``E in [lo, hi)`` labels for each bin defined by ``edges`` (MeV)."""
    return [f"E in [{edges[i]:.3g}, {edges[i + 1]:.3g}) MeV" for i in range(len(edges) - 1)]


def digitize_expr(value: pl.Expr, edges: np.ndarray) -> pl.Expr:
    """Bin index of ``value`` under arbitrary (possibly non-uniform) ``edges``.

    ``bin = (#interior edges <= value)``, clipped to ``[0, n_bins-1]`` — matches
    ``np.digitize(value, edges[1:-1])`` and works for the quantile energy edges.
    Vectorized as a sum of boolean comparisons; no per-row Python.
    """
    interior = [float(e) for e in edges[1:-1]]
    n_bins = len(edges) - 1
    idx = pl.lit(0, dtype=pl.Int32)
    for e in interior:
        idx = idx + (value >= e).cast(pl.Int32)
    return idx.clip(0, n_bins - 1)


def event_energy_bins(lf: pl.LazyFrame, edges: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Per-event incident-energy bin: ``(event_ids, bin_idx)`` numpy arrays.

    Incident energy is ``max(pre_E)`` per event (the primary). One bounded
    streaming ``group_by``; the tiny per-event result is digitized in numpy.
    """
    per_event = (
        lf.group_by("event_id")
        .agg(pl.col("pre_E").max().alias("incident_E"))
        .collect(engine="streaming")
        .sort("event_id")
    )
    event_ids = per_event["event_id"].to_numpy()
    incident = per_event["incident_E"].to_numpy()
    bin_idx = np.clip(np.digitize(incident, edges[1:-1]), 0, len(edges) - 2)
    return event_ids, bin_idx.astype(np.int64)