giant/tests/test_objectives.py

"""Tests for `giant/model/objectives.py` — the generator/objective registry
(gitea #32) that replaced bare `generator in ("flow", "ddpm", "wgan")`
string checks scattered across models.py/sample.py/builders.py/
stage2_inputs.py/trainers.py."""

import pytest
import torch

from giant.config import ConditioningAxisConfig
from giant.constants import COND_DIM, CONT_SLOT_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM
from giant.model.network import (
    OBJECTIVE_REGISTRY,
    DdpmObjective,
    FlowObjective,
    Stage1Model,
    Stage2Autoregressive,
    Stage2OneShot,
    WganObjective,
    build_objective,
)
from giant.model.schedule import CosineSchedule, flow_matching_loss, flow_matching_loss_secondary

_PHYS_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)


def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]:
    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


# ── registry ─────────────────────────────────────────────────────────────


def test_registry_has_exactly_the_three_known_objectives():
    assert set(OBJECTIVE_REGISTRY) == {"flow", "ddpm", "wgan"}


def test_build_objective_returns_correct_concrete_type():
    assert isinstance(build_objective("flow"), FlowObjective)
    assert isinstance(build_objective("ddpm"), DdpmObjective)
    assert isinstance(build_objective("wgan"), WganObjective)


def test_build_objective_unknown_name_raises():
    with pytest.raises(ValueError, match="unknown generator/objective"):
        build_objective("bogus")


def test_build_objective_filters_kwargs_by_signature():
    # FlowObjective takes no constructor args — n_steps (a DdpmObjective-only
    # kwarg) must be silently dropped, not raise a TypeError.
    build_objective("flow", n_steps=500)
    ddpm = build_objective("ddpm", n_steps=250)
    assert isinstance(ddpm, DdpmObjective)
    assert ddpm.n_steps == 250


# ── flags ────────────────────────────────────────────────────────────────


def test_flow_objective_flags():
    obj = build_objective("flow")
    assert obj.needs_time is True
    assert obj.is_adversarial is False
    assert obj.folds_type_slice is False
    assert obj.supports_stage2_decoder is True


def test_ddpm_objective_flags():
    obj = build_objective("ddpm")
    assert obj.needs_time is True
    assert obj.is_adversarial is False
    assert obj.folds_type_slice is False
    assert obj.supports_stage2_decoder is False


def test_wgan_objective_flags():
    obj = build_objective("wgan")
    assert obj.needs_time is False
    assert obj.is_adversarial is True
    assert obj.folds_type_slice is True
    assert obj.supports_stage2_decoder is True


# ── trunk_in_dim ─────────────────────────────────────────────────────────


def test_trunk_in_dim_flow_and_ddpm_pass_through_out_dim():
    assert build_objective("flow").trunk_in_dim(out_dim=9, noise_dim=8) == 9
    assert build_objective("ddpm").trunk_in_dim(out_dim=9, noise_dim=8) == 9


def test_trunk_in_dim_wgan_uses_noise_dim():
    assert build_objective("wgan").trunk_in_dim(out_dim=9, noise_dim=8) == 8


# ── ddpm schedule ────────────────────────────────────────────────────────


def test_ddpm_build_schedule_has_requested_length():
    schedule = build_objective("ddpm").build_schedule(n_steps=17, device=torch.device("cpu"))
    assert isinstance(schedule, CosineSchedule)
    assert schedule.T == 17


def test_flow_and_wgan_build_schedule_is_none():
    assert build_objective("flow").build_schedule(100, torch.device("cpu")) is None
    assert build_objective("wgan").build_schedule(100, torch.device("cpu")) is None


# ── stage1_loss parity ──────────────────────────────────────────────────


def test_flow_objective_stage1_loss_matches_direct_call():
    torch.manual_seed(0)
    model = Stage1Model(
        pdg_vocab=3, mat_vocab=2, particle_cfg=_PHYS_CFG, material_cfg=_PHYS_CFG, hidden_dim=16, n_res_blocks=1
    )
    cond_cont, cond_cat = _cond(4)
    x1 = torch.randn(4, X_DIM)

    torch.manual_seed(1)
    expected = flow_matching_loss(model, x1, cond_cont, cond_cat)
    torch.manual_seed(1)
    actual = build_objective("flow").stage1_loss(model, x1, cond_cont, cond_cat)
    assert torch.allclose(actual, expected)


def test_ddpm_objective_stage1_loss_matches_direct_call():
    torch.manual_seed(0)
    model = Stage1Model(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=_PHYS_CFG,
        material_cfg=_PHYS_CFG,
        hidden_dim=16,
        n_res_blocks=1,
        generator="ddpm",
    )
    cond_cont, cond_cat = _cond(4)
    x1 = torch.randn(4, X_DIM)
    objective = build_objective("ddpm", n_steps=50)
    schedule = objective.build_schedule(50, torch.device("cpu"))
    assert isinstance(schedule, CosineSchedule)

    torch.manual_seed(1)
    expected = schedule.loss(model, x1, cond_cont, cond_cat)
    torch.manual_seed(1)
    actual = objective.stage1_loss(model, x1, cond_cont, cond_cat, schedule=schedule)
    assert torch.allclose(actual, expected)


def test_ddpm_objective_stage1_loss_requires_a_schedule():
    model = Stage1Model(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=_PHYS_CFG,
        material_cfg=_PHYS_CFG,
        hidden_dim=16,
        n_res_blocks=1,
        generator="ddpm",
    )
    cond_cont, cond_cat = _cond(4)
    with pytest.raises(AssertionError):
        build_objective("ddpm").stage1_loss(model, torch.randn(4, X_DIM), cond_cont, cond_cat, schedule=None)


# ── stage2_loss dispatch ─────────────────────────────────────────────────


def test_flow_objective_stage2_loss_one_shot_matches_direct_call():
    torch.manual_seed(0)
    B, k_max = 4, 5
    sec_dim = k_max * SEC_SLOT_DIM
    model = Stage2OneShot(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=_PHYS_CFG,
        material_cfg=_PHYS_CFG,
        hidden_dim=16,
        n_res_blocks=1,
        generator="flow",
        sec_dim=sec_dim,
        k_max=k_max,
    )
    cond_cont, cond_cat = _cond(B)
    stage1_ctx = torch.randn(B, X_DIM)
    x1_s2 = torch.randn(B, sec_dim)
    sec_mask = torch.ones(B, k_max, dtype=torch.bool)

    torch.manual_seed(1)
    expected = flow_matching_loss_secondary(model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None)
    torch.manual_seed(1)
    actual = build_objective("flow").stage2_loss(
        model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None, ar_inputs=None
    )
    assert torch.allclose(actual, expected)


def test_flow_objective_stage2_loss_dispatches_to_ar_when_ar_inputs_given():
    torch.manual_seed(0)
    B, k_max = 4, 5
    model = Stage2Autoregressive(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=_PHYS_CFG,
        material_cfg=_PHYS_CFG,
        hidden_dim=16,
        n_res_blocks=1,
        generator="flow",
        k_max=k_max,
    )
    cond_cont, cond_cat = _cond(B)
    stage1_ctx = torch.randn(B, X_DIM)
    token_dim = CONT_SLOT_DIM + PARTICLE_PHYS_DIM
    x1_s2 = torch.randn(B, k_max, token_dim)
    sec_mask = torch.ones(B, k_max, dtype=torch.bool)
    ar_inputs = {
        "history_feat": torch.randn(B, k_max, token_dim),
        "has_prev": torch.ones(B, k_max, dtype=torch.bool),
        "remaining_frac": torch.rand(B, k_max),
        "slot_idx": torch.linspace(0, 1, k_max).unsqueeze(0).expand(B, -1),
    }

    loss = build_objective("flow").stage2_loss(
        model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None, ar_inputs=ar_inputs
    )
    assert loss.dim() == 0
    assert torch.isfinite(loss)


def test_ddpm_objective_stage2_loss_not_implemented():
    dummy_model = torch.nn.Module()
    dummy = torch.zeros(1)
    with pytest.raises(NotImplementedError):
        build_objective("ddpm").stage2_loss(
            dummy_model, dummy, dummy, dummy, dummy, torch.ones(1, 1, dtype=torch.bool), type_dim=None
        )


def test_wgan_objective_has_no_loss_methods():
    dummy_model = torch.nn.Module()
    dummy = torch.zeros(1)
    objective = build_objective("wgan")
    with pytest.raises(NotImplementedError):
        objective.stage1_loss(dummy_model, dummy, dummy, dummy)
    with pytest.raises(NotImplementedError):
        objective.stage2_loss(
            dummy_model, dummy, dummy, dummy, dummy, torch.ones(1, 1, dtype=torch.bool), type_dim=None
        )