giant/tests/test_phase2.py

"""Tests for Phase 2: secondary particle prediction."""

import numpy as np
import pytest
import torch

from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, X_DIM
from giant.model.network import DenoisingMLP, SecondaryDecoder
from giant.model.schedule import flow_matching_loss_secondary
from giant.sample import sample_secondaries


# ── helpers ──────────────────────────────────────────────────────────────────


def _stage1(pdg=3, mat=2, conditioning="embedding"):
    return DenoisingMLP(
        pdg_vocab=pdg,
        mat_vocab=mat,
        hidden_dim=32,
        n_blocks=2,
        conditioning=conditioning,
    )


def _sec_decoder(pdg=3, mat=2, conditioning="embedding"):
    return SecondaryDecoder(
        pdg_vocab=pdg,
        mat_vocab=mat,
        hidden_dim=32,
        n_blocks=2,
        conditioning=conditioning,
    )


def _cond(B=8, pdg=3, mat=2):
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.stack(
        [torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
    )
    return cond_cont, cond_cat


# ── DenoisingMLP Phase-2 additions ───────────────────────────────────────────


def test_predict_n_sec_shape():
    B = 8
    model = _stage1()
    cond_cont, cond_cat = _cond(B)
    logits = model.predict_n_sec(cond_cont, cond_cat)
    assert logits.shape == (B, K_MAX + 1)


def test_predict_n_sec_no_nan():
    B = 8
    model = _stage1()
    cond_cont, cond_cat = _cond(B)
    logits = model.predict_n_sec(cond_cont, cond_cat)
    assert torch.isfinite(logits).all()


@pytest.mark.parametrize("conditioning", ["embedding", "physical"])
def test_no_pdg_embedding_weight_method(conditioning):
    """The Stage-2 species output no longer needs a shared embedding table."""
    model = _stage1(pdg=5, mat=2, conditioning=conditioning)
    assert not hasattr(model, "pdg_embedding_weight")


def test_condition_encoder_physical_mode_has_no_embedding_tables():
    model = _stage1(pdg=5, mat=2, conditioning="physical")
    assert not hasattr(model.cond_enc, "pdg_emb")
    assert not hasattr(model.cond_enc, "mat_emb")
    assert hasattr(model.cond_enc, "particle_mlp")
    assert hasattr(model.cond_enc, "material_mlp")


def test_condition_encoder_embedding_mode_has_embedding_tables():
    model = _stage1(pdg=5, mat=2, conditioning="embedding")
    assert hasattr(model.cond_enc, "pdg_emb")
    assert hasattr(model.cond_enc, "mat_emb")
    assert not hasattr(model.cond_enc, "particle_mlp")


# ── SecondaryDecoder ──────────────────────────────────────────────────────────


@pytest.mark.parametrize("conditioning", ["embedding", "physical"])
def test_sec_decoder_output_shape(conditioning):
    B = 8
    decoder = _sec_decoder(conditioning=conditioning)
    x_t = torch.randn(B, SEC_DIM)
    t = torch.rand(B)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    out = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
    assert out.shape == (B, SEC_DIM)


def test_sec_decoder_no_nan():
    B = 4
    decoder = _sec_decoder()
    x_t = torch.randn(B, SEC_DIM)
    t = torch.rand(B)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    out = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
    assert torch.isfinite(out).all()


def test_sec_decoder_gradients():
    B = 4
    decoder = _sec_decoder()
    x_t = torch.randn(B, SEC_DIM)
    t = torch.rand(B)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward()
    for name, p in decoder.named_parameters():
        assert p.grad is not None, f"no grad for {name}"


# ── masked flow matching loss ─────────────────────────────────────────────────


def test_flow_matching_loss_secondary_scalar():
    B, pdg, mat = 8, 3, 2
    decoder = _sec_decoder(pdg, mat)
    x1 = torch.randn(B, SEC_DIM)
    cond_cont, cond_cat = _cond(B, pdg, mat)
    stage1_out = torch.randn(B, X_DIM)
    sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
    loss = flow_matching_loss_secondary(
        decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
    )
    assert loss.shape == ()
    assert loss.item() >= 0.0


def test_flow_matching_loss_secondary_mask_zeros_padding():
    """Loss with all-zero mask (no valid secondaries) should be 0."""
    B, pdg, mat = 4, 3, 2
    decoder = _sec_decoder(pdg, mat)
    x1 = torch.randn(B, SEC_DIM)
    cond_cont, cond_cat = _cond(B, pdg, mat)
    stage1_out = torch.randn(B, X_DIM)
    sec_mask = torch.zeros(B, K_MAX, dtype=torch.bool)
    loss = flow_matching_loss_secondary(
        decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
    )
    assert loss.item() == pytest.approx(0.0, abs=1e-6)


def test_flow_matching_loss_secondary_has_grad():
    B, pdg, mat = 4, 3, 2
    decoder = _sec_decoder(pdg, mat)
    x1 = torch.randn(B, SEC_DIM)
    cond_cont, cond_cat = _cond(B, pdg, mat)
    stage1_out = torch.randn(B, X_DIM)
    sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
    flow_matching_loss_secondary(
        decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
    ).backward()
    assert any(p.grad is not None for p in decoder.parameters())


# ── sampling ──────────────────────────────────────────────────────────────────


def test_sample_secondaries_shapes():
    B, pdg, mat = 6, 3, 2
    decoder = _sec_decoder(pdg, mat)
    cond_cont, cond_cat = _cond(B, pdg, mat)
    stage1_out = torch.randn(B, X_DIM)
    n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
    sec_cont, sec_phys, sec_valid = sample_secondaries(
        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=3
    )
    assert sec_cont.shape == (B, K_MAX, 4)
    assert sec_phys.shape == (B, K_MAX, PARTICLE_PHYS_DIM)
    assert sec_valid.shape == (B, K_MAX)
    assert sec_valid.dtype == torch.bool


def test_sample_secondaries_valid_mask_matches_n_sec():
    B, pdg, mat = 4, 3, 2
    decoder = _sec_decoder(pdg, mat)
    cond_cont, cond_cat = _cond(B, pdg, mat)
    stage1_out = torch.randn(B, X_DIM)
    n_sec_pred = torch.tensor([0, 1, 3, K_MAX])
    _, _, sec_valid = sample_secondaries(
        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
    )
    for i, n in enumerate(n_sec_pred.tolist()):
        assert sec_valid[i, :n].all()
        assert not sec_valid[i, n:].any()


# ── encode_secondaries round-trip ─────────────────────────────────────────────


def test_encode_secondaries_energy_conservation():
    """Decoded stick-breaking fractions must sum to ≈ e_sec."""
    from giant.data.transforms import encode_secondaries

    rng = np.random.default_rng(42)
    N = 50
    n_sec = rng.integers(1, 5, size=N)
    e_sec = rng.uniform(0.1, 10.0, size=N).astype(np.float32)

    sec_E_list = np.zeros((N, K_MAX), dtype=np.float32)
    sec_dir_list = np.zeros((N, K_MAX, 3), dtype=np.float32)
    sec_dir_list[:, :, 2] = 1.0
    sec_pdg_list = np.zeros((N, K_MAX), dtype=np.int64)
    sec_valid = np.zeros((N, K_MAX), dtype=bool)
    for i in range(N):
        k = n_sec[i]
        energies = rng.dirichlet(np.ones(k)) * e_sec[i]
        energies = np.sort(energies)[::-1]
        sec_E_list[i, :k] = energies.astype(np.float32)
        sec_valid[i, :k] = True
        sec_pdg_list[i, :k] = 22  # photon — resolvable by giant.particles

    pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
    pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)

    sec_cont = encode_secondaries(
        sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
    )
    assert sec_cont.shape == (N, K_MAX, 6)
    assert np.isfinite(sec_cont).all()


def test_encode_secondaries_direction_encoding():
    """Local-frame secondary directions should be unit vectors for valid slots."""
    from giant.data.transforms import encode_secondaries

    rng = np.random.default_rng(7)
    N = 20
    e_sec = np.ones(N, dtype=np.float32) * 5.0
    sec_E_list = np.zeros((N, K_MAX), dtype=np.float32)
    sec_E_list[:, 0] = 3.0
    sec_E_list[:, 1] = 2.0
    sec_dir_list = rng.standard_normal((N, K_MAX, 3)).astype(np.float32)
    norms = np.linalg.norm(sec_dir_list, axis=-1, keepdims=True)
    sec_dir_list /= np.where(norms > 0, norms, 1.0)
    sec_valid = np.zeros((N, K_MAX), dtype=bool)
    sec_valid[:, :2] = True

    pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
    sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir)

    # dir columns are sec_cont[:, :, 1:4]
    local_dirs = sec_cont[:, :2, 1:4]  # (N, 2, 3) — valid slots only
    norms_out = np.linalg.norm(local_dirs, axis=-1)
    np.testing.assert_allclose(norms_out, 1.0, atol=1e-5)


def test_encode_secondaries_physical_columns_without_pdg_list():
    """Omitting sec_pdg_list zero-fills the physical columns (no crash)."""
    from giant.data.transforms import encode_secondaries

    N = 3
    e_sec = np.ones(N, dtype=np.float32)
    sec_E_list = np.zeros((N, K_MAX), dtype=np.float32)
    sec_dir_list = np.zeros((N, K_MAX, 3), dtype=np.float32)
    sec_dir_list[:, :, 2] = 1.0
    sec_valid = np.zeros((N, K_MAX), dtype=bool)
    pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)

    sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir)
    np.testing.assert_allclose(sec_cont[:, :, 4:6], 0.0)


def test_encode_secondaries_physical_columns_match_ground_truth_pdg():
    """log_mass/charge for a valid slot match giant.particles for that PDG."""
    from giant.data.transforms import encode_secondaries, log_transform
    from giant.particles import particle_mass_charge

    N = 1
    e_sec = np.array([5.0], dtype=np.float32)
    sec_E_list = np.zeros((N, K_MAX), dtype=np.float32)
    sec_E_list[0, 0] = 5.0
    sec_dir_list = np.zeros((N, K_MAX, 3), dtype=np.float32)
    sec_dir_list[0, 0] = [0, 0, 1]
    sec_pdg_list = np.zeros((N, K_MAX), dtype=np.int64)
    sec_pdg_list[0, 0] = 11  # electron
    sec_valid = np.zeros((N, K_MAX), dtype=bool)
    sec_valid[0, 0] = True
    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)

    sec_cont = encode_secondaries(
        sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
    )
    mass, charge = particle_mass_charge(11)
    assert sec_cont[0, 0, 4] == pytest.approx(log_transform(np.array([mass]))[0])
    assert sec_cont[0, 0, 5] == pytest.approx(charge)


# ── decode_secondaries: exact energy conservation ────────────────────────────


def _random_sec_cont(rng, N, stick_logit_scale=1.0):
    sec_cont = rng.standard_normal((N, K_MAX, 6)).astype(np.float32)
    sec_cont[:, :, 0] *= stick_logit_scale
    dirs = sec_cont[:, :, 1:4]
    dirs /= np.linalg.norm(dirs, axis=-1, keepdims=True)
    return sec_cont


def test_decode_secondaries_valid_slots_sum_to_e_sec():
    """The valid slots' energies must sum to exactly e_sec, not just <= e_sec.

    Rows with n_sec=0 are excluded: there's no slot to put the budget in, so
    valid_sum is correctly 0 regardless of e_sec there (see
    test_decode_secondaries_zero_n_sec_has_zero_energy) — the shortfall in
    that case is handled downstream (e.g. rollout.py dumps it into edep).
    """
    from giant.data.transforms import decode_secondaries

    rng = np.random.default_rng(0)
    N = 200
    sec_cont = _random_sec_cont(rng, N)
    n_sec = rng.integers(0, K_MAX + 1, size=N)
    e_sec = rng.uniform(0.0, 50.0, size=N).astype(np.float32)
    pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)

    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
        sec_cont, n_sec, e_sec, pre_dir
    )

    valid_sum = (sec_E * sec_valid).sum(axis=1)
    has_secondaries = n_sec > 0
    np.testing.assert_allclose(
        valid_sum[has_secondaries],
        e_sec[has_secondaries],
        atol=1e-3,
        rtol=1e-5,
    )


def test_decode_secondaries_zero_n_sec_has_zero_energy():
    """n_sec=0 rows get no secondaries and no forced energy assignment."""
    from giant.data.transforms import decode_secondaries

    rng = np.random.default_rng(1)
    N = 10
    sec_cont = _random_sec_cont(rng, N)
    n_sec = np.zeros(N, dtype=np.int64)
    e_sec = rng.uniform(1.0, 10.0, size=N).astype(np.float32)
    pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)

    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
        sec_cont, n_sec, e_sec, pre_dir
    )

    assert not sec_valid.any()
    np.testing.assert_allclose(sec_E, 0.0)


def test_decode_secondaries_degenerate_row_falls_back_to_even_split():
    """All-zero stick fractions for the valid slots fall back to an even split."""
    from giant.data.transforms import decode_secondaries

    rng = np.random.default_rng(2)
    N = 4
    sec_cont = _random_sec_cont(rng, N)
    # Drive every valid slot's stick-breaking fraction to ~0 (huge negative logit).
    n_sec = np.array([0, 1, 3, K_MAX])
    for i, k in enumerate(n_sec):
        sec_cont[i, :k, 0] = -80.0
    e_sec = np.array([0.0, 4.0, 9.0, 30.0], dtype=np.float32)
    pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)

    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
        sec_cont, n_sec, e_sec, pre_dir
    )

    for i, k in enumerate(n_sec):
        if k == 0:
            continue
        np.testing.assert_allclose(sec_E[i, :k], e_sec[i] / k, atol=1e-4)
        np.testing.assert_allclose(sec_E[i, :k].sum(), e_sec[i], atol=1e-3)


def test_decode_secondaries_rescale_preserves_relative_shares():
    """Rescaling should keep each valid slot's *share* of the budget unchanged.

    A shortfall shouldn't get dumped into whichever slot is last by energy
    rank — it should be spread proportionally, i.e. sec_E[i] / sec_E[j] for
    two valid slots must match before and after the e_sec rescale.
    """
    from giant.data.transforms import decode_secondaries

    rng = np.random.default_rng(3)
    N = 1
    sec_cont = _random_sec_cont(rng, N)
    n_sec = np.array([4])
    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)

    sec_E_small, _, _, _, sec_valid = decode_secondaries(
        sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir
    )
    sec_E_large, _, _, _, _ = decode_secondaries(
        sec_cont, n_sec, np.array([50.0], dtype=np.float32), pre_dir
    )

    ratio_small = sec_E_small[0, :4] / sec_E_small[0, 0]
    ratio_large = sec_E_large[0, :4] / sec_E_large[0, 0]
    np.testing.assert_allclose(ratio_small, ratio_large, rtol=1e-4)


def test_decode_secondaries_mass_charge_round_trip_with_normalizer():
    from giant.data.transforms import Normalizer, decode_secondaries, encode_secondaries

    N = 1
    e_sec = np.array([5.0], dtype=np.float32)
    sec_E_list = np.zeros((N, K_MAX), dtype=np.float32)
    sec_E_list[0, 0] = 5.0
    sec_dir_list = np.zeros((N, K_MAX, 3), dtype=np.float32)
    sec_dir_list[0, 0] = [0, 0, 1]
    sec_pdg_list = np.zeros((N, K_MAX), dtype=np.int64)
    sec_pdg_list[0, 0] = 2212  # proton
    sec_valid = np.zeros((N, K_MAX), dtype=bool)
    sec_valid[0, 0] = True
    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)

    sec_cont = encode_secondaries(
        sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
    )
    norm = Normalizer()
    norm.mean = np.array([-2.0, 0.5], dtype=np.float32)
    norm.std = np.array([3.0, 1.5], dtype=np.float32)
    sec_cont_normed = sec_cont.copy()
    sec_cont_normed[:, :, 4:6] = norm.transform(
        sec_cont[:, :, 4:6].reshape(-1, 2)
    ).reshape(N, K_MAX, 2)

    n_sec = np.array([1])
    _, _, sec_mass, sec_charge, _ = decode_secondaries(
        sec_cont_normed, n_sec, e_sec, pre_dir, sec_phys_normalizer=norm
    )
    assert sec_mass[0, 0] == pytest.approx(938.27208943, abs=1e-2)
    assert sec_charge[0, 0] == pytest.approx(1.0, abs=1e-4)