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