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, EMB_DIM, K_MAX, 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, snap_type_to_pdg_idx


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


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


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


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


def test_pdg_embedding_weight_shape():
    model = _stage1(pdg=5, mat=2)
    w = model.pdg_embedding_weight()
    assert w.shape == (5, EMB_DIM)


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


def test_sec_decoder_output_shape():
    B = 8
    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 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_type_emb, 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_type_emb.shape == (B, K_MAX, EMB_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()


def test_snap_type_to_pdg_idx_shape():
    B, pdg_vocab = 4, 5
    emb_weight = torch.randn(pdg_vocab, EMB_DIM)
    sec_type_emb = torch.randn(B, K_MAX, EMB_DIM)
    idx = snap_type_to_pdg_idx(sec_type_emb, emb_weight)
    assert idx.shape == (B, K_MAX)
    assert idx.dtype == torch.int64
    assert (idx >= 0).all() and (idx < pdg_vocab).all()


# ── 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_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

    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)
    assert sec_cont.shape == (N, K_MAX, 4)
    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)


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


def _random_sec_cont(rng, N, stick_logit_scale=1.0):
    sec_cont = rng.standard_normal((N, K_MAX, 4)).astype(np.float32)
    sec_cont[:, :, 0] *= stick_logit_scale
    dirs = sec_cont[:, :, 1:]
    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)
    sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
    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, _sec_pdg, sec_valid = decode_secondaries(
        sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22}
    )

    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)
    sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
    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, _sec_pdg, sec_valid = decode_secondaries(
        sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22}
    )

    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
    sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
    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, _sec_pdg, sec_valid = decode_secondaries(
        sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22}
    )

    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_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)

    sec_E_small, _, _, sec_valid = decode_secondaries(
        sec_cont,
        sec_pdg_pred,
        n_sec,
        np.array([5.0], dtype=np.float32),
        pre_dir,
        {0: 22},
    )
    sec_E_large, _, _, _ = decode_secondaries(
        sec_cont,
        sec_pdg_pred,
        n_sec,
        np.array([50.0], dtype=np.float32),
        pre_dir,
        {0: 22},
    )

    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)