"""Tests for giant/train.py.""" import copy import csv import tempfile from pathlib import Path import pytest import torch from giant.constants import ( COND_DIM, CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM, ) from giant.model.network import build_critics, build_models from giant.train import ( FlowDDPMStageTrainer, WGANStageTrainer, _ar_has_prev, _assemble_stage2_ar_inputs, _assemble_stage2_ar_target, _assemble_stage2_real, _build_stage_trainers, _gumbel_tau, _relax_onehot_type_slice, _remaining_energy_fraction, _shift_prev, _stick_fraction, _type_repr, _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"] == {} # --- AR helper functions (v0.3.0 step 5, docs/v0.3.0-design.md §6) --------- def test_stick_fraction_matches_sigmoid_of_logit(): sec_cont = torch.zeros(2, 3, SEC_SLOT_DIM) sec_cont[..., 0] = torch.tensor([[0.0, 2.0, -2.0], [1.0, -1.0, 0.0]]) frac = _stick_fraction(sec_cont) assert torch.allclose(frac, torch.sigmoid(sec_cont[..., 0])) def test_remaining_energy_fraction_hand_computed(): fraction = torch.tensor([[0.5, 0.5, 1.0]]) remaining = _remaining_energy_fraction(fraction) assert torch.allclose(remaining, torch.tensor([[1.0, 0.5, 0.25]])) def test_shift_prev_shifts_and_zero_pads_slot0(): x = torch.arange(2 * 4 * 3).reshape(2, 4, 3).float() shifted = _shift_prev(x) assert torch.all(shifted[:, 0] == 0) assert torch.equal(shifted[:, 1:], x[:, :-1]) def test_ar_has_prev_false_only_at_slot_zero(): has_prev = _ar_has_prev(5, torch.device("cpu")) assert has_prev.shape == (1, 5) assert has_prev.tolist() == [[False, True, True, True, True]] @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"]) def test_type_repr_shapes_and_values(target): B, K, emb_dim = 3, 4, 6 sec_cont = torch.randn(B, K, SEC_SLOT_DIM) sec_type_idx = torch.randint(0, emb_dim, (B, K)) cond_enc = torch.nn.Module() if target == "embedding": cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim) repr_ = _type_repr(sec_type_idx, sec_cont, {"target": target}, cond_enc, emb_dim) expected_width = PARTICLE_PHYS_DIM if target == "physical" else emb_dim assert repr_.shape == (B, K, expected_width) if target == "physical": assert torch.equal( repr_, sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM] ) if target == "onehot": assert torch.all(repr_.sum(-1) == 1.0) @pytest.mark.parametrize( "target,generator", [ ("physical", "flow"), ("physical", "wgan"), ("onehot", "flow"), ("onehot", "wgan"), ("embedding", "flow"), ("embedding", "wgan"), ], ) def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened( target, generator ): """Regression test tying the refactor together: _assemble_stage2_real is now defined as _assemble_stage2_ar_target(...).flatten(1).""" B, emb_dim = 4, 6 sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM) sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX)) cond_enc = torch.nn.Module() if target == "embedding": cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim) particle_type_cfg = {"target": target} flat = _assemble_stage2_real( sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim ) unflat = _assemble_stage2_ar_target( sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim ) assert torch.equal(unflat.flatten(1), flat) def test_assemble_stage2_ar_inputs_shapes_and_history_feat_width(): B, emb_dim = 3, 6 sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM) sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX)) cond_enc = torch.nn.Module() out = _assemble_stage2_ar_inputs( sec_cont, sec_type_idx, {"target": "physical"}, cond_enc, emb_dim ) assert out["history_feat"].shape == (B, K_MAX, CONT_SLOT_DIM + PARTICLE_PHYS_DIM) assert out["has_prev"].shape == (B, K_MAX) assert out["remaining_frac"].shape == (B, K_MAX) assert out["slot_idx"].shape == (B, K_MAX) assert torch.all(out["slot_idx"][:, 0] == 0.0) assert torch.all(out["slot_idx"][:, -1] == 1.0) def test_relax_onehot_type_slice_grad_probe_populates_both_norms(): B, k_max, cont_dim, type_dim = 4, K_MAX, CONT_SLOT_DIM, 6 x_flat = torch.randn(B, k_max * (cont_dim + type_dim), requires_grad=True) grad_probe: dict[str, float] = {} out = _relax_onehot_type_slice( x_flat, k_max, cont_dim, type_dim, tau=0.5, grad_probe=grad_probe ) out.sum().backward() assert grad_probe["cont"] >= 0.0 assert grad_probe["type"] >= 0.0 def test_relax_onehot_type_slice_grad_probe_none_is_backward_compatible(): B, k_max, cont_dim, type_dim = 4, K_MAX, CONT_SLOT_DIM, 6 x_flat = torch.randn(B, k_max * (cont_dim + type_dim), requires_grad=True) out = _relax_onehot_type_slice(x_flat, k_max, cont_dim, type_dim, tau=0.5) out.sum().backward() assert x_flat.grad is not None # --- 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) sec_type_idx = torch.zeros(batch_size, K_MAX, dtype=torch.long) batches.append( (cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_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, }, ), ), ( "stage2_onehot_target_wgan", lambda cfg: cfg["stage2_model"].__setitem__( "particle_type", {"target": "onehot", "lambda": 1.0} ), ), ( "stage2_onehot_target_flow", lambda cfg: ( cfg["stage2_model"].__setitem__("generator", "flow"), cfg["stage2_model"].__setitem__( "particle_type", {"target": "onehot", "lambda": 1.0} ), ), ), ( "stage2_embedding_target_wgan", lambda cfg: ( cfg["conditioning"]["particle"].__setitem__("type", "embedding"), cfg["conditioning"]["material"].__setitem__("type", "embedding"), cfg["stage2_model"].__setitem__( "particle_type", {"target": "embedding", "lambda": 1.0} ), ), ), ( "stage2_embedding_target_flow", lambda cfg: ( cfg["conditioning"]["particle"].__setitem__("type", "embedding"), cfg["conditioning"]["material"].__setitem__("type", "embedding"), cfg["stage2_model"].__setitem__("generator", "flow"), cfg["stage2_model"].__setitem__( "particle_type", {"target": "embedding", "lambda": 1.0} ), ), ), ( "ar_wgan_onehot", lambda cfg: ( cfg["stage2_model"].__setitem__("decoder", "autoregressive"), cfg["stage2_model"].__setitem__( "particle_type", {"target": "onehot", "lambda": 1.0} ), ), ), ( "ar_wgan_physical", lambda cfg: cfg["stage2_model"].__setitem__("decoder", "autoregressive"), ), ( "ar_flow_onehot", lambda cfg: ( cfg["stage2_model"].__setitem__("decoder", "autoregressive"), cfg["stage2_model"].__setitem__("generator", "flow"), cfg["stage2_model"].__setitem__( "particle_type", {"target": "onehot", "lambda": 1.0} ), ), ), ( "ar_flow_embedding", lambda cfg: ( cfg["conditioning"]["particle"].__setitem__("type", "embedding"), cfg["conditioning"]["material"].__setitem__("type", "embedding"), cfg["stage2_model"].__setitem__("decoder", "autoregressive"), cfg["stage2_model"].__setitem__("generator", "flow"), cfg["stage2_model"].__setitem__( "particle_type", {"target": "embedding", "lambda": 1.0} ), ), ), ( "ar_stage2_only", lambda cfg: ( cfg["stage1_model"].__setitem__("active", False), cfg["stage2_model"].__setitem__("decoder", "autoregressive"), ), ), ( "ar_mixed_stage1_wgan_stage2_flow_onehot", lambda cfg: ( cfg["stage1_model"].__setitem__("generator", "wgan"), cfg["stage2_model"].__setitem__("generator", "flow"), cfg["stage2_model"].__setitem__("decoder", "autoregressive"), cfg["stage2_model"].__setitem__( "particle_type", {"target": "onehot", "lambda": 1.0} ), ), ), ], ) 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"), ) # --- AR trainer wiring (v0.3.0 step 5) -------------------------------------- def test_build_stage_trainers_rejects_scheduled_teacher_forcing(): cfg = _base_cfg() cfg["stage2_model"]["decoder"] = "autoregressive" cfg["stage2_model"]["autoregressive"] = { "history": "markov", "teacher_forcing": "scheduled", } model_config = _model_config(cfg) models = build_models(model_config) critics = build_critics(model_config) with pytest.raises(NotImplementedError): _build_stage_trainers( cfg, models, critics, torch.device("cpu"), total_train_batches=4 ) def test_ar_wgan_onehot_grad_norm_instrumentation_populates_metrics(): """§11.4 differentiability validation-obligation instrumentation: the trunk-gradient-norm-by-slice columns must appear and actually fire for generator='wgan' + particle_type.target='onehot' under decoder= 'autoregressive' (added at v0.3.0 step 5 per the design doc's instruction to accrue evidence during the architecture comparison).""" cfg = _base_cfg() cfg["stage2_model"]["decoder"] = "autoregressive" cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0} with tempfile.TemporaryDirectory() as tmp: out_dir = Path(tmp) / "run" _run_train(cfg, out_dir) with open(out_dir / "metrics.csv", newline="") as f: rows = list(csv.DictReader(f)) assert "stage2_train_grad_norm_type_slice" in rows[0] assert "stage2_train_grad_norm_cont_slice" in rows[0] assert any(float(r["stage2_train_grad_norm_type_slice"]) > 0 for r in rows) assert any(float(r["stage2_train_grad_norm_cont_slice"]) > 0 for r in rows) def test_wgan_onehot_one_shot_also_gets_grad_norm_instrumentation(): """The instrumentation is decoder-agnostic — one_shot + wgan + onehot must populate the same columns.""" cfg = _base_cfg() cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0} with tempfile.TemporaryDirectory() as tmp: out_dir = Path(tmp) / "run" _run_train(cfg, out_dir) with open(out_dir / "metrics.csv", newline="") as f: rows = list(csv.DictReader(f)) assert any(float(r["stage2_train_grad_norm_type_slice"]) > 0 for r in rows) assert any(float(r["stage2_train_grad_norm_cont_slice"]) > 0 for r in rows) def test_wgan_physical_omits_grad_norm_slice_columns(): cfg = _base_cfg() # default stage2_model has no particle_type -> "physical" 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 "stage2_train_grad_norm_type_slice" not in header assert "stage2_train_grad_norm_cont_slice" not in header