giant/giant/analysis/context.py

"""Shared analysis context (the ``prep`` step): fixed bin edges + group sets.

Every histogram in the catalog bins against **fixed** edges so each compute job
is a single streaming pass with no min/max range scan. Those edges — plus the
energy-bin quantiles, the top PDG species and the material list to stratify by,
and the shower depth/transverse ranges — are resolved *once* here, on the submit
node, from a hash-subsample plus a few cheap exact ``group_by`` passes, and shipped
in ``shared.json``. Tiny and self-describing; no per-event arrays.

plotstyle-free (runs on the submit node, but also importable by workers).
"""

from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from pathlib import Path

import numpy as np
import polars as pl

from giant.analysis.grouping import energy_bin_edges
from giant.analysis.reduce import (
    attach_entry_axis,
    depth_expr,
    entry_axis,
    transverse_expr,
)
from giant.analysis.sources import Side, open_side, physical_steps, secondaries
from giant.analysis.variables import RANGED_VARS


@dataclass
class Context:
    """Resolved bin edges and grouping sets shared by every compute job."""

    n_marginal_bins: int
    var_ranges: dict[str, tuple[float, float]]  # ranged var -> (lo, hi)
    energy_edges: list[float]
    top_pdgs: list[int]
    materials: list[str]
    depth_edges: list[float]
    transverse_edges: list[float]
    sec_energy_range: tuple[float, float]
    n_sec_bins: int
    n_events: dict[str, int] = field(default_factory=dict)

    # -- (de)serialization -------------------------------------------------
    def save(self, path: str | Path) -> None:
        Path(path).write_text(json.dumps(asdict(self), indent=2))

    @classmethod
    def load(cls, path: str | Path) -> "Context":
        d = json.loads(Path(path).read_text())
        d["var_ranges"] = {k: tuple(v) for k, v in d["var_ranges"].items()}
        d["sec_energy_range"] = tuple(d["sec_energy_range"])
        return cls(**d)

    # -- convenience -------------------------------------------------------
    def marginal_edges(self, var: str) -> np.ndarray:
        lo, hi = self.var_ranges[var]
        return np.linspace(lo, hi, self.n_marginal_bins + 1)


_LO_Q, _HI_Q = 0.001, 0.999


def _row_subsample(lf: pl.LazyFrame, sample_rows: int, seed: int) -> pl.LazyFrame:
    """Hash-subsample ~``sample_rows`` rows (for range estimation only)."""
    n_total = lf.select(pl.len()).collect(engine="streaming").item()
    if n_total <= sample_rows:
        return lf
    threshold = int(sample_rows / n_total * 2**32)
    return lf.filter((pl.col("pre_E").hash(seed=seed) % 2**32) < threshold)


def _combined_quantiles(
    r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_q: float
) -> tuple[float, float]:
    """Robust (lo_q, hi_q) range over the union of two value samples."""
    both = np.concatenate([r_vals, t_vals])
    lo, hi = float(np.quantile(both, lo_q)), float(np.quantile(both, hi_q))
    if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
        lo, hi = lo - 0.5, hi + 0.5
    return lo, hi


def build_context(
    rollout: str | Path | pl.LazyFrame,
    reference: str | Path | pl.LazyFrame,
    *,
    n_energy_bins: int = 4,
    n_marginal_bins: int = 50,
    n_sec_bins: int = 40,
    top_k_pdg: int = 6,
    sample_rows: int = 1_000_000,
    seed: int = 0,
) -> Context:
    """Resolve the shared context from the two files (the ``prep`` step)."""
    r_all = open_side(rollout, Side.rollout)
    t_all = open_side(reference, Side.reference)
    r_lf = physical_steps(r_all, Side.rollout)
    t_lf = physical_steps(t_all, Side.reference)

    # Ranged marginal variables: robust ranges over a shared row subsample.
    exprs = [e.alias(n) for n, (_, e) in RANGED_VARS.items()]
    r_s = (
        _row_subsample(r_lf, sample_rows, seed)
        .select(exprs)
        .collect(engine="streaming")
    )
    t_s = (
        _row_subsample(t_lf, sample_rows, seed)
        .select(exprs)
        .collect(engine="streaming")
    )
    var_ranges = {
        name: _combined_quantiles(
            r_s[name].to_numpy(), t_s[name].to_numpy(), _LO_Q, _HI_Q
        )
        for name in RANGED_VARS
    }

    # Energy-bin edges from exact per-event incident energies (cheap group_by).
    def _incident(lf: pl.LazyFrame) -> np.ndarray:
        return (
            lf.group_by("event_id")
            .agg(pl.col("pre_E").max())
            .collect(engine="streaming")["pre_E"]
            .to_numpy()
        )

    r_inc, t_inc = _incident(r_lf), _incident(t_lf)
    energy_edges = energy_bin_edges(np.concatenate([r_inc, t_inc]), n_energy_bins)

    # Top PDG species and material list (cheap single-column group_bys).
    def _counts(lf: pl.LazyFrame, col: str) -> pl.DataFrame:
        return lf.group_by(col).agg(pl.len().alias("n")).collect(engine="streaming")

    pdg_counts = (
        pl.concat([_counts(r_lf, "pdg"), _counts(t_lf, "pdg")])
        .group_by("pdg")
        .agg(pl.col("n").sum())
        .sort("n", descending=True)
    )
    top_pdgs = [int(x) for x in pdg_counts["pdg"].to_list()[:top_k_pdg]]
    materials = sorted(
        set(_counts(r_lf, "material")["material"].to_list())
        | set(_counts(t_lf, "material")["material"].to_list())
    )

    # Shower depth / transverse ranges from a subsampled proxy.
    def _proxy(lf: pl.LazyFrame) -> tuple[np.ndarray, np.ndarray]:
        ea = entry_axis(lf)
        sub = (
            attach_entry_axis(_row_subsample(lf, sample_rows, seed), ea)
            .select(depth_expr().alias("d"), transverse_expr().alias("t"))
            .collect(engine="streaming")
        )
        return sub["d"].to_numpy(), sub["t"].to_numpy()

    r_d, r_t = _proxy(r_lf)
    t_d, t_t = _proxy(t_lf)
    d_lo, d_hi = _combined_quantiles(r_d, t_d, _LO_Q, _HI_Q)
    depth_edges = np.linspace(d_lo, d_hi, n_marginal_bins + 1)
    t_hi = max(float(np.quantile(np.concatenate([r_t, t_t]), _HI_Q)), 1e-6)
    transverse_edges = np.linspace(0.0, t_hi, n_marginal_bins + 1)

    # Secondary energy range.
    r_se = secondaries(r_lf, Side.rollout).select("energy")
    t_se = secondaries(t_all, Side.reference).select("energy")
    r_se = _row_sample_col(r_se, sample_rows, seed)
    t_se = _row_sample_col(t_se, sample_rows, seed)
    sec_energy_range = _combined_quantiles(r_se, t_se, _LO_Q, _HI_Q)

    return Context(
        n_marginal_bins=n_marginal_bins,
        var_ranges=var_ranges,
        energy_edges=[float(x) for x in energy_edges],
        top_pdgs=top_pdgs,
        materials=materials,
        depth_edges=[float(x) for x in depth_edges],
        transverse_edges=[float(x) for x in transverse_edges],
        sec_energy_range=sec_energy_range,
        n_sec_bins=n_sec_bins,
        n_events={
            "rollout": len(r_inc),
            "reference": len(t_inc),
        },
    )


def _row_sample_col(lf: pl.LazyFrame, sample_rows: int, seed: int) -> np.ndarray:
    """Collect a subsample of a single-column ``energy`` LazyFrame to numpy."""
    vals = lf.collect(engine="streaming")["energy"].to_numpy()
    if len(vals) > sample_rows:
        rng = np.random.default_rng(seed)
        vals = vals[rng.choice(len(vals), size=sample_rows, replace=False)]
    return vals