from datetime import datetime from pathlib import Path import pytest from giant import config as gconfig _CONFIGS_DIR = Path(__file__).resolve().parents[1] / "configs" def _write_toml(path, git_hash=None, extra=""): meta = f'\n[meta]\ngit_hash = "{git_hash}"\n' if git_hash is not None else "" path.write_text(extra + meta) # --------------------------------------------------------------------------- # Conditioning enum # --------------------------------------------------------------------------- def test_conditioning_enum_has_onehot(): assert gconfig.Conditioning.onehot == "onehot" assert {c.value for c in gconfig.Conditioning} == { "physical", "embedding", "onehot", } # --------------------------------------------------------------------------- # Config dataclasses (issues.md Issue 1) # --------------------------------------------------------------------------- def test_giant_config_to_dict_matches_default_config(): """DEFAULT_CONFIG is generated from GiantConfig().to_dict() (not hand-maintained), so the two cannot structurally drift apart — but this pins the *equality* too, catching e.g. a stray in-place mutation of DEFAULT_CONFIG added elsewhere after import.""" assert gconfig.GiantConfig().to_dict() == gconfig.DEFAULT_CONFIG @pytest.mark.parametrize( "cls", [ gconfig.ConditioningAxisConfig, gconfig.ConditioningConfig, gconfig.FlowConfig, gconfig.DdpmConfig, gconfig.Stage1WganConfig, gconfig.Stage2WganConfig, gconfig.RouterConfig, gconfig.Stage2RouterConfig, gconfig.TrunkConfig, gconfig.NSecConfig, gconfig.ParticleTypeConfig, gconfig.AutoregressiveConfig, gconfig.HeadConfig, gconfig.Stage1HeadsConfig, gconfig.Stage2HeadsConfig, gconfig.Stage1ModelConfig, gconfig.Stage2ModelConfig, gconfig.TrainConfig, gconfig.GiantConfig, ], ) def test_config_dataclass_from_dict_round_trips_through_to_dict(cls): assert cls.from_dict(cls().to_dict()) == cls() assert cls.from_dict(None) == cls() def test_stage2_model_config_defaults_match_documented_v030_intent(): """The two keys issues.md Issue 1 found drifted between DEFAULT_CONFIG and build_models/StageSpec.from_config's own .get(key, default) fallbacks — pinned directly against the dataclass that is now their shared single source of truth.""" spec = gconfig.Stage2ModelConfig() assert spec.decoder == "autoregressive" assert spec.particle_type.target == "onehot" def test_trunk_config_defaults_to_resmlp_for_both_stages(): """gitea #33: a v0.2-migrated / pre-existing config with no `trunk` key at all must reproduce today's behaviour exactly.""" assert gconfig.Stage1ModelConfig().trunk.type == "resmlp" assert gconfig.Stage2ModelConfig().trunk.type == "resmlp" assert gconfig.DEFAULT_CONFIG["stage1_model"]["trunk"]["type"] == "resmlp" assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["type"] == "resmlp" def test_trunk_config_defaults_block_conditioning_to_add_for_both_stages(): """gitea #34: a pre-existing config with no `block_conditioning` key must reproduce today's additive-bias behaviour exactly.""" assert gconfig.Stage1ModelConfig().trunk.block_conditioning == "add" assert gconfig.Stage2ModelConfig().trunk.block_conditioning == "add" assert gconfig.DEFAULT_CONFIG["stage1_model"]["trunk"]["block_conditioning"] == "add" assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["block_conditioning"] == "add" def test_heads_config_defaults_reproduce_pre_gitea_36_hardcoded_shape(): """gitea #36: a pre-existing config with no `heads` key must reproduce today's hardcoded `hidden_dim // 2`, one-hidden-layer architecture exactly.""" assert gconfig.Stage1ModelConfig().heads.n_sec.hidden_ratio == 0.5 assert gconfig.Stage1ModelConfig().heads.n_sec.depth == 2 assert gconfig.Stage2ModelConfig().heads.n_sec.hidden_ratio == 0.5 assert gconfig.Stage2ModelConfig().heads.n_sec.depth == 2 assert gconfig.Stage2ModelConfig().heads.type.hidden_ratio == 0.5 assert gconfig.Stage2ModelConfig().heads.type.depth == 2 assert gconfig.DEFAULT_CONFIG["stage1_model"]["heads"]["n_sec"] == {"hidden_ratio": 0.5, "depth": 2} assert gconfig.DEFAULT_CONFIG["stage2_model"]["heads"]["n_sec"] == {"hidden_ratio": 0.5, "depth": 2} assert gconfig.DEFAULT_CONFIG["stage2_model"]["heads"]["type"] == {"hidden_ratio": 0.5, "depth": 2} def test_particle_type_config_n_classes_defaults_to_zero_and_round_trips(): """gitea #29: n_classes=0 means "inherit conditioning.particle.emb_dim" — the default must stay 0 so an existing config.toml with no stage2_model.particle_type.n_classes key reproduces pre-#29 behavior.""" assert gconfig.ParticleTypeConfig().n_classes == 0 spec = gconfig.ParticleTypeConfig.from_dict({"n_classes": 32}) assert spec.n_classes == 32 assert spec.to_dict()["n_classes"] == 32 def test_router_config_extra_round_trips_composed_axis_keys(): d = {"enabled": True, "type": "composed", "axis0_type": "energy", "axis0_n_experts": 4} router = gconfig.RouterConfig.from_dict(d) assert router.enabled is True assert router.extra == {"axis0_type": "energy", "axis0_n_experts": 4} assert router.to_dict()["axis0_type"] == "energy" def test_stage2_router_config_tie_to_stage1_not_leaked_into_extra(): router = gconfig.Stage2RouterConfig.from_dict({"tie_to_stage1": True}) assert router.tie_to_stage1 is True assert "tie_to_stage1" not in router.extra def test_stage1_router_config_has_no_tie_to_stage1_key(): """Stage 1's router schema must not gain stage 2's tie_to_stage1 key — that would change every future run's saved config.toml shape.""" assert "tie_to_stage1" not in gconfig.RouterConfig().to_dict() def test_n_sec_config_owner_defaults_to_stage2(): n_sec = gconfig.NSecConfig() assert n_sec.owner == "stage2" def test_n_sec_config_owner_round_trips(): n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "owner": "stage1"}) assert n_sec.owner == "stage1" assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "stop_sampling": "greedy"} def test_n_sec_config_stop_sampling_defaults_to_greedy(): assert gconfig.NSecConfig().stop_sampling == "greedy" def test_n_sec_config_stop_sampling_round_trips(): n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "stop_sampling": "sample"}) assert n_sec.stop_sampling == "sample" assert n_sec.to_dict()["stop_sampling"] == "sample" # --------------------------------------------------------------------------- # _deep_merge # --------------------------------------------------------------------------- def test_deep_merge_leaf_override_keeps_untouched_siblings(): base = {"a": 1, "b": {"c": 2, "d": 3}} result = gconfig._deep_merge(base, {"b": {"c": 99}}) assert result == {"a": 1, "b": {"c": 99, "d": 3}} def test_deep_merge_recurses_at_multiple_levels(): base = { "stage1_model": { "hidden_dim": 256, "router": {"enabled": False, "type": "energy", "n_experts": 4}, } } result = gconfig._deep_merge(base, {"stage1_model": {"router": {"enabled": True}}}) assert result["stage1_model"]["hidden_dim"] == 256 assert result["stage1_model"]["router"] == { "enabled": True, "type": "energy", "n_experts": 4, } def test_deep_merge_does_not_mutate_base(): base = {"a": {"b": 1}} gconfig._deep_merge(base, {"a": {"b": 2}}) assert base == {"a": {"b": 1}} def test_deep_merge_non_dict_override_replaces_wholesale(): base = {"a": {"b": 1}} result = gconfig._deep_merge(base, {"a": 5}) assert result == {"a": 5} # --------------------------------------------------------------------------- # migrate_config # --------------------------------------------------------------------------- def test_migrate_config_already_v3_returned_unchanged(): cfg = {"meta": {"config_version": 3}, "stage1_model": {"generator": "flow"}} result = gconfig.migrate_config(cfg) assert result == cfg result["stage1_model"]["generator"] = "wgan" assert cfg["stage1_model"]["generator"] == "flow" # deep-copied, not aliased def test_migrate_config_empty_dict_still_injects_hardcoded_v02_facts(): # No [train]/[model] at all still counts as "v0.2" (config_version # absent) — the hardcoded architectural facts fire unconditionally. new = gconfig.migrate_config({}) assert "train" not in new assert new["conditioning"]["out_dim"] == 128 assert new["conditioning"]["particle"]["n_layers"] == 2 assert new["conditioning"]["material"]["n_layers"] == 2 assert new["stage1_model"]["active"] is True assert new["stage1_model"]["flow"]["time_dim"] == 64 assert new["stage1_model"]["ddpm"]["time_dim"] == 64 assert new["stage2_model"]["active"] is True assert new["stage2_model"]["flow"]["time_dim"] == 64 assert new["stage2_model"]["ddpm"]["time_dim"] == 64 assert new["stage2_model"]["context_dim"] == 64 assert new["stage2_model"]["decoder"] == "one_shot" assert new["stage2_model"]["particle_type"]["target"] == "physical" assert new["meta"] == {"config_version": 3} def test_migrate_config_mode_maps_to_both_stage_generators(): new = gconfig.migrate_config({"train": {"mode": "wgan"}}) assert new["stage1_model"]["generator"] == "wgan" assert new["stage2_model"]["generator"] == "wgan" def test_migrate_config_lambda_nsec_and_lambda_s2(): new = gconfig.migrate_config({"train": {"lambda_nsec": 0.2, "lambda_s2": 2.0}}) assert new["stage2_model"]["n_sec"]["lambda"] == 0.2 assert new["stage2_model"]["lambda"] == 2.0 def test_migrate_config_wgan_knobs_map_to_both_stages(): new = gconfig.migrate_config({"train": {"n_critic": 3, "gp_weight": 5.0, "critic_lr": 1e-4}}) for stage in ("stage1_model", "stage2_model"): assert new[stage]["wgan"]["n_critic"] == 3 assert new[stage]["wgan"]["gp_weight"] == 5.0 assert new[stage]["wgan"]["critic_lr"] == 1e-4 def test_migrate_config_model_hidden_dim_n_blocks_dropout_map_to_both_stages(): new = gconfig.migrate_config({"model": {"hidden_dim": 128, "n_blocks": 4, "dropout": 0.2}}) for stage in ("stage1_model", "stage2_model"): assert new[stage]["hidden_dim"] == 128 assert new[stage]["n_res_blocks"] == 4 assert new[stage]["dropout"] == 0.2 def test_migrate_config_emb_dim_and_conditioning_map_to_both_axes(): new = gconfig.migrate_config({"model": {"emb_dim": 32, "conditioning": "embedding"}}) for axis in ("particle", "material"): assert new["conditioning"][axis]["emb_dim"] == 32 assert new["conditioning"][axis]["type"] == "embedding" def test_migrate_config_noise_dim_maps_to_both_stages_wgan(): new = gconfig.migrate_config({"model": {"noise_dim": 128}}) assert new["stage1_model"]["wgan"]["noise_dim"] == 128 assert new["stage2_model"]["wgan"]["noise_dim"] == 128 def test_migrate_config_k_max_maps_to_stage2_only(): new = gconfig.migrate_config({"model": {"k_max": 20}}) assert new["stage2_model"]["k_max"] == 20 assert "k_max" not in new.get("stage1_model", {}) def test_migrate_config_router_copied_to_both_stages_with_tie_to_stage1_false(): new = gconfig.migrate_config( { "model": { "router": { "enabled": True, "type": "energy", "n_experts": 10, "temperature": 0.05, } } } ) assert new["stage1_model"]["router"] == { "enabled": True, "type": "energy", "n_experts": 10, "temperature": 0.05, } assert new["stage2_model"]["router"] == { "enabled": True, "type": "energy", "n_experts": 10, "temperature": 0.05, "tie_to_stage1": False, } def test_migrate_config_router_nonzero_expert_dims_raises(): cfg = {"model": {"router": {"enabled": True, "expert_hidden_dim": 128, "expert_n_blocks": 0}}} try: gconfig.migrate_config(cfg) assert False, "expected ValueError" except ValueError as e: assert "expert_hidden_dim" in str(e) def test_migrate_config_router_zero_expert_dims_dropped_silently(): new = gconfig.migrate_config( { "model": { "router": { "enabled": True, "type": "energy", "n_experts": 4, "expert_hidden_dim": 0, "expert_n_blocks": 0, } } } ) assert "expert_hidden_dim" not in new["stage1_model"]["router"] assert "expert_n_blocks" not in new["stage1_model"]["router"] def test_migrate_config_train_passthrough_is_exact(): new = gconfig.migrate_config({"train": {"epochs": 7, "batch_size": 999, "seed": 3}}) assert new["train"] == {"epochs": 7, "batch_size": 999, "seed": 3} def test_migrate_config_preserves_meta_git_hash(): new = gconfig.migrate_config({"meta": {"git_hash": "abc123"}}) assert new["meta"] == {"git_hash": "abc123", "config_version": 3} def test_migrate_config_real_router_fixture_raises_on_nonzero_expert_dims(): cfg = gconfig.load_toml(_CONFIGS_DIR / "router_energy_n10_embedding.toml") try: gconfig.migrate_config(cfg) assert False, "expected ValueError" except ValueError as e: assert "expert_hidden_dim" in str(e) def test_migrate_config_real_wgan_fixture(): cfg = gconfig.load_toml(_CONFIGS_DIR / "wgan_h128_b4_physical.toml") new = gconfig.migrate_config(cfg) assert new["stage1_model"]["generator"] == "wgan" assert new["stage2_model"]["generator"] == "wgan" for stage in ("stage1_model", "stage2_model"): assert new[stage]["hidden_dim"] == 128 assert new[stage]["n_res_blocks"] == 4 assert new[stage]["dropout"] == 0.0 assert new["conditioning"]["particle"]["type"] == "physical" assert new["conditioning"]["material"]["type"] == "physical" assert new["train"]["epochs"] == 30 assert new["train"]["warmup_epochs"] == 3 # --------------------------------------------------------------------------- # merge_cli_overrides # --------------------------------------------------------------------------- def test_merge_cli_overrides_defaults_only_matches_default_config(): cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {}) assert cfg == gconfig.DEFAULT_CONFIG assert cfg is not gconfig.DEFAULT_CONFIG def test_merge_cli_overrides_nested_override_keeps_siblings(): cfg = gconfig.merge_cli_overrides( gconfig.DEFAULT_CONFIG, None, {"stage1_model": {"router": {"enabled": True}}}, ) assert cfg["stage1_model"]["router"]["enabled"] is True assert cfg["stage1_model"]["router"]["type"] == "energy" # default preserved assert cfg["stage1_model"]["hidden_dim"] == 256 # untouched sibling section def test_merge_cli_overrides_file_then_explicit_override_precedence(tmp_path, monkeypatch): monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123") path = tmp_path / "config.toml" _write_toml( path, git_hash="abc123", extra="[train]\nepochs = 5\n\n[model]\nhidden_dim = 64\n", ) cfg = gconfig.merge_cli_overrides( gconfig.DEFAULT_CONFIG, path, {"stage1_model": {"hidden_dim": 128}}, ) assert cfg["train"]["epochs"] == 5 # from file assert cfg["stage1_model"]["hidden_dim"] == 128 # explicit override wins over file assert cfg["stage2_model"]["hidden_dim"] == 64 # migrated from file, not overridden def test_merge_cli_overrides_migrates_v2_file_transparently(tmp_path, monkeypatch): monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123") path = tmp_path / "config.toml" _write_toml( path, git_hash="abc123", extra='[train]\nmode = "wgan"\n\n[model]\nconditioning = "embedding"\n', ) cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) assert cfg["stage1_model"]["generator"] == "wgan" assert cfg["stage2_model"]["generator"] == "wgan" assert cfg["conditioning"]["particle"]["type"] == "embedding" # hardcoded v0.2 fact still applied even though it's not a CLI-settable key assert cfg["conditioning"]["particle"]["n_layers"] == 2 def test_merge_cli_overrides_warns_on_git_hash_mismatch(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "current999") path = tmp_path / "config.toml" _write_toml(path, git_hash="old111", extra="[train]\nepochs = 5\n") gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) captured = capsys.readouterr() assert "warning" in captured.err assert "old111" in captured.err assert "current999" in captured.err def test_merge_cli_overrides_no_warning_on_matching_git_hash(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "same123") path = tmp_path / "config.toml" _write_toml(path, git_hash="same123", extra="[train]\nepochs = 5\n") gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) assert capsys.readouterr().err == "" def test_merge_cli_overrides_no_warning_when_git_hash_unknown(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "unknown") path = tmp_path / "config.toml" _write_toml(path, git_hash="abc123", extra="[train]\nepochs = 5\n") gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) assert capsys.readouterr().err == "" def test_merge_cli_overrides_no_warning_when_meta_section_absent(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "current999") path = tmp_path / "config.toml" path.write_text("[train]\nepochs = 5\n") gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) assert capsys.readouterr().err == "" def test_merge_cli_overrides_real_default_toml_fixture(monkeypatch): monkeypatch.setattr(gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243") cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / "default.toml", {}) assert cfg["stage1_model"]["generator"] == "flow" assert cfg["stage1_model"]["hidden_dim"] == 256 assert cfg["stage2_model"]["hidden_dim"] == 256 assert cfg["conditioning"]["particle"]["emb_dim"] == 16 assert cfg["conditioning"]["particle"]["n_layers"] == 2 # migrated hardcoded fact assert cfg["train"]["epochs"] == 100 assert cfg["stage2_model"]["decoder"] == "one_shot" # --------------------------------------------------------------------------- # save_config # --------------------------------------------------------------------------- def test_save_config_round_trips_multi_level_nesting(tmp_path): cfg = { "stage1_model": { "hidden_dim": 256, "router": { "enabled": True, "type": "energy", "n_experts": 4, }, }, "train": {"epochs": 100}, } meta = {"config_version": 3, "git_hash": "abc123"} gconfig.save_config(cfg, tmp_path, meta) loaded = gconfig.load_toml(tmp_path / "config.toml") assert loaded["stage1_model"]["hidden_dim"] == 256 assert loaded["stage1_model"]["router"] == { "enabled": True, "type": "energy", "n_experts": 4, } assert loaded["train"] == {"epochs": 100} assert loaded["meta"] == meta def test_save_config_round_trips_three_level_nesting(tmp_path): cfg = { "stage2_model": { "decoder": "autoregressive", "n_sec": {"mode": "head", "lambda": 0.1}, "router": {"tie_to_stage1": True}, } } gconfig.save_config(cfg, tmp_path, {"config_version": 3}) loaded = gconfig.load_toml(tmp_path / "config.toml") assert loaded["stage2_model"]["decoder"] == "autoregressive" assert loaded["stage2_model"]["n_sec"] == {"mode": "head", "lambda": 0.1} assert loaded["stage2_model"]["router"] == {"tie_to_stage1": True} # --------------------------------------------------------------------------- # default_out_dir_name # --------------------------------------------------------------------------- _NOW = datetime(2026, 7, 29, 14, 30) def _cfg_with(**dotted_overrides): """Build a full DEFAULT_CONFIG-shaped dict with dotted-path overrides applied via _deep_merge, e.g. _cfg_with(**{"stage1_model.hidden_dim": 512}).""" overrides: dict = {} for dotted, value in dotted_overrides.items(): gconfig._set_path(overrides, dotted, value) return gconfig._deep_merge(gconfig.DEFAULT_CONFIG, overrides) def test_default_out_dir_name_all_defaults_is_just_the_timestamp(): assert gconfig.default_out_dir_name(gconfig.DEFAULT_CONFIG, now=_NOW) == "20260729_1430" def test_default_out_dir_name_stage1_generator_shown_bare_no_prefix(): cfg = _cfg_with(**{"stage1_model.generator": "wgan"}) assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_wgan" def test_default_out_dir_name_stage2_decoder_shown(): cfg = _cfg_with(**{"stage2_model.decoder": "one_shot"}) assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_dec-one_shot" def test_default_out_dir_name_particle_type_target_shown(): cfg = _cfg_with(**{"stage2_model.particle_type.target": "physical"}) assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_pt-physical" def test_default_out_dir_name_particle_conditioning_embedding_abbreviated(): cfg = _cfg_with(**{"conditioning.particle.type": "embedding"}) assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_cemb" def test_default_out_dir_name_stage1_router_shown_as_unit(): cfg = _cfg_with( **{ "stage1_model.router.enabled": True, "stage1_model.router.type": "energy", "stage1_model.router.n_experts": 8, } ) assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8" def test_default_out_dir_name_stage2_router_shown_as_unit_distinct_from_stage1(): cfg = _cfg_with( **{ "stage2_model.router.enabled": True, "stage2_model.router.type": "pdg", "stage2_model.router.n_experts": 3, } ) assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s2r-pdg3" def test_default_out_dir_name_router_disabled_omitted_even_if_subfields_nondefault(): cfg = _cfg_with( **{ "stage1_model.router.enabled": False, "stage1_model.router.type": "pdg", "stage1_model.router.n_experts": 8, } ) assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430" def test_default_out_dir_name_router_gumbel_shown_when_enabled(): cfg = _cfg_with( **{ "stage1_model.router.enabled": True, "stage1_model.router.type": "energy", "stage1_model.router.n_experts": 8, "stage1_model.router.gumbel": True, } ) assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8_s1gum" def test_default_out_dir_name_overflow_caps_and_hashes_remainder(): cfg = _cfg_with( **{ "stage1_model.generator": "wgan", "stage2_model.generator": "flow", "stage2_model.decoder": "one_shot", "stage2_model.autoregressive.history": "attention", "stage2_model.particle_type.target": "physical", "stage1_model.router.enabled": True, "stage1_model.router.type": "energy", "stage1_model.router.n_experts": 8, "stage2_model.router.enabled": True, "stage2_model.router.type": "pdg", "stage2_model.router.n_experts": 3, } ) name = gconfig.default_out_dir_name(cfg, now=_NOW) # First 6 by priority: stage1_generator, stage2_generator, stage2_decoder, # stage2_history, particle_type_target, stage1_router — stage2_router # overflows into the hash suffix. assert name.startswith("20260729_1430_wgan_s2-flow_dec-one_shot_hist-attention_pt-physical_s1r-energy8_+") def test_default_out_dir_name_overflow_hash_is_deterministic_and_value_sensitive(): overrides = { "stage1_model.generator": "wgan", "stage2_model.generator": "flow", "stage2_model.decoder": "one_shot", "stage2_model.autoregressive.history": "attention", "stage2_model.particle_type.target": "physical", "stage1_model.router.enabled": True, "stage1_model.router.type": "energy", "stage1_model.router.n_experts": 8, "train.seed": 3, } name_a = gconfig.default_out_dir_name(_cfg_with(**overrides), now=_NOW) name_b = gconfig.default_out_dir_name(_cfg_with(**overrides), now=_NOW) assert name_a == name_b changed = dict(overrides, **{"train.seed": 99}) name_c = gconfig.default_out_dir_name(_cfg_with(**changed), now=_NOW) assert name_c != name_a # --------------------------------------------------------------------------- # validate_config # --------------------------------------------------------------------------- def test_validate_config_default_config_passes(): gconfig.validate_config(gconfig.DEFAULT_CONFIG) # must not raise def test_validate_config_embedding_target_requires_embedding_conditioning(): cfg = _cfg_with( **{ "stage2_model.particle_type.target": "embedding", "conditioning.particle.type": "physical", } ) try: gconfig.validate_config(cfg) assert False, "expected ValueError" except ValueError as e: assert "embedding" in str(e) def test_validate_config_embedding_target_passes_with_embedding_conditioning(): cfg = _cfg_with( **{ "stage2_model.particle_type.target": "embedding", "conditioning.particle.type": "embedding", "conditioning.material.type": "embedding", } ) gconfig.validate_config(cfg) # must not raise def test_validate_config_mixed_particle_material_conditioning_is_valid(): """The particle and material conditioning axes are configured independently and may mix freely — e.g. material "physical" with particle "embedding" — and the data pipeline (giant/data/transforms.py) now implements that end-to-end, so validate_config must not reject it.""" cfg = _cfg_with( **{ "conditioning.particle.type": "physical", "conditioning.material.type": "embedding", } ) gconfig.validate_config(cfg) # must not raise def test_validate_config_pdg_router_incompatible_with_physical_conditioning(): cfg = _cfg_with( **{ "stage1_model.router.enabled": True, "stage1_model.router.type": "pdg", "conditioning.particle.type": "physical", } ) try: gconfig.validate_config(cfg) assert False, "expected ValueError" except ValueError as e: assert "pdg" in str(e) def test_validate_config_tie_to_stage1_requires_stage1_active(): cfg = _cfg_with( **{ "stage2_model.router.tie_to_stage1": True, "stage1_model.active": False, } ) try: gconfig.validate_config(cfg) assert False, "expected ValueError" except ValueError as e: assert "tie_to_stage1" in str(e) def test_validate_config_stop_token_accepted_under_autoregressive(): """DEFAULT_CONFIG's stage2_model.decoder is already "autoregressive" (see test_stage2_model_config_defaults_match_documented_v030_intent), so mode="stop_token" alone must not raise.""" cfg = _cfg_with(**{"stage2_model.n_sec.mode": "stop_token"}) gconfig.validate_config(cfg) # must not raise def test_validate_config_stop_token_rejected_under_one_shot(): cfg = _cfg_with( **{ "stage2_model.n_sec.mode": "stop_token", "stage2_model.decoder": "one_shot", } ) try: gconfig.validate_config(cfg) assert False, "expected ValueError" except ValueError as e: assert "stop_token" in str(e) and "autoregressive" in str(e) def test_validate_config_stop_token_rejected_for_stage1_owner(): cfg = _cfg_with( **{ "stage2_model.n_sec.mode": "stop_token", "stage2_model.n_sec.owner": "stage1", } ) try: gconfig.validate_config(cfg) assert False, "expected ValueError" except ValueError as e: assert "stop_token" in str(e) and "owner" in str(e) def test_validate_config_bad_stop_sampling_rejected(): cfg = _cfg_with(**{"stage2_model.n_sec.stop_sampling": "bogus"}) try: gconfig.validate_config(cfg) assert False, "expected ValueError" except ValueError as e: assert "stop_sampling" in str(e) def test_validate_config_stage1_context_sampled_not_implemented(): cfg = _cfg_with(**{"stage2_model.stage1_context": "sampled"}) try: gconfig.validate_config(cfg) assert False, "expected ValueError" except ValueError as e: assert "sampled" in str(e) def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint(): """'n_sec.mode = "truth" is invalid for a rollout-capable checkpoint' — both stages active means giant rollout could load this checkpoint, but 'truth' has no ground truth to draw n_sec from at rollout time.""" cfg = _cfg_with( **{ "stage2_model.n_sec.mode": "truth", "stage1_model.active": True, "stage2_model.active": True, } ) try: gconfig.validate_config(cfg) assert False, "expected ValueError" except ValueError as e: assert "n_sec.mode" in str(e) and "truth" in str(e) def test_validate_config_n_sec_truth_allowed_for_stage2_only_checkpoint(): """'truth' is exactly the standalone stage-2 evaluation mode the design doc carves out — stage1_model.active = false must still pass.""" cfg = _cfg_with( **{ "stage2_model.n_sec.mode": "truth", "stage1_model.active": False, "stage2_model.active": True, } ) gconfig.validate_config(cfg) # must not raise def test_validate_config_n_sec_truth_allowed_when_stage2_inactive(): cfg = _cfg_with( **{ "stage2_model.n_sec.mode": "truth", "stage1_model.active": True, "stage2_model.active": False, } ) gconfig.validate_config(cfg) # must not raise def test_validate_config_ar_default_markov_always_passes(): """DEFAULT_CONFIG already has decoder='autoregressive', history='markov', teacher_forcing='always' — must not raise (v0.3.0 step 5; see also test_validate_config_default_config_passes).""" cfg = _cfg_with(**{"stage2_model.decoder": "autoregressive"}) gconfig.validate_config(cfg) # must not raise def test_validate_config_ar_order_energy_desc_passes(): """'energy_desc' is the only implemented order — must not raise.""" cfg = _cfg_with( **{ "stage2_model.decoder": "autoregressive", "stage2_model.autoregressive.order": "energy_desc", } ) gconfig.validate_config(cfg) # must not raise def test_validate_config_ar_order_invalid_value_rejected(): cfg = _cfg_with( **{ "stage2_model.decoder": "autoregressive", "stage2_model.autoregressive.order": "energy_asc", } ) try: gconfig.validate_config(cfg) assert False, "expected ValueError" except ValueError as e: assert "order" in str(e) def test_validate_config_ar_history_attention_passes(): """v0.3.0 step 7 implements history='attention' — must not raise.""" cfg = _cfg_with( **{ "stage2_model.decoder": "autoregressive", "stage2_model.autoregressive.history": "attention", } ) gconfig.validate_config(cfg) # must not raise @pytest.mark.parametrize("teacher_forcing", ["scheduled", "never"]) def test_validate_config_ar_teacher_forcing_scheduled_or_never_passes(teacher_forcing): """v0.3.0 step 7 implements teacher_forcing in {'scheduled', 'never'} — must not raise.""" cfg = _cfg_with( **{ "stage2_model.decoder": "autoregressive", "stage2_model.autoregressive.teacher_forcing": teacher_forcing, } ) gconfig.validate_config(cfg) # must not raise def test_validate_config_ar_history_invalid_value_rejected(): cfg = _cfg_with( **{ "stage2_model.decoder": "autoregressive", "stage2_model.autoregressive.history": "bogus", } ) try: gconfig.validate_config(cfg) assert False, "expected ValueError" except ValueError as e: assert "history" in str(e) def test_validate_config_ar_teacher_forcing_invalid_value_rejected(): cfg = _cfg_with( **{ "stage2_model.decoder": "autoregressive", "stage2_model.autoregressive.teacher_forcing": "bogus", } ) try: gconfig.validate_config(cfg) assert False, "expected ValueError" except ValueError as e: assert "teacher_forcing" in str(e) def test_validate_config_ar_checks_skipped_under_one_shot(): """order/history/teacher_forcing values that would fail under AR are irrelevant (and unchecked) when decoder='one_shot'.""" cfg = _cfg_with( **{ "stage2_model.decoder": "one_shot", "stage2_model.autoregressive.order": "bogus", "stage2_model.autoregressive.history": "attention", "stage2_model.autoregressive.teacher_forcing": "scheduled", } ) gconfig.validate_config(cfg) # must not raise # --------------------------------------------------------------------------- # validate_config_keys / merge_cli_overrides unknown-key rejection # --------------------------------------------------------------------------- def test_validate_config_keys_default_config_passes(): gconfig.validate_config_keys(gconfig.DEFAULT_CONFIG) # must not raise def test_validate_config_keys_rejects_unknown_top_level_key(): cfg = _cfg_with(**{"bogus_section.foo": 1}) try: gconfig.validate_config_keys(cfg) assert False, "expected ValueError" except ValueError as e: assert "bogus_section" in str(e) def test_validate_config_keys_rejects_unknown_nested_key_with_close_match_hint(): cfg = _cfg_with(**{"stage1_model.n_res_block": 12}) try: gconfig.validate_config_keys(cfg) assert False, "expected ValueError" except ValueError as e: assert "stage1_model.n_res_block" in str(e) assert "n_res_blocks" in str(e) def test_validate_config_keys_allows_composed_router_axis_keys(): cfg = _cfg_with( **{ "stage1_model.router.enabled": True, "stage1_model.router.type": "composed", "stage1_model.router.axis0_type": "energy", "stage1_model.router.axis0_n_experts": 4, "stage1_model.router.axis1_type": "pdg", "stage1_model.router.axis1_emb_dim": 8, } ) gconfig.validate_config_keys(cfg) # must not raise def test_validate_config_keys_allows_centers_init(): cfg = _cfg_with(**{"stage1_model.router.centers_init": [-1.0, 0.0, 1.0]}) gconfig.validate_config_keys(cfg) # must not raise def test_validate_config_keys_rejects_unrelated_unknown_router_key(): cfg = _cfg_with(**{"stage1_model.router.n_expert": 4}) # typo for n_experts try: gconfig.validate_config_keys(cfg) assert False, "expected ValueError" except ValueError as e: assert "stage1_model.router.n_expert" in str(e) assert "n_experts" in str(e) def test_validate_config_keys_skips_meta_section(): cfg = _cfg_with() cfg["meta"] = {"config_version": 3, "git_hash": "abc123"} gconfig.validate_config_keys(cfg) # must not raise def test_validate_config_keys_allows_trunk_type(): cfg = _cfg_with(**{"stage1_model.trunk.type": "resmlp", "stage2_model.trunk.type": "resmlp"}) gconfig.validate_config_keys(cfg) # must not raise def test_validate_config_keys_rejects_unknown_trunk_key(): cfg = _cfg_with(**{"stage1_model.trunk.type_o": "resmlp"}) # typo for type try: gconfig.validate_config_keys(cfg) assert False, "expected ValueError" except ValueError as e: assert "stage1_model.trunk.type_o" in str(e) assert "type" in str(e) def test_validate_config_keys_allows_block_conditioning(): cfg = _cfg_with( **{ "stage1_model.trunk.block_conditioning": "film", "stage2_model.trunk.block_conditioning": "adaln", } ) gconfig.validate_config_keys(cfg) # must not raise def test_validate_config_keys_rejects_unknown_block_conditioning_key(): cfg = _cfg_with(**{"stage1_model.trunk.block_conditioning_o": "film"}) # typo try: gconfig.validate_config_keys(cfg) assert False, "expected ValueError" except ValueError as e: assert "stage1_model.trunk.block_conditioning_o" in str(e) assert "block_conditioning" in str(e) def test_merge_cli_overrides_rejects_typo_in_toml_file(tmp_path): path = tmp_path / "config.toml" path.write_text("[meta]\nconfig_version = 3\n\n[stage1_model]\nn_res_block = 12\n") try: gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) assert False, "expected ValueError" except ValueError as e: assert "n_res_block" in str(e) def test_merge_cli_overrides_rejects_typo_in_cli_overrides(): try: gconfig.merge_cli_overrides( gconfig.DEFAULT_CONFIG, None, {"stage1_model": {"n_res_block": 12}}, ) assert False, "expected ValueError" except ValueError as e: assert "n_res_block" in str(e) @pytest.mark.parametrize("fixture_name", ["default.toml", "wgan_h128_b4_physical.toml"]) def test_merge_cli_overrides_real_config_fixtures_pass_key_validation(fixture_name, monkeypatch): monkeypatch.setattr(gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243") gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / fixture_name, {}) # must not raise # --------------------------------------------------------------------------- # overrides_from_flags (issues.md Issue 3): the flag -> config-path table # shared by `giant train`/`giant new-run`. Each test below pins one # precedence rule directly, without CliRunner — see also # tests/test_cli_train_overrides.py for the thin end-to-end smoke coverage. # --------------------------------------------------------------------------- def test_overrides_from_flags_empty_values_yield_empty_overrides(): assert gconfig.overrides_from_flags({}) == {} assert gconfig.overrides_from_flags({"epochs": None, "hidden_dim": None}) == {} def test_overrides_from_flags_train_block_passthrough(): overrides = gconfig.overrides_from_flags({"epochs": 5, "lr": 1e-3, "hidden_dim": None}) assert overrides == {"train": {"epochs": 5, "lr": 1e-3}} @pytest.mark.parametrize( ("shorthand", "explicit", "path_key"), [ ("hidden_dim", "stage1_hidden_dim", "hidden_dim"), ("n_blocks", "stage1_n_res_blocks", "n_res_blocks"), ("dropout", "stage1_dropout", "dropout"), ], ) def test_overrides_from_flags_stage1_explicit_overrides_shorthand(shorthand, explicit, path_key): overrides = gconfig.overrides_from_flags({shorthand: 1, explicit: 2}) assert overrides["stage1_model"][path_key] == 2 @pytest.mark.parametrize( ("shorthand", "path_key"), [("hidden_dim", "hidden_dim"), ("n_blocks", "n_res_blocks"), ("dropout", "dropout")], ) def test_overrides_from_flags_stage1_shorthand_alone(shorthand, path_key): overrides = gconfig.overrides_from_flags({shorthand: 7}) assert overrides["stage1_model"][path_key] == 7 def test_overrides_from_flags_stage2_only_knobs(): overrides = gconfig.overrides_from_flags( { "stage2_hidden_dim": 32, "stage2_n_res_blocks": 4, "stage2_dropout": 0.1, "stage2_decoder": "one_shot", "stage2_k_max": 8, "stage2_context_dim": 16, "stage2_stage1_context": "sampled", } ) assert overrides["stage2_model"] == { "hidden_dim": 32, "n_res_blocks": 4, "dropout": 0.1, "decoder": "one_shot", "k_max": 8, "context_dim": 16, "stage1_context": "sampled", } assert "stage1_model" not in overrides def test_overrides_from_flags_mode_fans_to_both_stages(): overrides = gconfig.overrides_from_flags({"mode": "wgan"}) assert overrides["stage1_model"]["generator"] == "wgan" assert overrides["stage2_model"]["generator"] == "wgan" def test_overrides_from_flags_stage1_generator_overrides_mode_for_stage1_only(): overrides = gconfig.overrides_from_flags({"mode": "wgan", "stage1_generator": "flow"}) assert overrides["stage1_model"]["generator"] == "flow" assert overrides["stage2_model"]["generator"] == "wgan" def test_overrides_from_flags_stage2_generator_overrides_mode_for_stage2_only(): overrides = gconfig.overrides_from_flags({"mode": "wgan", "stage2_generator": "flow"}) assert overrides["stage1_model"]["generator"] == "wgan" assert overrides["stage2_model"]["generator"] == "flow" def test_overrides_from_flags_emb_dim_sets_both_conditioning_axes(): overrides = gconfig.overrides_from_flags({"emb_dim": 24}) assert overrides["conditioning"]["particle"]["emb_dim"] == 24 assert overrides["conditioning"]["material"]["emb_dim"] == 24 def test_overrides_from_flags_conditioning_sets_both_axes_type(): overrides = gconfig.overrides_from_flags({"conditioning": "onehot"}) assert overrides["conditioning"]["particle"]["type"] == "onehot" assert overrides["conditioning"]["material"]["type"] == "onehot" def test_overrides_from_flags_router_config_only_touches_stage1(): overrides = gconfig.overrides_from_flags({"router_config": {"enabled": True, "type": "energy"}}) assert overrides["stage1_model"]["router"] == {"enabled": True, "type": "energy"} assert "stage2_model" not in overrides @pytest.mark.parametrize( ("shared", "stage1_specific", "stage2_specific", "path_key"), [ ("n_critic", "stage1_n_critic", "stage2_n_critic", "n_critic"), ("gp_weight", "stage1_gp_weight", "stage2_gp_weight", "gp_weight"), ("noise_dim", "stage1_noise_dim", "stage2_noise_dim", "noise_dim"), ("critic_lr", "stage1_critic_lr", "stage2_critic_lr", "critic_lr"), ], ) def test_overrides_from_flags_wgan_knobs_split_per_stage(shared, stage1_specific, stage2_specific, path_key): overrides = gconfig.overrides_from_flags({shared: 5.0, stage1_specific: 3.0}) assert overrides["stage1_model"]["wgan"][path_key] == 3.0 assert overrides["stage2_model"]["wgan"][path_key] == 5.0 overrides = gconfig.overrides_from_flags({shared: 5.0, stage2_specific: 2.5}) assert overrides["stage1_model"]["wgan"][path_key] == 5.0 assert overrides["stage2_model"]["wgan"][path_key] == 2.5 @pytest.mark.parametrize( ("stage_flag", "stage_model", "path_key"), [ ("stage1_critic_hidden_dim", "stage1_model", "critic_hidden_dim"), ("stage1_critic_n_res_blocks", "stage1_model", "critic_n_res_blocks"), ("stage2_critic_hidden_dim", "stage2_model", "critic_hidden_dim"), ("stage2_critic_n_res_blocks", "stage2_model", "critic_n_res_blocks"), ], ) def test_overrides_from_flags_critic_sizing_is_stage_scoped_only(stage_flag, stage_model, path_key): """critic_hidden_dim/critic_n_res_blocks are architectural per-stage knobs (gitea #28) — unlike n_critic/gp_weight/noise_dim/critic_lr above, there is deliberately no shared alias that fans out to both stages.""" overrides = gconfig.overrides_from_flags({stage_flag: 32}) assert overrides == {stage_model: {"wgan": {path_key: 32}}} # --------------------------------------------------------------------------- # checkpoint config-mismatch warnings (unchanged surface, still exercised) # --------------------------------------------------------------------------- def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "current999") ckpt_path = tmp_path / "best.pt" ckpt_path.write_bytes(b"") _write_toml(tmp_path / "config.toml", git_hash="old111", extra="[train]\nepochs = 5\n") gconfig.warn_if_checkpoint_config_mismatch(ckpt_path) captured = capsys.readouterr() assert "warning" in captured.err assert "old111" in captured.err assert "current999" in captured.err def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "current999") ckpt_path = tmp_path / "best.pt" ckpt_path.write_bytes(b"") gconfig.warn_if_checkpoint_config_mismatch(ckpt_path) assert capsys.readouterr().err == "" def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "same123") ckpt_path = tmp_path / "best.pt" ckpt_path.write_bytes(b"") _write_toml(tmp_path / "config.toml", git_hash="same123", extra="[train]\nepochs = 5\n") gconfig.warn_if_checkpoint_config_mismatch(ckpt_path) assert capsys.readouterr().err == ""