giant/tests/test_analysis_prediction.py

"""Tests for giant.analysis.prediction (paired truth/pred frames for `giant predict`
output) and the `prediction` family of catalog specs."""

from __future__ import annotations

import numpy as np
import polars as pl
import pytest

from giant.analysis.catalog import Bundle, get_spec
from giant.analysis.context import Context, build_context
from giant.analysis.prediction import (
    PAIRED_SCALARS,
    PredictionSpec,
    open_prediction,
    paired_frame,
    paired_secondaries,
    prediction_secondaries,
)
from giant.analysis.reduce import hist2d
from giant.analysis.sources import RolloutSpec
from tests.test_analysis_reduce import _reference_frame, _rollout_frame


def _global_prediction_frame() -> pl.LazyFrame:
    """A `--coord global --truth` predict parquet, as a LazyFrame (schema per
    `giant.cli.predict`'s global-coord table, `giant/cli.py:1310-1379`)."""
    return pl.DataFrame(
        {
            "event_id": [1, 1, 2],
            "pdg": [11, 11, 22],
            "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],
            "material": ["G4_PbWO4", "G4_PbWO4", "G4_Pb"],
            "layer_id": [0, 1, 0],
            "n_sec": [1, 0, 2],
            "n_sec_pred": [1, 0, 1],
            # predicted (unprefixed) values
            "step_length": [1.2, 0.9, 1.1],
            "delta_e": [42.0, 29.0, 31.0],
            "edep": [35.0, 29.0, 25.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],
            "post_x": [0.0, 0.0, 0.0],
            "post_y": [0.0, 0.0, 0.0],
            "post_z": [1.2, 1.9, 1.1],
            "sec_pdg_list": [[22], [], [22]],
            "sec_E_list": [[5.0], [], [4.0]],
            "sec_dx_list": [[0.0], [], [0.0]],
            "sec_dy_list": [[0.0], [], [0.0]],
            "sec_dz_list": [[1.0], [], [1.0]],
            # truth
            "true_step_length": [1.0, 1.0, 1.0],
            "true_delta_e": [40.0, 30.0, 30.0],
            "true_edep": [40.0, 30.0, 20.0],
            "true_post_E": [60.0, 30.0, 20.0],
            "true_post_dx": [0.0, 0.0, 0.0],
            "true_post_dy": [0.0, 0.0, 0.0],
            "true_post_dz": [1.0, 1.0, 1.0],
            "true_post_x": [0.0, 0.0, 0.0],
            "true_post_y": [0.0, 0.0, 0.0],
            "true_post_z": [1.0, 2.0, 1.0],
            "true_e_sec": [0.0, 0.0, 10.0],
            "process": ["compt", "phot", "compt"],
            "true_sec_pdg_list": [[22], [], [22, 11]],
            "true_sec_E_list": [[6.0], [], [7.0, 3.0]],
            "true_sec_dx_list": [[0.0], [], [0.0, 1.0]],
            "true_sec_dy_list": [[0.0], [], [0.0, 0.0]],
            "true_sec_dz_list": [[1.0], [], [1.0, 0.0]],
        }
    ).lazy()


def _local_prediction_frame() -> pl.LazyFrame:
    """A `--coord local` predict parquet — always paired, never has secondaries."""
    return pl.DataFrame(
        {
            "event_id": [1, 2],
            "pdg": [11, 22],
            "pre_x": [0.0, 0.0],
            "pre_y": [0.0, 0.0],
            "pre_z": [0.0, 0.0],
            "pre_E": [100.0, 50.0],
            "pre_dx": [0.0, 0.0],
            "pre_dy": [0.0, 0.0],
            "pre_dz": [1.0, 1.0],
            "material": ["G4_PbWO4", "G4_Pb"],
            "layer_id": [0, 0],
            "n_sec": [1, 0],
            # ALR logits: [edep_logit, sec_logit] -> softmax([z1,z2,0]) * pre_E
            "pred_log_step_length": [np.log(1.2 + 1e-6), np.log(0.9 + 1e-6)],
            "pred_edep_logit": [1.0, 0.5],
            "pred_sec_logit": [0.0, -1.0],
            "pred_post_dx": [0.0, 0.0],
            "pred_post_dy": [0.0, 0.0],
            "pred_post_dz": [1.0, 1.0],
            "pred_travel_dx": [0.0, 0.0],
            "pred_travel_dy": [0.0, 0.0],
            "pred_travel_dz": [1.0, 1.0],
            "true_log_step_length": [np.log(1.0 + 1e-6), np.log(1.0 + 1e-6)],
            "true_edep_logit": [0.8, 0.6],
            "true_sec_logit": [0.2, -2.0],
            "true_post_dx": [0.0, 0.0],
            "true_post_dy": [0.0, 0.0],
            "true_post_dz": [1.0, 1.0],
            "true_travel_dx": [0.0, 0.0],
            "true_travel_dy": [0.0, 0.0],
            "true_travel_dz": [1.0, 1.0],
        }
    ).lazy()


def test_open_prediction_detects_coord_and_truth():
    g = open_prediction(_global_prediction_frame())
    assert g.coord == "global" and g.has_truth

    loc = open_prediction(_local_prediction_frame())
    assert loc.coord == "local" and loc.has_truth


def test_paired_frame_global_matches_source_columns():
    lf = _global_prediction_frame()
    p = paired_frame(lf, "global", has_truth=True).collect()
    assert p["pred_step_length"].to_list() == [1.2, 0.9, 1.1]
    assert p["true_step_length"].to_list() == [1.0, 1.0, 1.0]
    assert p["pred_edep"].to_list() == [35.0, 29.0, 25.0]
    assert p["true_edep"].to_list() == [40.0, 30.0, 20.0]
    # post_E isn't written directly for the prediction (energy conservation:
    # pre_E - delta_e); truth carries it verbatim.
    assert p["pred_post_E"].to_list() == pytest.approx([100.0 - 42.0, 60.0 - 29.0, 50.0 - 31.0])
    assert p["true_post_E"].to_list() == [60.0, 30.0, 20.0]
    # cos_scatter: pre_dir . post_dir, both (0,0,1) here -> 1.0
    assert p["pred_cos_scatter"].to_list() == pytest.approx([1.0, 1.0, 1.0])
    assert p["true_cos_scatter"].to_list() == pytest.approx([1.0, 1.0, 1.0])


def test_paired_frame_local_decodes_energy_simplex():
    lf = _local_prediction_frame()
    p = paired_frame(lf, "local", has_truth=True).collect()
    # softmax([1.0, 0.0, 0.0]) * 100 for row 0's pred edep
    z = np.exp([1.0, 0.0, 0.0])
    expected_edep_0 = (z[0] / z.sum()) * 100.0
    assert p["pred_edep"][0] == pytest.approx(expected_edep_0)
    assert p["pred_step_length"][0] == pytest.approx(1.2, abs=1e-4)
    # local coord never has a meaningful cos_travel (no reconstructed post_pos)
    assert "cos_travel" not in [c.rsplit("_", 1)[-1] for c in ["pred_cos_travel"] if c in p.columns] or True
    assert "pred_cos_travel" not in p.columns


def test_prediction_secondaries_and_pairing():
    lf = _global_prediction_frame()
    true_sec = prediction_secondaries(lf, "true").collect()
    pred_sec = prediction_secondaries(lf, "pred").collect()
    assert true_sec["pdg"].to_list() == [22, 22, 11]
    assert pred_sec["pdg"].to_list() == [22, 22]

    pairs = paired_secondaries(lf).collect()
    # event 1: 1 true, 1 pred -> paired (22, 22); event 2: 2 true, 1 pred -> paired rank0 only (22, 22)
    assert pairs["true_pdg"].to_list() == [22, 22]
    assert pairs["pred_pdg"].to_list() == [22, 22]


def test_hist2d_basic():
    lf = pl.DataFrame({"x": [0.1, 0.5, 0.9, 0.5], "y": [0.1, 0.9, 0.9, 0.1]}).lazy()
    edges = np.linspace(0.0, 1.0, 3)  # 2 bins: [0,0.5), [0.5,1]
    mat = hist2d(lf, pl.col("x"), pl.col("y"), edges, edges)
    assert mat.sum() == 4
    assert mat.shape == (2, 2)


def _ctx_with_predictions(n_marginal_bins: int = 10) -> Context:
    return build_context(
        [RolloutSpec("rollout", _rollout_frame())],
        _reference_frame(),
        predictions=[PredictionSpec("pred", _global_prediction_frame())],
        n_energy_bins=2,
        n_marginal_bins=n_marginal_bins,
        top_k_pdg=3,
        sample_rows=1000,
    )


def test_build_context_resolves_prediction_ranges():
    ctx = _ctx_with_predictions()
    assert "edep" in ctx.pred_var_ranges
    assert "edep" in ctx.pred_residual_ranges
    assert ctx.pred_top_sec_pdgs  # secondaries present in the fixture


def test_prediction_specs_compute_valid_reduced():
    ctx = _ctx_with_predictions()
    bundle = Bundle.open(
        [RolloutSpec("rollout", _rollout_frame())],
        _reference_frame(),
        ctx,
        predictions=[PredictionSpec("pred", _global_prediction_frame())],
    )
    for spec_id in (
        "pred_marginal_edep",
        "pred_scatter_edep",
        "pred_residual_edep",
        "pred_relative_residual_edep",
        "pred_residual_profile_edep",
        "pred_ks_summary",
        "pred_bias_summary",
        "pred_rmse_summary",
        "pred_n_sec_confusion",
        "pred_sec_species_confusion",
        "pred_dir_alignment_post",
        "pred_dir_alignment_travel",
        "pred_constraint_violations",
        "pred_correlation_delta",
    ):
        spec = get_spec(spec_id)
        r = spec.finalize([spec.compute_partial(bundle)], ctx)
        assert r.id == spec_id
        assert r.kind != "unavailable", f"{spec_id} unexpectedly unavailable"
        assert "pred" in r.payload["series"]


def test_prediction_specs_unavailable_without_predictions():
    ctx = _ctx_with_predictions()
    bundle = Bundle.open([RolloutSpec("rollout", _rollout_frame())], _reference_frame(), ctx)
    for spec_id in ("pred_marginal_edep", "pred_scatter_edep", "pred_n_sec_confusion", "pred_ks_summary"):
        spec = get_spec(spec_id)
        r = spec.finalize([spec.compute_partial(bundle)], ctx)
        assert r.kind == "unavailable"
        assert r.payload["note"]


@pytest.mark.parametrize(
    "spec_id",
    ["pred_marginal_edep", "pred_scatter_edep", "pred_n_sec_confusion", "pred_ks_summary", "pred_correlation_delta"],
)
def test_prediction_chunked_matches_unchunked(spec_id: str):
    ctx = _ctx_with_predictions()
    specs = [RolloutSpec("rollout", _rollout_frame())]
    preds = [PredictionSpec("pred", _global_prediction_frame())]
    spec = get_spec(spec_id)

    unchunked_bundle = Bundle.open(specs, _reference_frame(), ctx, predictions=preds)
    unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], ctx)

    n_chunks = 2
    parts = [
        spec.compute_partial(Bundle.open(specs, _reference_frame(), ctx, chunk=(k, n_chunks), predictions=preds))
        for k in range(n_chunks)
    ]
    chunked = spec.finalize(parts, ctx)

    assert chunked.kind == unchunked.kind
    _assert_close(unchunked.payload, chunked.payload)


def _assert_close(a, b) -> None:
    """Recursively compare two JSON-shaped payloads (float-tolerant)."""
    if isinstance(a, dict):
        assert set(a) == set(b)
        for k in a:
            _assert_close(a[k], b[k])
    elif isinstance(a, list):
        assert len(a) == len(b)
        for x, y in zip(a, b):
            _assert_close(x, y)
    elif isinstance(a, float):
        assert np.isclose(a, b, atol=1e-9) or (np.isnan(a) and np.isnan(b))
    else:
        assert a == b


def test_paired_scalars_are_subset_of_all_vars():
    assert set(PAIRED_SCALARS) <= {"step_length", "edep", "delta_e", "post_E"}