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.prediction import PredictionSpec, open_prediction, paired_vars_for_coord, prediction_secondaries
from giant.analysis.reduce import (
    attach_entry_axis,
    depth_expr,
    entry_axis,
    transverse_expr,
)
from giant.analysis.sources import RolloutSpec, 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)
    # -- giant predict (paired truth/pred comparison) — empty when no
    # --prediction was given to `prep`, so old shared.json files still load.
    pred_var_ranges: dict[str, tuple[float, float]] = field(default_factory=dict)
    pred_residual_ranges: dict[str, tuple[float, float]] = field(default_factory=dict)
    pred_n_sec_cap: int = 10
    pred_top_sec_pdgs: list[int] = field(default_factory=list)

    # -- (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"])
        if "pred_var_ranges" in d:
            d["pred_var_ranges"] = {k: tuple(v) for k, v in d["pred_var_ranges"].items()}
        if "pred_residual_ranges" in d:
            d["pred_residual_ranges"] = {k: tuple(v) for k, v in d["pred_residual_ranges"].items()}
        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)

    def pred_marginal_edges(self, var: str) -> np.ndarray:
        lo, hi = self.pred_var_ranges[var]
        return np.linspace(lo, hi, self.n_marginal_bins + 1)

    def pred_residual_edges(self, var: str) -> np.ndarray:
        lo, hi = self.pred_residual_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(vals: list[np.ndarray], lo_q: float, hi_q: float) -> tuple[float, float]:
    """Robust (lo_q, hi_q) range over the union of several value samples."""
    both = np.concatenate(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(
    rollouts: list[RolloutSpec],
    reference: str | Path | pl.LazyFrame,
    *,
    predictions: list[PredictionSpec] | None = None,
    n_energy_bins: int = 4,
    n_marginal_bins: int = 50,
    n_sec_bins: int = 40,
    top_k_pdg: int = 6,
    pred_n_sec_cap: int = 10,
    top_k_sec_pdg: int = 8,
    sample_rows: int = 1_000_000,
    seed: int = 0,
) -> Context:
    """Resolve the shared context from the reference + every rollout (the ``prep`` step).

    Every range/quantile below is the union of the reference and *all*
    rollouts, so a single set of fixed bin edges/group sets is valid for
    every series a compute job streams over.
    """
    t_all = open_side(reference, Side.reference)
    t_lf = physical_steps(t_all, Side.reference)
    r_lfs = {rs.name: physical_steps(open_side(rs.source, Side.rollout), Side.rollout) for rs in rollouts}

    # Ranged marginal variables: robust ranges over a shared row subsample.
    exprs = [e.alias(n) for n, (_, e) in RANGED_VARS.items()]
    t_s = _row_subsample(t_lf, sample_rows, seed).select(exprs).collect(engine="streaming")
    r_s = {
        name: _row_subsample(lf, sample_rows, seed).select(exprs).collect(engine="streaming")
        for name, lf in r_lfs.items()
    }
    var_ranges = {
        name: _combined_quantiles([t_s[name].to_numpy(), *(df[name].to_numpy() for df in r_s.values())], _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()

    t_inc = _incident(t_lf)
    r_inc = {name: _incident(lf) for name, lf in r_lfs.items()}
    energy_edges = energy_bin_edges(np.concatenate([t_inc, *r_inc.values()]), 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(t_lf, "pdg"), *(_counts(lf, "pdg") for lf in r_lfs.values())])
        .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]]
    material_set: set[str] = set(_counts(t_lf, "material")["material"].to_list())
    for lf in r_lfs.values():
        material_set |= set(_counts(lf, "material")["material"].to_list())
    materials = sorted(material_set)

    # 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()

    t_d, t_t = _proxy(t_lf)
    r_proxy = {name: _proxy(lf) for name, lf in r_lfs.items()}
    d_lo, d_hi = _combined_quantiles([t_d, *(p[0] for p in r_proxy.values())], _LO_Q, _HI_Q)
    depth_edges = np.linspace(d_lo, d_hi, n_marginal_bins + 1)
    t_hi = max(float(np.quantile(np.concatenate([t_t, *(p[1] for p in r_proxy.values())]), _HI_Q)), 1e-6)
    transverse_edges = np.linspace(0.0, t_hi, n_marginal_bins + 1)

    # Secondary energy range.
    t_se = _row_sample_col(secondaries(t_all, Side.reference).select("energy"), sample_rows, seed)
    r_se = {
        name: _row_sample_col(secondaries(lf, Side.rollout).select("energy"), sample_rows, seed)
        for name, lf in r_lfs.items()
    }
    sec_energy_range = _combined_quantiles([t_se, *r_se.values()], _LO_Q, _HI_Q)

    # giant predict: paired truth/pred ranges + residual ranges + secondary
    # species vocab, all over the union of every prediction's `paired` frame.
    pred_var_ranges: dict[str, tuple[float, float]] = {}
    pred_residual_ranges: dict[str, tuple[float, float]] = {}
    top_sec_pdgs: list[int] = []
    if predictions:
        sides = {ps.name: open_prediction(ps.source) for ps in predictions}
        present_vars = sorted(set().union(*(paired_vars_for_coord(s.coord) for s in sides.values())))
        for var in present_vars:
            true_samples, pred_samples, residual_samples = [], [], []
            for s in sides.values():
                if var not in paired_vars_for_coord(s.coord):
                    continue
                cols = [f"pred_{var}"] + ([f"true_{var}"] if s.has_truth else [])
                sample = _row_subsample(s.paired.select(cols), sample_rows, seed).collect(engine="streaming")
                pred_samples.append(sample[f"pred_{var}"].to_numpy())
                if s.has_truth:
                    true_samples.append(sample[f"true_{var}"].to_numpy())
                    residual_samples.append(sample[f"pred_{var}"].to_numpy() - sample[f"true_{var}"].to_numpy())
            pred_var_ranges[var] = _combined_quantiles([*true_samples, *pred_samples], _LO_Q, _HI_Q)
            if residual_samples:
                pred_residual_ranges[var] = _combined_quantiles(residual_samples, _LO_Q, _HI_Q)

        sec_pdg_counts: dict[int, int] = {}
        for s in sides.values():
            if s.coord != "global" or not s.has_truth:
                continue
            for prefix in ("true", "pred"):
                counts = (
                    prediction_secondaries(s.lf, prefix)
                    .group_by("pdg")
                    .agg(pl.len().alias("n"))
                    .collect(engine="streaming")
                )
                for pdg, n in zip(counts["pdg"].to_list(), counts["n"].to_list()):
                    sec_pdg_counts[pdg] = sec_pdg_counts.get(pdg, 0) + n
        top_sec_pdgs = [pdg for pdg, _ in sorted(sec_pdg_counts.items(), key=lambda kv: -kv[1])[:top_k_sec_pdg]]

    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],
        pred_var_ranges=pred_var_ranges,
        pred_residual_ranges=pred_residual_ranges,
        pred_n_sec_cap=pred_n_sec_cap,
        pred_top_sec_pdgs=top_sec_pdgs,
        sec_energy_range=sec_energy_range,
        n_sec_bins=n_sec_bins,
        n_events={
            "reference": len(t_inc),
            **{name: len(arr) for name, arr in r_inc.items()},
        },
    )


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