giant/tests/test_analysis_reduce.py

"""Tests for the streaming compute primitives (giant.analysis.reduce/sources/grouping)."""

from __future__ import annotations

import numpy as np
import polars as pl

from giant.analysis import grouping as G
from giant.analysis import reduce as R
from giant.analysis.sources import (
    SYNTHETIC_TERMINATION_REASONS,
    Side,
    open_side,
    physical_steps,
    secondaries,
)
from giant.data.loader import EVENT_ID_FILE_STRIDE


def _rollout_frame() -> pl.LazyFrame:
    # event 1: primary (2 steps) + 1 secondary track + 1 escaped bookkeeping row
    # event 2: primary (1 step)
    return pl.DataFrame(
        {
            "event_id": [1, 1, 1, 1, 2],
            "track_id": [0, 0, 1, 0, 0],
            "parent_id": [-1, -1, 0, -1, -1],
            "generation": [0, 0, 1, 0, 0],
            "step_no": [0, 1, 0, 99, 0],
            "pdg": [11, 11, 22, 11, 11],
            "pre_x": [0.0, 0.0, 0.0, 0.0, 0.0],
            "pre_y": [0.0, 0.0, 0.0, 0.0, 0.0],
            "pre_z": [0.0, 1.0, 1.0, 2.0, 0.0],
            "pre_E": [100.0, 60.0, 20.0, 30.0, 50.0],
            "pre_dx": [0.0, 0.0, 1.0, 0.0, 0.0],
            "pre_dy": [0.0, 0.0, 0.0, 0.0, 0.0],
            "pre_dz": [1.0, 1.0, 0.0, 1.0, 1.0],
            "post_x": [0.0, 0.0, 1.0, 0.0, 0.0],
            "post_y": [0.0, 0.0, 0.0, 0.0, 0.0],
            "post_z": [1.0, 2.0, 1.0, 2.0, 1.0],
            "post_E": [60.0, 30.0, 0.0, 0.0, 20.0],
            "post_dx": [0.0, 0.0, 1.0, 0.0, 0.0],
            "post_dy": [0.0, 0.0, 0.0, 0.0, 0.0],
            "post_dz": [1.0, 1.0, 0.0, 1.0, 1.0],
            "edep": [40.0, 30.0, 20.0, 0.0, 30.0],
            "step_length": [1.0, 1.0, 1.0, 0.0, 1.0],
            "material": ["G4_PbWO4"] * 5,
            "layer_id": [0, 1, 1, -1, 0],
            "n_sec_pred": [1, 0, 0, 0, 0],
            "termination_reason": [
                "",
                "natural_end",
                "natural_end",
                "escaped",
                "natural_end",
            ],
        }
    ).lazy()


def _reference_frame() -> pl.LazyFrame:
    return pl.DataFrame(
        {
            "event_id": [1, 1, 2],
            "track_id": [0, 0, 0],
            "step_no": [0, 1, 0],
            "pdg": [11, 11, 11],
            "pre_x": [0.0, 0.0, 0.0],
            "pre_y": [0.0, 0.0, 0.0],
            "pre_z": [0.0, 1.0, 0.0],
            "pre_E": [100.0, 60.0, 50.0],
            "pre_dx": [0.0, 0.0, 0.0],
            "pre_dy": [0.0, 0.0, 0.0],
            "pre_dz": [1.0, 1.0, 1.0],
            "post_x": [0.0, 0.0, 0.0],
            "post_y": [0.0, 0.0, 0.0],
            "post_z": [1.0, 2.0, 1.0],
            "post_E": [60.0, 30.0, 20.0],
            "post_dx": [0.0, 0.0, 0.0],
            "post_dy": [0.0, 0.0, 0.0],
            "post_dz": [1.0, 1.0, 1.0],
            "edep": [40.0, 30.0, 30.0],
            "step_length": [1.0, 1.0, 1.0],
            "material": ["G4_PbWO4", "G4_PbWO4", "G4_Pb"],
            "layer_id": [0, 1, 0],
            "sec_E_list": [[20.0], [], [10.0]],
            "sec_pdg_list": [[22], [], [22]],
            "sec_dx_list": [[1.0], [], [0.0]],
            "sec_dy_list": [[0.0], [], [0.0]],
            "sec_dz_list": [[0.0], [], [1.0]],
        }
    ).lazy()


def test_hist1d_overall_and_grouped():
    lf = _rollout_frame()
    edges = np.linspace(0.0, 50.0, 6)  # width 10
    h = R.hist1d(lf, pl.col("edep"), edges)
    # edep values: 40,30,20,0,30 -> bins [0),[10),[20),[30),[40)
    assert h[0].tolist() == [1, 0, 1, 2, 1]
    # grouped by pdg: pdg 22 has a single edep=20
    hg = R.hist1d(lf, pl.col("edep"), edges, group=pl.col("pdg"))
    assert hg[22].tolist() == [0, 0, 1, 0, 0]
    assert hg[11].sum() == 4


def test_hist1d_clamps_extreme_values_and_drops_nan():
    # A rollout can emit a wildly out-of-range step_length (or an inf/NaN); the
    # fixed-edge binning must clamp rather than overflow the i32 bin cast.
    lf = pl.DataFrame({"x": [5.0, 1.0725e10, float("inf"), -float("inf"), float("nan"), None]}).lazy()
    edges = np.linspace(0.0, 50.0, 6)  # width 10
    h = R.hist1d(lf, pl.col("x"), edges)
    # 5 -> bin 0; 1e10 and +inf -> top bin; -inf -> bin 0; NaN/null dropped
    assert h[0].tolist() == [2, 0, 0, 0, 2]


def test_profile_partial_clamps_extreme_values_and_drops_nan():
    lf = pl.DataFrame(
        {
            "event_id": [1, 1, 1, 1],
            "z": [5.0, 1.0725e10, float("nan"), 45.0],
            "w": [1.0, 2.0, 4.0, 8.0],
        }
    ).lazy()
    edges = np.linspace(0.0, 50.0, 6)
    ev, mat = R.profile_partial(lf, pl.col("z"), edges, pl.col("w"))
    assert ev.tolist() == [1]
    # 1e10 clamps into the top bin alongside 45; the NaN row's weight is dropped
    assert mat[0].tolist() == [1.0, 0.0, 0.0, 0.0, 10.0]


def test_physical_steps_drops_synthetic_rollout_rows_only():
    lf = _rollout_frame()
    phys = physical_steps(lf, Side.rollout).collect()
    assert phys.height == 4  # dropped the escaped bookkeeping row
    assert "escaped" not in phys["termination_reason"].to_list()
    assert SYNTHETIC_TERMINATION_REASONS  # non-empty guard
    # reference passes through unchanged
    ref = _reference_frame()
    assert physical_steps(ref, Side.reference).collect().height == ref.collect().height


def test_event_scalars_totals_include_all_rows():
    lf = _rollout_frame()
    es = R.event_scalars(lf).sort("event_id")
    row1 = es.filter(pl.col("event_id") == 1).to_dicts()[0]
    assert row1["total_edep"] == 90.0  # 40+30+20+0
    assert row1["incident_E"] == 100.0
    assert row1["n_steps"] == 4


def test_secondaries_rollout_vs_reference_align():
    r = secondaries(_rollout_frame(), Side.rollout).collect().sort("event_id")
    assert r["energy"].to_list() == [20.0]  # only the generation>0, step_no==0 row
    assert r["pdg"].to_list() == [22]
    t = secondaries(_reference_frame(), Side.reference).collect().sort("event_id")
    # two secondaries (event 1 and event 2); empty list dropped
    assert sorted(t["energy"].to_list()) == [10.0, 20.0]
    assert t["pdg"].to_list() == [22, 22]


def test_sec_count_by_event_zero_fills_events_with_no_secondaries():
    r_phys = physical_steps(_rollout_frame(), Side.rollout)
    r_sec = secondaries(_rollout_frame(), Side.rollout)
    ev, n = R.sec_count_by_event(r_phys, r_sec)
    # event 1 has one secondary track; event 2 has none and must still appear (as 0),
    # not silently drop out of a plain group_by on the secondaries frame alone.
    assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 0}

    t_all = _reference_frame()
    t_sec = secondaries(t_all, Side.reference)
    ev, n = R.sec_count_by_event(t_all, t_sec)
    assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 1}


def test_leakage_fraction():
    frac = R.leakage_fraction(_rollout_frame())
    # event 1: escaped pre_E=30, deposited=90 -> 30/120 = 0.25; event 2: 0
    assert sorted(round(f, 6) for f in frac) == [0.0, 0.25]


def test_weighted_profile_matches_manual_bincount():
    lf = _rollout_frame()
    ea = R.entry_axis(lf)
    lf2 = R.attach_entry_axis(lf, ea)
    edges = np.linspace(0.0, 3.0, 4)  # depth bins along +z
    mean, std = R.weighted_profile(lf2, R.depth_expr(), edges, pl.col("edep"))
    assert mean.shape == (3,)
    # totals conserved: sum over bins == mean total edep per event
    assert np.isclose(mean.sum() * 1, (90.0 + 30.0) / 2)  # 2 events


def test_energy_bins_edges_and_event_map():
    incident = np.array([100.0, 100.0, 1000.0, 1000.0])
    edges = G.energy_bin_edges(incident, n_bins=2)
    assert len(edges) == 3 and edges[0] <= 100.0 < edges[-1]
    ids, bins = G.event_energy_bins(_rollout_frame(), edges)
    assert set(bins.tolist()) <= {0, 1}
    assert len(ids) == 2


def test_digitize_expr_matches_numpy():
    edges = np.array([0.0, 10.0, 100.0, 1000.0])
    df = pl.DataFrame({"v": [5.0, 50.0, 500.0, 2000.0]})
    got = df.select(G.digitize_expr(pl.col("v"), edges).alias("b"))["b"].to_list()
    assert got == np.digitize([5.0, 50.0, 500.0, 2000.0], edges[1:-1]).tolist()


def test_pdg_and_material_labels():
    assert G.pdg_label(22) == "gamma"
    assert G.pdg_label(999999) == "999999"
    assert G.material_label("G4_PbWO4") == "PbWO4"


def _write_shard(path, event_ids, edeps):
    pl.DataFrame({"event_id": event_ids, "pdg": [11] * len(event_ids), "edep": edeps}).write_parquet(path)


def test_open_side_reference_offsets_event_ids_across_shards(tmp_path):
    # Each shard is a separate Geant4 job whose own event_id numbering restarts
    # from 0 — a naive multi-shard scan collides on event_id across shards.
    _write_shard(tmp_path / "a.parquet", [0, 1], [1.0, 2.0])
    _write_shard(tmp_path / "b.parquet", [0, 1], [3.0, 4.0])
    df = open_side(tmp_path, Side.reference).sort("event_id").collect()
    assert df["event_id"].to_list() == [0, 1, EVENT_ID_FILE_STRIDE, EVENT_ID_FILE_STRIDE + 1]
    assert df["edep"].to_list() == [1.0, 2.0, 3.0, 4.0]
    assert "__source_path" not in df.columns


def test_open_side_reference_single_file_unchanged(tmp_path):
    _write_shard(tmp_path / "only.parquet", [0, 1], [1.0, 2.0])
    df = open_side(tmp_path / "only.parquet", Side.reference).sort("event_id").collect()
    assert df["event_id"].to_list() == [0, 1]
    assert "__source_path" not in df.columns


def test_open_side_reference_manifest(tmp_path):
    _write_shard(tmp_path / "a.parquet", [0, 1], [1.0, 2.0])
    _write_shard(tmp_path / "b.parquet", [0, 1], [3.0, 4.0])
    manifest = tmp_path / "shards.manifest"
    manifest.write_text("a.parquet\nb.parquet\n")
    df = open_side(manifest, Side.reference).sort("event_id").collect()
    assert df["event_id"].to_list() == [0, 1, EVENT_ID_FILE_STRIDE, EVENT_ID_FILE_STRIDE + 1]
    assert df["edep"].to_list() == [1.0, 2.0, 3.0, 4.0]