"""Tests for giant/training/.""" import copy import csv import math import tempfile from pathlib import Path from unittest.mock import MagicMock, patch import numpy as np import pytest import torch import torch.nn.functional as F from giant.config import ParticleTypeConfig from giant.constants import ( COND_DIM, CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM, ) from giant.checkpoint_io import load_for_inference from giant.data.dataset import StepBatch from giant.data.transforms import Normalizer from giant.model.network import Stage2Autoregressive, build_critics, build_models from giant.sample import sample_stage1 as trainers_sample_stage1 from giant.training import ( FlowDDPMStageTrainer, StageSpec, WGANStageTrainer, build_checkpoint, build_stage_trainers, init_stages_from_checkpoints, train, ) from giant.training.metrics import _wandb_run_config from giant.training.trainers import _type_class_weight_vector from giant.training.stage2_inputs import ( _ar_has_prev, _assemble_stage2_ar_inputs, _assemble_stage2_ar_target, _assemble_stage2_real, _gumbel_tau, _relax_onehot_type_slice, _remaining_energy_fraction, _shift_prev, _stage2_tf_prob, _stick_fraction, _stop_target_and_mask, _type_repr, ) 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) --------- 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]] def test_stop_target_and_mask_hand_computed(): # k_max=5; n_sec=0 (no real secondaries, stop slot is 0), n_sec=2 # (stop slot is 2), n_sec=5 (== k_max: no in-range stop slot at all). n_sec = torch.tensor([0, 2, 5]) target, mask = _stop_target_and_mask(n_sec, 5, torch.device("cpu")) assert target.tolist() == [ [1, 0, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 0, 0], ] assert mask.tolist() == [ [True, False, False, False, False], [True, True, True, False, False], [True, True, True, True, True], ] # --- _stage2_tf_prob (v0.3.0 step 7) ----------- def test_stage2_tf_prob_always_is_constant_one(): assert _stage2_tf_prob("always", 1.0, 0.0, 0, 10) == 1.0 assert _stage2_tf_prob("always", 1.0, 0.0, 9, 10) == 1.0 def test_stage2_tf_prob_never_is_constant_zero(): assert _stage2_tf_prob("never", 1.0, 1.0, 0, 10) == 0.0 assert _stage2_tf_prob("never", 1.0, 1.0, 9, 10) == 0.0 def test_stage2_tf_prob_scheduled_interpolates_linearly(): assert _stage2_tf_prob("scheduled", 1.0, 0.0, 0, 11) == 1.0 assert abs(_stage2_tf_prob("scheduled", 1.0, 0.0, 5, 11) - 0.5) < 1e-9 assert _stage2_tf_prob("scheduled", 1.0, 0.0, 10, 11) == 0.0 def test_stage2_tf_prob_scheduled_clamps_beyond_total_epochs(): end = _stage2_tf_prob("scheduled", 1.0, 0.0, 10, 11) beyond = _stage2_tf_prob("scheduled", 1.0, 0.0, 50, 11) assert beyond == end def test_stage2_tf_prob_scheduled_handles_single_epoch(): # total_epochs=1 is guarded to a denominator of 1 internally (like # _gumbel_tau's total_steps=0 guard) — epoch=0 gives zero progress. assert _stage2_tf_prob("scheduled", 1.0, 0.0, 0, 1) == 1.0 @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, ParticleTypeConfig(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 = ParticleTypeConfig(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, ParticleTypeConfig(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}, # Explicit, not relying on the fallback default (which is # "onehot", matching DEFAULT_CONFIG — see issues.md Issue 1): # the "physical"-labelled cases below (and this fixture's own # comment history) intend this as the base "physical" case, # with "*_onehot"/"*_embedding" cases opting in explicitly. "particle_type": {"target": "physical", "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, "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(StepBatch(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, normalizer_dict=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=normalizer_dict or {"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), 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 spec = StageSpec( name="stage1", is_stage2=False, generator="wgan", n_critic=1000, # never a generator step in this test ema_decay=0.0, steps_per_epoch=4, ) trainer = WGANStageTrainer(spec, models["stage1"], critics["stage1"], 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(): spec = StageSpec(name="stage2", is_stage2=True, generator="ddpm", ddpm_n_steps=50, ema_decay=0.0) with pytest.raises(NotImplementedError): FlowDDPMStageTrainer(spec, torch.nn.Linear(1, 1), torch.device("cpu")) # --- gitea #44: class-balanced secondary particle-type loss ----------------- def test_type_class_weight_vector_none_scheme_returns_none(): assert _type_class_weight_vector({0: 100, 1: 5}, n_classes=2, scheme="none") is None def test_type_class_weight_vector_raises_without_counts(): with pytest.raises(ValueError, match="class_counts"): _type_class_weight_vector({}, n_classes=4, scheme="inverse_freq") def test_type_class_weight_vector_inverse_freq_favors_rare_class_and_has_mean_one(): weights = _type_class_weight_vector({0: 1000, 1: 10, 2: 1, 3: 1}, n_classes=4, scheme="inverse_freq") assert weights is not None assert len(weights) == 4 assert weights[1] > weights[0] # rarer class -> larger weight assert math.isclose(sum(weights) / len(weights), 1.0, rel_tol=1e-9) def test_type_class_weight_vector_missing_index_clamps_to_count_one(): # n_classes=3 but only index 0 was ever observed (e.g. a tiny dataset) — # indices 1/2 must not divide by zero. weights = _type_class_weight_vector({0: 10}, n_classes=3, scheme="inverse_freq") assert weights is not None assert all(math.isfinite(w) for w in weights) def _onehot_flow_stage2_setup(): """A built stage-2 model + a batch, under target='onehot' + generator='flow' (mirrors the 'stage2_onehot_target_flow' case in test_train_end_to_end).""" cfg = _base_cfg() cfg["stage2_model"]["generator"] = "flow" cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0} model_config = _model_config(cfg) model = build_models(model_config)["stage2"] assert model is not None batch = _fake_batches(1, 8)[0] device = torch.device("cpu") cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1 # Mostly class 0 (common), a few slot 1's set to class 1 (rare) — # PARTICLE_CFG's emb_dim=8, n_classes=0 (inherit) -> 8 type classes. sec_type_idx = torch.zeros(8, K_MAX, dtype=torch.long) sec_type_idx[:, :2] = 1 sec_mask = torch.ones(8, K_MAX, dtype=torch.bool) return model, cond_cont, cond_cat, x1_s1, sec_type_idx, sec_mask, device def test_flow_ddpm_trainer_type_loss_none_leaves_weight_unset(): model, *_ = _onehot_flow_stage2_setup() spec = StageSpec( name="stage2", is_stage2=True, generator="flow", particle_type=ParticleTypeConfig(target="onehot", class_weighting="none"), particle_type_n_classes=8, ema_decay=0.0, ) trainer = FlowDDPMStageTrainer(spec, model, torch.device("cpu")) assert trainer.type_class_weights is None def test_flow_ddpm_trainer_type_loss_matches_manual_weighted_cross_entropy(): model, cond_cont, cond_cat, x1_s1, sec_type_idx, sec_mask, device = _onehot_flow_stage2_setup() class_counts = {0: 1000, 1: 10, 2: 1, 3: 1, 4: 1, 5: 1, 6: 1, 7: 1} weights = _type_class_weight_vector(class_counts, n_classes=8, scheme="inverse_freq") spec = StageSpec( name="stage2", is_stage2=True, generator="flow", particle_type=ParticleTypeConfig(target="onehot", class_weighting="inverse_freq"), particle_type_n_classes=8, type_class_weights=weights, ema_decay=0.0, ) trainer = FlowDDPMStageTrainer(spec, model, device) assert trainer.type_class_weights is not None stage1_ctx = trainer._stage1_context(x1_s1, cond_cont, cond_cat, epoch=None) with torch.no_grad(): type_out = model.predict_type(cond_cont, cond_cat, stage1_ctx) weight_t = torch.tensor(weights) ce = F.cross_entropy(type_out.transpose(1, 2), sec_type_idx, weight=weight_t, reduction="none") expected = (ce * sec_mask.float()).sum() / sec_mask.float().sum().clamp(min=1) l_type, _ = trainer._type_loss(cond_cont, cond_cat, stage1_ctx, sec_type_idx, sec_mask, device) assert torch.allclose(l_type, expected, atol=1e-6) # Unweighted trainer, same model/batch — the two losses must differ # (the batch mixes the common and rare classes, so weighting changes the # per-slot contributions), confirming the weight is actually plumbed in. spec_none = StageSpec( name="stage2", is_stage2=True, generator="flow", particle_type=ParticleTypeConfig(target="onehot", class_weighting="none"), particle_type_n_classes=8, ema_decay=0.0, ) trainer_none = FlowDDPMStageTrainer(spec_none, model, device) with torch.no_grad(): l_type_none, _ = trainer_none._type_loss(cond_cont, cond_cat, stage1_ctx, sec_type_idx, sec_mask, device) assert not torch.allclose(l_type, l_type_none) def test_build_stage_trainers_threads_sec_type_class_counts_into_weights(): cfg = _base_cfg() cfg["stage2_model"]["generator"] = "flow" cfg["stage2_model"]["particle_type"] = { "target": "onehot", "lambda": 1.0, "class_weighting": "inverse_freq", } model_config = _model_config(cfg) models = build_models(model_config) critics = build_critics(model_config) class_counts = {i: 100 for i in range(8)} class_counts[1] = 1 # one rare class trainers = build_stage_trainers( cfg, models, critics, torch.device("cpu"), total_train_batches=4, sec_type_class_counts=class_counts ) stage2_trainer = trainers["stage2"] assert isinstance(stage2_trainer, FlowDDPMStageTrainer) weights = stage2_trainer.type_class_weights assert weights is not None assert weights[1] > weights[0] def test_build_stage_trainers_no_class_counts_with_none_weighting_is_fine(): """The overwhelmingly common case (class_weighting = 'none', the default): build_stage_trainers must not require sec_type_class_counts at all.""" cfg = _base_cfg() cfg["stage2_model"]["generator"] = "flow" cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}) model_config = _model_config(cfg) models = build_models(model_config) critics = build_critics(model_config) trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4) stage2_trainer = trainers["stage2"] assert isinstance(stage2_trainer, FlowDDPMStageTrainer) assert stage2_trainer.type_class_weights is None # --- gitea #42: freeze / init_from ------------------------------------------- def _state_dict_clone(module): return {k: v.clone() for k, v in module.state_dict().items()} def _assert_state_dicts_equal(before, after, label): for key, value in before.items(): assert torch.equal(value, after[key]), f"{label}: {key} changed while frozen" def test_frozen_flow_stage_trainer_step_does_not_update_model_or_ema(): cfg = _base_cfg() model_config = _model_config(cfg) model = build_models(model_config)["stage1"] assert model is not None spec = StageSpec(name="stage1", is_stage2=False, generator="flow", freeze=True, ema_decay=0.999, steps_per_epoch=4) trainer = FlowDDPMStageTrainer(spec, model, torch.device("cpu")) assert trainer.ema_model is not None model_before = _state_dict_clone(trainer.model) ema_before = _state_dict_clone(trainer.ema_model) for batch in _fake_batches(4, 8): trainer.step(batch, torch.device("cpu"), global_step=1) _assert_state_dicts_equal(model_before, trainer.model.state_dict(), "frozen flow model") _assert_state_dicts_equal(ema_before, trainer.ema_model.state_dict(), "frozen flow ema") def test_frozen_wgan_stage_trainer_step_does_not_update_generator_or_critic(): 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 spec = StageSpec( name="stage1", is_stage2=False, generator="wgan", freeze=True, n_critic=1, # a generator step every batch, so a bug would surface immediately ema_decay=0.999, steps_per_epoch=4, ) trainer = WGANStageTrainer(spec, models["stage1"], critics["stage1"], torch.device("cpu")) assert trainer.ema_model is not None model_before = _state_dict_clone(trainer.model) critic_before = _state_dict_clone(trainer.critic) ema_before = _state_dict_clone(trainer.ema_model) for global_step, batch in enumerate(_fake_batches(4, 8)): trainer.step(batch, torch.device("cpu"), global_step=global_step) _assert_state_dicts_equal(model_before, trainer.model.state_dict(), "frozen wgan generator") _assert_state_dicts_equal(critic_before, trainer.critic.state_dict(), "frozen wgan critic") _assert_state_dicts_equal(ema_before, trainer.ema_model.state_dict(), "frozen wgan ema") @pytest.mark.parametrize("stage1_generator", ["flow", "wgan"]) def test_train_end_to_end_frozen_stage1_unchanged_while_stage2_trains(stage1_generator): cfg = _base_cfg() cfg["stage1_model"]["generator"] = stage1_generator cfg["stage1_model"]["freeze"] = True model_config = _model_config(cfg) models = build_models(model_config) critics = build_critics(model_config) assert models["stage1"] is not None and models["stage2"] is not None stage1_before = _state_dict_clone(models["stage1"]) stage2_before = _state_dict_clone(models["stage2"]) with tempfile.TemporaryDirectory() as tmp: train( cfg=cfg, models=models, critics=critics, train_loader=_fake_batches(4, cfg["train"]["batch_size"]), val_loader=_fake_batches(2, cfg["train"]["batch_size"], seed=1), device=torch.device("cpu"), out_dir=Path(tmp) / "run", 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, ) _assert_state_dicts_equal(stage1_before, models["stage1"].state_dict(), "frozen stage1") stage2_after = models["stage2"].state_dict() assert any(not torch.equal(v, stage2_after[k]) for k, v in stage2_before.items()), ( "unfrozen stage2 should have trained" ) def test_init_stages_from_checkpoints_loads_matching_stage_and_ema_weights(tmp_path): cfg = _base_cfg() model_config = _model_config(cfg) source_models = build_models(model_config) source_critics = build_critics(model_config) source_trainers = build_stage_trainers(cfg, source_models, source_critics, torch.device("cpu"), 4) source_stage1_ema = source_trainers["stage1"].ema_model assert source_stage1_ema is not None # Diverge the source's EMA from its raw weights so a same-vs-different # check below actually distinguishes the two copy paths. for p in source_stage1_ema.parameters(): p.data.add_(1.0) ckpt_path = tmp_path / "source.pt" ckpt = build_checkpoint(source_trainers, epoch=1, global_step=1, best_val_loss=0.0, extras={}) torch.save(ckpt, ckpt_path) cfg2 = copy.deepcopy(cfg) cfg2["stage1_model"]["init_from"] = str(ckpt_path) dest_models = build_models(_model_config(cfg2)) dest_critics = build_critics(_model_config(cfg2)) dest_trainers = build_stage_trainers(cfg2, dest_models, dest_critics, torch.device("cpu"), 4) loaded = init_stages_from_checkpoints(dest_trainers) assert len(loaded) == 1 and "stage1" in loaded[0] dest_stage1_ema = dest_trainers["stage1"].ema_model assert dest_stage1_ema is not None _assert_state_dicts_equal( source_trainers["stage1"].model.state_dict(), dest_trainers["stage1"].model.state_dict(), "init_from raw" ) _assert_state_dicts_equal( source_stage1_ema.state_dict(), dest_stage1_ema.state_dict(), "init_from ema", ) # stage2 has no init_from set -- untouched fresh init, not the source's. stage2_matches_source = all( torch.equal(v, dest_trainers["stage2"].model.state_dict()[k]) for k, v in source_trainers["stage2"].model.state_dict().items() ) assert not stage2_matches_source def test_run_train_job_stage1_init_from_freeze_produces_rollout_capable_checkpoint(tmp_path): """The exact scenario gitea #42 exists for: retrain stage 2 alone against a fixed, known-good stage 1, and still get a checkpoint giant rollout can load (checkpoint_io.load_for_inference with require_stage2=True).""" normalizer_dict = { "cond": Normalizer().fit(np.zeros((1, COND_DIM), dtype=np.float32)).to_dict(), "target": Normalizer().fit(np.zeros((1, X_DIM), dtype=np.float32)).to_dict(), "sec_phys": Normalizer().fit(np.zeros((1, 2), dtype=np.float32)).to_dict(), } cfg = _base_cfg() source_out = tmp_path / "source" _run_train(cfg, source_out, normalizer_dict=normalizer_dict) source_ckpt = torch.load(source_out / "best.pt", weights_only=False) cfg2 = copy.deepcopy(cfg) cfg2["stage1_model"]["init_from"] = str(source_out / "best.pt") cfg2["stage1_model"]["freeze"] = True retrain_out = tmp_path / "retrain" _run_train(cfg2, retrain_out, normalizer_dict=normalizer_dict) ctx = load_for_inference(retrain_out / "best.pt", torch.device("cpu"), "rollout", require_stage2=True) assert ctx.stage1 is not None and ctx.stage2 is not None retrain_ckpt = torch.load(retrain_out / "best.pt", weights_only=False) for key, value in source_ckpt["model"].items(): assert torch.equal(value, retrain_ckpt["model"][key]), f"frozen stage1 {key} drifted across the retrain" def test_stage_spec_from_config_omitted_decoder_and_particle_type_match_default_config(): """Regression for issues.md Issue 1: StageSpec.from_config's own fallback defaults for stage2_model.decoder/particle_type must equal DEFAULT_CONFIG's ("autoregressive" / "onehot"), not the old, now-wrong ("one_shot" / "physical") literals a .get(key, default) call used to supply when a hand-built cfg omitted these keys.""" cfg = _base_cfg() del cfg["stage2_model"]["decoder"] del cfg["stage2_model"]["particle_type"] spec = StageSpec.from_config(cfg, "stage2", is_stage2=True, steps_per_epoch=1) assert spec.decoder == "autoregressive" assert spec.particle_type.target == "onehot" def _routed_stage1_trainer(lambda_balance, lambda_proc, lambda_entropy): cfg = _base_cfg() cfg["stage1_model"]["router"] = { "enabled": True, "type": "energy", "n_experts": 3, "temperature": 0.5, "learn_centers": True, "lambda_balance": lambda_balance, "lambda_proc": lambda_proc, "lambda_entropy": lambda_entropy, } model_config = _model_config(cfg) models = build_models(model_config) critics = build_critics(model_config) trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4) return trainers["stage1"] def test_router_aux_losses_skipped_when_lambda_zero_but_run_when_positive(): """Gitea #31: FlowDDPMStageTrainer._compute must not call router.balance_loss/classify_loss/entropy_loss when the corresponding lambda is 0 (the default) -- those calls do their own router.gate(...) forward pass that is wasted once the term is masked out of the total loss anyway. Checked both ways: zero lambdas must skip all three calls, positive lambdas must still make them (the guard must not accidentally suppress the real path).""" batch = _fake_batches(1, 4)[0] device = torch.device("cpu") trainer_zero = _routed_stage1_trainer(0.0, 0.0, 0.0) router_zero = trainer_zero.router router_zero.balance_loss = MagicMock(wraps=router_zero.balance_loss) router_zero.classify_loss = MagicMock(wraps=router_zero.classify_loss) router_zero.entropy_loss = MagicMock(wraps=router_zero.entropy_loss) stats_zero = trainer_zero.step(batch, device, global_step=1) assert router_zero.balance_loss.call_count == 0 assert router_zero.classify_loss.call_count == 0 assert router_zero.entropy_loss.call_count == 0 assert stats_zero["loss_balance"] == 0.0 assert stats_zero["loss_proc"] == 0.0 assert stats_zero["loss_entropy"] == 0.0 trainer_pos = _routed_stage1_trainer(0.1, 0.1, 0.01) router_pos = trainer_pos.router router_pos.balance_loss = MagicMock(wraps=router_pos.balance_loss) router_pos.classify_loss = MagicMock(wraps=router_pos.classify_loss) router_pos.entropy_loss = MagicMock(wraps=router_pos.entropy_loss) trainer_pos.step(batch, device, global_step=1) assert router_pos.balance_loss.call_count == 1 assert router_pos.classify_loss.call_count == 1 assert router_pos.entropy_loss.call_count == 1 # --- AR trainer wiring (v0.3.0 step 5) -------------------------------------- @pytest.mark.parametrize("teacher_forcing", ["always", "scheduled", "never"]) @pytest.mark.parametrize("history", ["markov", "attention"]) @pytest.mark.parametrize("stage2_generator", ["wgan", "flow"]) def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(teacher_forcing, history, stage2_generator): """v0.3.0 step 7: history='attention' and teacher_forcing in {'scheduled', 'never'} must actually train — a stage-2 AR trainer.step() must run and produce a finite loss, for every {history} x {teacher_forcing} x {generator} combination.""" cfg = _base_cfg() cfg["stage2_model"]["decoder"] = "autoregressive" cfg["stage2_model"]["generator"] = stage2_generator cfg["stage2_model"]["autoregressive"] = { "history": history, "teacher_forcing": teacher_forcing, "tf_p_start": 1.0, "tf_p_end": 0.0, "attn_n_heads": 2, "attn_n_layers": 1, } model_config = _model_config(cfg) models = build_models(model_config) critics = build_critics(model_config) trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4) trainer = trainers["stage2"] batch = _fake_batches(1, 4)[0] stats = trainer.step(batch, torch.device("cpu"), global_step=1) loss_key = "g_loss" if stage2_generator == "wgan" else "loss" assert math.isfinite(stats[loss_key]) @pytest.mark.parametrize("stage2_generator", ["wgan", "flow"]) def test_train_end_to_end_ar_attention_history_scheduled_teacher_forcing( stage2_generator, ): """Full `train()` run (not just one `trainer.step()` call) with history='attention' AND teacher_forcing='scheduled' together — the combination v0.3.0 step 7 exists to land — must complete and write a checkpoint + metrics.csv with finite losses throughout.""" cfg = _base_cfg() cfg["stage2_model"]["decoder"] = "autoregressive" cfg["stage2_model"]["generator"] = stage2_generator cfg["stage2_model"]["autoregressive"] = { "history": "attention", "teacher_forcing": "scheduled", "tf_p_start": 1.0, "tf_p_end": 0.0, "attn_n_heads": 2, "attn_n_layers": 1, } with tempfile.TemporaryDirectory() as tmp: out_dir = Path(tmp) / "run" _run_train(cfg, out_dir) assert (out_dir / "last.pt").exists() with open(out_dir / "metrics.csv", newline="") as f: rows = list(csv.DictReader(f)) assert len(rows) == cfg["train"]["epochs"] loss_col = "stage2/train/g_loss" if stage2_generator == "wgan" else "stage2/train/loss" assert all(math.isfinite(float(r[loss_col])) for r in rows) def test_ar_wgan_onehot_grad_norm_instrumentation_populates_metrics(): """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 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) # --- n_sec.mode = "stop_token" (gitea #40) ---------------------------------- def _stop_token_cfg(): cfg = _base_cfg() cfg["stage2_model"]["decoder"] = "autoregressive" cfg["stage2_model"]["n_sec"] = {"mode": "stop_token", "lambda": 0.1} return cfg @pytest.mark.parametrize("stage2_generator", ["wgan", "flow"]) def test_build_stage_trainers_stop_token_step_runs(stage2_generator): """A stop_token AR stage-2 trainer.step() must run and emit a finite loss_stop for both non-adversarial (flow) and WGAN generators — the two trainer subclasses wire the stop head's BCE term in independently.""" cfg = _stop_token_cfg() cfg["stage2_model"]["generator"] = stage2_generator model_config = _model_config(cfg) models = build_models(model_config) critics = build_critics(model_config) trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4) trainer = trainers["stage2"] batch = _fake_batches(1, 4)[0] stats = trainer.step(batch, torch.device("cpu"), global_step=1) assert math.isfinite(stats["loss_stop"]) assert math.isfinite(stats["stop_acc"]) def test_stop_token_model_has_stop_head_not_n_sec_head(): cfg = _stop_token_cfg() model_config = _model_config(cfg) models = build_models(model_config) stage2 = models["stage2"] assert isinstance(stage2, Stage2Autoregressive) assert stage2.n_sec_head is None assert stage2.stop_head is not None def test_head_mode_model_has_n_sec_head_not_stop_head(): """Sanity check on the other side of the gate — the default 'head' mode must be unaffected by the stop_head plumbing.""" cfg = _base_cfg() cfg["stage2_model"]["decoder"] = "autoregressive" model_config = _model_config(cfg) models = build_models(model_config) stage2 = models["stage2"] assert isinstance(stage2, Stage2Autoregressive) assert stage2.n_sec_head is not None assert stage2.stop_head is None def test_train_end_to_end_stop_token(): """Full train() run with n_sec.mode='stop_token' must complete and write a checkpoint + metrics.csv with finite losses throughout.""" cfg = _stop_token_cfg() with tempfile.TemporaryDirectory() as tmp: out_dir = Path(tmp) / "run" _run_train(cfg, out_dir) assert (out_dir / "last.pt").exists() with open(out_dir / "metrics.csv", newline="") as f: rows = list(csv.DictReader(f)) assert len(rows) == cfg["train"]["epochs"] assert all(math.isfinite(float(r["stage2/train/loss_stop"])) for r in rows) def test_wgan_physical_omits_grad_norm_slice_columns(): cfg = _base_cfg() # _base_cfg's stage2_model.particle_type.target is "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 # --- stage2_model.stage1_context = "sampled" (gitea #41) -------------------- def _sampled_ctx_cfg(ema_decay=0.999): cfg = _base_cfg() cfg["stage1_model"]["generator"] = "flow" cfg["stage2_model"]["generator"] = "flow" cfg["stage2_model"]["stage1_context"] = "sampled" cfg["stage2_model"]["ctx_p_start"] = 0.0 cfg["stage2_model"]["ctx_p_end"] = 0.0 cfg["train"]["ema_decay"] = ema_decay return cfg def _build_sampled_trainers(cfg): model_config = _model_config(cfg) models = build_models(model_config) critics = build_critics(model_config) return build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4) def test_build_stage_trainers_attaches_stage1_only_under_sampled(): trainers = _build_sampled_trainers(_sampled_ctx_cfg()) assert trainers["stage2"].stage1_source is trainers["stage1"] assert trainers["stage1"].stage1_source is None def test_build_stage_trainers_leaves_stage1_source_none_under_truth(): """Regression guard for the old silent no-op: 'truth' (the default) must never attach a stage1_source, so _stage1_context short-circuits without ever calling sample_stage1.""" cfg = _base_cfg() trainers = _build_sampled_trainers(cfg) assert trainers["stage2"].stage1_source is None def test_stage1_context_sampled_calls_sample_stage1_and_differs_from_truth(): cfg = _sampled_ctx_cfg() trainers = _build_sampled_trainers(cfg) stage1, stage2 = trainers["stage1"], trainers["stage2"] batch = _fake_batches(1, 4)[0] cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1 with patch("giant.training.trainers.sample_stage1", wraps=trainers_sample_stage1) as spy: ctx = stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=0) assert spy.call_count == 1 assert spy.call_args.args[0] is stage1.sampling_model() assert not torch.equal(ctx, x1_s1) def test_stage1_context_truth_default_never_calls_sample_stage1(): cfg = _base_cfg() trainers = _build_sampled_trainers(cfg) stage2 = trainers["stage2"] batch = _fake_batches(1, 4)[0] cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1 with patch("giant.training.trainers.sample_stage1", wraps=trainers_sample_stage1) as spy: ctx = stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=0) assert spy.call_count == 0 assert torch.equal(ctx, x1_s1) def test_stage1_context_val_epoch_none_uses_ground_truth_even_under_sampled(): cfg = _sampled_ctx_cfg() trainers = _build_sampled_trainers(cfg) stage2 = trainers["stage2"] batch = _fake_batches(1, 4)[0] cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1 with patch("giant.training.trainers.sample_stage1", wraps=trainers_sample_stage1) as spy: ctx = stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=None) assert spy.call_count == 0 assert torch.equal(ctx, x1_s1) def test_stage1_context_sampled_preserves_stage1_training_mode(): """Every sampler in giant/sample.py flips its model to .eval() as a side effect with no restore of its own (see sample_flow). Sampling from the RAW stage-1 model (ema_decay=0, so sampling_model() returns self.model, the same weights the stage-1 trainer is actively training on) must not silently leave it in eval mode for the rest of the epoch's stage-1 updates.""" cfg = _sampled_ctx_cfg(ema_decay=0.0) trainers = _build_sampled_trainers(cfg) stage1, stage2 = trainers["stage1"], trainers["stage2"] stage1.train_mode() assert stage1.model.training batch = _fake_batches(1, 4)[0] cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1 stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=0) assert stage1.model.training @pytest.mark.parametrize("stage2_generator", ["wgan", "flow"]) def test_build_stage_trainers_sampled_step_runs(stage2_generator): """Both trainer subclasses' call sites (FlowDDPMStageTrainer._compute, WGANStageTrainer.step) must run end to end under 'sampled' and produce a finite loss.""" cfg = _sampled_ctx_cfg() cfg["stage2_model"]["generator"] = stage2_generator trainers = _build_sampled_trainers(cfg) trainer = trainers["stage2"] batch = _fake_batches(1, 4)[0] stats = trainer.step(batch, torch.device("cpu"), global_step=1) loss_key = "g_loss" if stage2_generator == "wgan" else "loss" assert math.isfinite(stats[loss_key]) def test_train_end_to_end_stage1_context_sampled(): """Full train() run with stage1_context='sampled' must complete and write a checkpoint + metrics.csv with finite losses throughout.""" cfg = _sampled_ctx_cfg() with tempfile.TemporaryDirectory() as tmp: out_dir = Path(tmp) / "run" _run_train(cfg, out_dir) assert (out_dir / "last.pt").exists() with open(out_dir / "metrics.csv", newline="") as f: rows = list(csv.DictReader(f)) assert len(rows) == cfg["train"]["epochs"] assert all(math.isfinite(float(r["stage2/train/loss"])) for r in rows)