giant/tests/test_network.py

import torch
from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
from giant.model.network import (
    ConditionEncoder,
    SinusoidalEmbedding,
    Stage1Model,
    Stage2OneShot,
    cat_col_layout,
    stage2_trunk_sec_dim,
    stage2_type_dim,
)

PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
ONEHOT_PARTICLE_CFG = {"type": "onehot", "emb_dim": 6, "n_layers": 1}
ONEHOT_MATERIAL_CFG = {"type": "onehot", "emb_dim": 4, "n_layers": 1}


def test_sinusoidal_embedding_shape():
    emb = SinusoidalEmbedding(64)
    t = torch.rand(16)
    assert emb(t).shape == (16, 64)


def test_sinusoidal_embedding_batch_1():
    emb = SinusoidalEmbedding(32)
    t = torch.tensor([0.5])
    assert emb(t).shape == (1, 32)


def test_stage1_model_output_shape():
    B = 8
    model = Stage1Model(
        pdg_vocab=5,
        mat_vocab=3,
        particle_cfg=PARTICLE_CFG,
        material_cfg=MATERIAL_CFG,
        n_sec_head_k_max=15,
    )
    x_t = torch.randn(B, 9)
    t = torch.rand(B)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.stack(
        [
            torch.randint(0, 5, (B,)),
            torch.randint(0, 3, (B,)),
        ],
        dim=1,
    )
    out = model(x_t, cond_cont, cond_cat, t=t)
    assert out.shape == (B, 9)


def test_stage1_model_gradients_flow():
    B = 4
    model = Stage1Model(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=PARTICLE_CFG,
        material_cfg=MATERIAL_CFG,
        hidden_dim=32,
        n_res_blocks=2,
        n_sec_head_k_max=15,
    )
    x_t = torch.randn(B, 9)
    t = torch.rand(B)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    # Both paths must be exercised to get gradients through all parameters.
    flow_loss = model(x_t, cond_cont, cond_cat, t=t).sum()
    nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum()
    (flow_loss + nsec_loss).backward()
    for name, p in model.named_parameters():
        assert p.grad is not None, f"no grad for {name}"


def test_stage1_model_no_n_sec_head_by_default():
    """Fresh v0.3.0 construction (no n_sec_head_k_max) has no n_sec head —
    decision 1 (docs/v0.3.0-design.md §2) moves it to stage 2."""
    model = Stage1Model(
        pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG
    )
    assert model.n_sec_head is None


# --- cat_col_layout / stage2_type_dim / stage2_trunk_sec_dim ---------------


def test_cat_col_layout_neither_onehot():
    assert cat_col_layout("physical", "embedding") == (None, None)


def test_cat_col_layout_particle_only():
    assert cat_col_layout("onehot", "physical") == (2, None)


def test_cat_col_layout_material_only():
    assert cat_col_layout("physical", "onehot") == (None, 2)


def test_cat_col_layout_both_onehot_particle_then_material():
    assert cat_col_layout("onehot", "onehot") == (2, 3)


def test_stage2_type_dim_physical_is_particle_phys_dim():
    assert stage2_type_dim({"target": "physical"}, emb_dim=16) == PARTICLE_PHYS_DIM


def test_stage2_type_dim_onehot_and_embedding_are_emb_dim():
    assert stage2_type_dim({"target": "onehot"}, emb_dim=16) == 16
    assert stage2_type_dim({"target": "embedding"}, emb_dim=16) == 16


def test_stage2_trunk_sec_dim_physical_matches_v02_sec_dim():
    k_max = 15
    assert (
        stage2_trunk_sec_dim({"target": "physical"}, "flow", k_max, emb_dim=16)
        == k_max * SEC_SLOT_DIM
    )
    assert (
        stage2_trunk_sec_dim({"target": "physical"}, "wgan", k_max, emb_dim=16)
        == k_max * SEC_SLOT_DIM
    )


def test_stage2_trunk_sec_dim_onehot_wgan_folds_type_in():
    k_max = 15
    assert stage2_trunk_sec_dim(
        {"target": "onehot"}, "wgan", k_max, emb_dim=16
    ) == k_max * (CONT_SLOT_DIM + 16)


def test_stage2_trunk_sec_dim_onehot_flow_excludes_type():
    k_max = 15
    assert (
        stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16)
        == k_max * CONT_SLOT_DIM
    )


# --- ConditionEncoder onehot mode -------------------------------------------


def test_condition_encoder_onehot_forward_shape_and_gradients():
    B = 8
    particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
    material_emb_dim = int(ONEHOT_MATERIAL_CFG["emb_dim"])
    enc = ConditionEncoder(
        pdg_vocab=5,
        mat_vocab=3,
        particle_cfg=ONEHOT_PARTICLE_CFG,
        material_cfg=ONEHOT_MATERIAL_CFG,
        out_dim=32,
    )
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.stack(
        [
            torch.randint(0, 5, (B,)),
            torch.randint(0, 3, (B,)),
            torch.randint(0, particle_emb_dim, (B,)),
            torch.randint(0, material_emb_dim, (B,)),
        ],
        dim=1,
    )
    out = enc(cond_cont, cond_cat)
    assert out.shape == (B, 32)
    # onehot itself is unlearned, but the fusion MLP downstream still has
    # gradients — the encoder as a whole must still be trainable.
    out.sum().backward()
    assert enc.mlp[0].weight.grad is not None


def test_condition_encoder_onehot_is_a_true_one_hot_vector():
    """The onehot axis feeds a fixed, unlearned one-hot into the fusion MLP —
    verify the concatenated input segment really is one-hot, not e.g. an
    accidentally-learned embedding."""
    B = 4
    particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
    enc = ConditionEncoder(
        pdg_vocab=5,
        mat_vocab=3,
        particle_cfg=ONEHOT_PARTICLE_CFG,
        material_cfg={"type": "physical", "emb_dim": 4, "n_layers": 1},
        out_dim=16,
    )
    cond_cont = torch.zeros(B, COND_DIM)
    idx = torch.tensor([0, 1, 2, 5])
    cond_cat = torch.stack(
        [
            torch.zeros(B, dtype=torch.long),
            torch.zeros(B, dtype=torch.long),
            idx.clamp(max=particle_emb_dim - 1),
        ],
        dim=1,
    )
    pdg_e = enc._particle_embed(cond_cont, cond_cat)
    assert pdg_e.shape == (B, particle_emb_dim)
    assert torch.all(pdg_e.sum(dim=-1) == 1.0)


# --- Stage2OneShot particle_type architecture (docs/v0.3.0-design.md decision 2) --


def _build_stage2(target: str, generator: str, emb_dim: int = 6) -> Stage2OneShot:
    particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1}
    if target != "physical":
        particle_cfg = dict(particle_cfg)
        if target == "embedding":
            particle_cfg["type"] = "embedding"
    k_max = 5
    sec_dim = stage2_trunk_sec_dim({"target": target}, generator, k_max, emb_dim)
    return Stage2OneShot(
        pdg_vocab=5,
        mat_vocab=3,
        particle_cfg=particle_cfg,
        material_cfg=MATERIAL_CFG,
        hidden_dim=16,
        n_res_blocks=1,
        cond_out_dim=16,
        context_dim=8,
        sec_dim=sec_dim,
        generator=generator,
        k_max=k_max,
        particle_type_cfg={"target": target, "lambda": 1.0},
    )


def test_stage2_oneshot_physical_has_no_type_head_regardless_of_generator():
    assert _build_stage2("physical", "flow").type_head is None
    assert _build_stage2("physical", "wgan").type_head is None


def test_stage2_oneshot_onehot_flow_has_type_head():
    model = _build_stage2("onehot", "flow")
    assert model.type_head is not None


def test_stage2_oneshot_onehot_wgan_has_no_type_head():
    """Under wgan the type slice is folded into forward()'s own output and
    relaxed via ST-Gumbel by the trainer — no separate head needed."""
    model = _build_stage2("onehot", "wgan")
    assert model.type_head is None


def test_stage2_oneshot_embedding_flow_has_type_head():
    model = _build_stage2("embedding", "flow")
    assert model.type_head is not None


def test_stage2_oneshot_predict_type_shape():
    B, k_max, emb_dim = 4, 5, 6
    model = _build_stage2("onehot", "flow", emb_dim=emb_dim)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
    out = model.predict_type(cond_cont, cond_cat, stage1_out)
    assert out.shape == (B, k_max, emb_dim)


def test_stage2_oneshot_predict_type_raises_when_no_type_head():
    model = _build_stage2("physical", "flow")
    cond_cont = torch.randn(2, COND_DIM)
    cond_cat = torch.zeros(2, 2, dtype=torch.long)
    stage1_out = torch.randn(2, 9)
    try:
        model.predict_type(cond_cont, cond_cat, stage1_out)
        raise AssertionError("expected RuntimeError")
    except RuntimeError:
        pass


def test_stage2_oneshot_forward_shape_onehot_wgan():
    B, k_max, emb_dim = 4, 5, 6
    model = _build_stage2("onehot", "wgan", emb_dim=emb_dim)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
    z = torch.randn(B, model.noise_dim)
    out = model(z, cond_cont, cond_cat, stage1_out)
    assert out.shape == (B, k_max * (CONT_SLOT_DIM + emb_dim))


def test_stage2_oneshot_forward_shape_onehot_flow_excludes_type():
    B, k_max, emb_dim = 4, 5, 6
    model = _build_stage2("onehot", "flow", emb_dim=emb_dim)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
    x_t = torch.randn(B, k_max * CONT_SLOT_DIM)
    t = torch.rand(B)
    out = model(x_t, cond_cont, cond_cat, stage1_out, t=t)
    assert out.shape == (B, k_max * CONT_SLOT_DIM)