giant/tests/test_train.py

"""Tests for giant/train.py."""

import copy
import tempfile
from pathlib import Path

import pytest
import torch

from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM
from giant.model.network import build_critics, build_models
from giant.train import (
    FlowDDPMStageTrainer,
    WGANStageTrainer,
    _gumbel_tau,
    _wandb_run_config,
    train,
)

PDG_VOCAB = 6
MAT_VOCAB = 3


def test_gumbel_tau_at_step_zero_is_start():
    assert _gumbel_tau(0, 1000, 1.0, 0.1) == 1.0


def test_gumbel_tau_at_total_steps_is_end():
    assert abs(_gumbel_tau(1000, 1000, 1.0, 0.1) - 0.1) < 1e-9


def test_gumbel_tau_interpolates_linearly_midway():
    assert abs(_gumbel_tau(500, 1000, 1.0, 0.1) - 0.55) < 1e-9


def test_gumbel_tau_clamps_beyond_total_steps():
    assert _gumbel_tau(5000, 1000, 1.0, 0.1) == _gumbel_tau(1000, 1000, 1.0, 0.1)


def test_gumbel_tau_handles_zero_total_steps():
    # total_steps=0 is guarded to 1 internally: step=0 gives zero progress
    # (still tau_start), any step>=1 immediately clamps to full progress.
    assert _gumbel_tau(0, 0, 1.0, 0.1) == 1.0
    assert abs(_gumbel_tau(1, 0, 1.0, 0.1) - 0.1) < 1e-9


def test_wandb_run_config_includes_full_cfg_and_param_counts():
    cfg = {
        "train": {"lr": 3e-4},
        "conditioning": {"out_dim": 128},
        "stage1_model": {"generator": "flow"},
        "stage2_model": {"generator": "wgan"},
    }
    wcfg = _wandb_run_config(
        cfg, model_config={"pdg_vocab": 3}, param_counts={"stage1": 100}
    )
    assert wcfg["train"] == {"lr": 3e-4}
    assert wcfg["stage1_model"] == {"generator": "flow"}
    assert wcfg["stage2_model"] == {"generator": "wgan"}
    assert wcfg["model_config"] == {"pdg_vocab": 3}
    assert wcfg["param_counts"] == {"stage1": 100}


def test_wandb_run_config_handles_missing_model_config():
    cfg = {"train": {}, "conditioning": {}, "stage1_model": {}, "stage2_model": {}}
    wcfg = _wandb_run_config(cfg, model_config=None, param_counts={})
    assert wcfg["model_config"] == {}


# --- end-to-end train() integration tests -----------------------------------

PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}


def _base_cfg():
    return {
        "conditioning": {
            "out_dim": 32,
            "share_stages": False,
            "particle": dict(PARTICLE_CFG),
            "material": dict(MATERIAL_CFG),
        },
        "stage1_model": {
            "active": True,
            "generator": "flow",
            "hidden_dim": 24,
            "n_res_blocks": 2,
            "dropout": 0.0,
            "lambda": 1.0,
            "flow": {"time_dim": 16},
            "ddpm": {"time_dim": 16, "n_steps": 50},
            "wgan": {
                "noise_dim": 16,
                "n_critic": 2,
                "gp_weight": 10.0,
                "critic_lr": 0.0,
            },
            "router": {"enabled": False},
        },
        "stage2_model": {
            "active": True,
            "decoder": "one_shot",
            "generator": "wgan",
            "hidden_dim": 24,
            "n_res_blocks": 2,
            "dropout": 0.0,
            "lambda": 1.0,
            "k_max": K_MAX,
            "context_dim": 16,
            "n_sec": {"mode": "head", "lambda": 0.1},
            "flow": {"time_dim": 16},
            "ddpm": {"time_dim": 16, "n_steps": 50},
            "wgan": {
                "noise_dim": 16,
                "n_critic": 2,
                "gp_weight": 10.0,
                "critic_lr": 0.0,
            },
            "router": {"enabled": False, "tie_to_stage1": False},
        },
        "train": {
            "epochs": 2,
            "batch_size": 8,
            "lr": 3e-4,
            "weight_decay": 0.01,
            "ema_decay": 0.999,
            "warmup_epochs": 0,
            "val_fraction": 0.1,
            "max_val_batches": 0,
            "num_workers": 0,
            "seed": 0,
            "validate_every": 0,
            "validate_steps": 2,
            "wandb": False,
        },
    }


def _fake_batches(n_batches, batch_size, seed=0):
    g = torch.Generator().manual_seed(seed)
    batches = []
    for _ in range(n_batches):
        cond_cont = torch.randn(batch_size, COND_DIM, generator=g)
        cond_cat = torch.stack(
            [
                torch.randint(0, PDG_VOCAB, (batch_size,), generator=g),
                torch.randint(0, MAT_VOCAB, (batch_size,), generator=g),
            ],
            dim=1,
        )
        x1 = torch.randn(batch_size, X_DIM, generator=g)
        n_sec = torch.randint(0, K_MAX, (batch_size,), generator=g)
        sec_cont = torch.randn(batch_size, K_MAX, SEC_SLOT_DIM, generator=g)
        proc_idx = torch.zeros(batch_size, dtype=torch.long)
        batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx))
    return batches


def _model_config(cfg):
    return {
        "pdg_vocab": PDG_VOCAB,
        "mat_vocab": MAT_VOCAB,
        "conditioning": cfg["conditioning"],
        "stage1_model": cfg["stage1_model"],
        "stage2_model": cfg["stage2_model"],
    }


def _run_train(cfg, out_dir, resume_path=None):
    model_config = _model_config(cfg)
    models = build_models(model_config)
    critics = build_critics(model_config)
    train_loader = _fake_batches(4, cfg["train"]["batch_size"])
    val_loader = _fake_batches(2, cfg["train"]["batch_size"], seed=1)
    train(
        cfg=cfg,
        models=models,
        critics=critics,
        train_loader=train_loader,
        val_loader=val_loader,
        device=torch.device("cpu"),
        out_dir=out_dir,
        normalizer_dict={"cond": {}, "target": {}, "sec_phys": {}},
        pdg_map={"22": 0},
        mat_map={"G4_AIR": 0},
        proc_map=None,
        model_config=model_config,
        total_train_batches=4,
        resume_path=resume_path,
    )


@pytest.mark.parametrize(
    "label,mutate",
    [
        ("both_flow", lambda cfg: None),
        ("both_wgan", lambda cfg: cfg["stage1_model"].__setitem__("generator", "wgan")),
        (
            "mixed_stage1_flow_stage2_wgan",
            lambda cfg: None,  # already the default
        ),
        (
            "mixed_stage1_wgan_stage2_flow",
            lambda cfg: (
                cfg["stage1_model"].__setitem__("generator", "wgan"),
                cfg["stage2_model"].__setitem__("generator", "flow"),
            ),
        ),
        ("stage1_only", lambda cfg: cfg["stage2_model"].__setitem__("active", False)),
        ("stage2_only", lambda cfg: cfg["stage1_model"].__setitem__("active", False)),
        (
            "both_ddpm_stage1_flow_stage2",
            lambda cfg: (
                cfg["stage1_model"].__setitem__("generator", "ddpm"),
                cfg["stage2_model"].__setitem__("generator", "flow"),
            ),
        ),
        (
            "routed_stage1_energy_gumbel",
            lambda cfg: cfg["stage1_model"].__setitem__(
                "router",
                {
                    "enabled": True,
                    "type": "energy",
                    "n_experts": 3,
                    "temperature": 0.5,
                    "learn_centers": True,
                    "lambda_balance": 0.1,
                    "lambda_entropy": 0.01,
                    "gumbel": True,
                    "gumbel_tau_start": 1.0,
                    "gumbel_tau_end": 0.1,
                },
            ),
        ),
    ],
)
def test_train_end_to_end(label, mutate):
    cfg = _base_cfg()
    mutate(cfg)
    with tempfile.TemporaryDirectory() as tmp:
        out_dir = Path(tmp) / "run"
        _run_train(cfg, out_dir)
        assert (out_dir / "last.pt").exists()
        assert (out_dir / "metrics.csv").exists()
        ckpt = torch.load(out_dir / "last.pt", weights_only=False)
        if cfg["stage1_model"]["active"]:
            assert "model" in ckpt
        else:
            assert "model" not in ckpt
        if cfg["stage2_model"]["active"]:
            assert "sec_decoder" in ckpt
        else:
            assert "sec_decoder" not in ckpt


def test_train_resume_continues_from_checkpoint():
    cfg = _base_cfg()
    with tempfile.TemporaryDirectory() as tmp:
        out_dir = Path(tmp) / "run"
        _run_train(cfg, out_dir)
        ckpt_before = torch.load(out_dir / "last.pt", weights_only=False)
        assert ckpt_before["epoch"] == 2

        cfg2 = copy.deepcopy(cfg)
        cfg2["train"]["epochs"] = 3
        _run_train(cfg2, out_dir, resume_path=out_dir / "last.pt")
        ckpt_after = torch.load(out_dir / "last.pt", weights_only=False)
        assert ckpt_after["epoch"] == 3
        assert ckpt_after["global_step"] > ckpt_before["global_step"]


def test_train_raises_when_no_active_stage():
    cfg = _base_cfg()
    cfg["stage1_model"]["active"] = False
    cfg["stage2_model"]["active"] = False
    model_config = _model_config(cfg)
    models = build_models(model_config)
    critics = build_critics(model_config)
    with tempfile.TemporaryDirectory() as tmp:
        with pytest.raises(ValueError, match="no active stage"):
            train(
                cfg=cfg,
                models=models,
                critics=critics,
                train_loader=_fake_batches(1, 8),
                val_loader=_fake_batches(1, 8),
                device=torch.device("cpu"),
                out_dir=Path(tmp) / "run",
                model_config=model_config,
                total_train_batches=1,
            )


def test_metrics_csv_columns_are_stage_prefixed():
    cfg = _base_cfg()
    with tempfile.TemporaryDirectory() as tmp:
        out_dir = Path(tmp) / "run"
        _run_train(cfg, out_dir)
        header = (out_dir / "metrics.csv").read_text().splitlines()[0].split(",")
        assert "stage1_train_loss" in header
        assert "stage2_train_d_loss" in header
        assert "val_loss" in header
        assert "epoch" in header


def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch():
    """Regression test: on a non-generator-step batch, if this stage's model
    has no n_sec_head (n_sec defaults to stage 2, decision 1), g_loss is a
    graph-less zero — .backward() must not be called on it."""
    cfg = _base_cfg()
    cfg["stage1_model"]["generator"] = "wgan"
    model_config = _model_config(cfg)
    models = build_models(model_config)
    critics = build_critics(model_config)
    assert models["stage1"] is not None and critics["stage1"] is not None
    trainer = WGANStageTrainer(
        name="stage1",
        model=models["stage1"],
        critic=critics["stage1"],
        is_stage2=False,
        lambda_weight=1.0,
        n_sec_lambda=0.1,
        n_critic=1000,  # never a generator step in this test
        gp_weight=10.0,
        lr=3e-4,
        critic_lr=0.0,
        ema_decay=0.0,
        warmup_epochs=0,
        epochs=1,
        steps_per_epoch=4,
        device=torch.device("cpu"),
    )
    assert trainer.model.n_sec_head is None
    batch = _fake_batches(1, 8)[0]
    stats = trainer.step(batch, torch.device("cpu"), global_step=1)
    assert stats["did_g_step"] is False


def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
    with pytest.raises(NotImplementedError):
        FlowDDPMStageTrainer(
            name="stage2",
            model=torch.nn.Linear(1, 1),
            is_stage2=True,
            generator="ddpm",
            lambda_weight=1.0,
            n_sec_lambda=0.1,
            lambda_balance=0.0,
            lambda_proc=0.0,
            lambda_entropy=0.0,
            gumbel_tau_start=1.0,
            gumbel_tau_end=0.1,
            lr=3e-4,
            weight_decay=0.01,
            ema_decay=0.0,
            warmup_epochs=0,
            epochs=1,
            steps_per_epoch=1,
            ddpm_n_steps=50,
            device=torch.device("cpu"),
        )