Decouple secondary-species vocabulary from conditioning.particle.emb_dim (gitea #29)

conditioning.particle.emb_dim and stage2_model.particle_type.target="onehot"'s
class count were silently the same number everywhere (pipeline.py's PDG
top-N map build, Stage2OneShot/Stage2Autoregressive's type head, StageSpec's
training loss width, the checkpoint's shared pdg_topn_map), fixing the
secondary-species vocabulary at whatever width the unrelated
physical-conditioning MLP happened to use — the exact vocabulary the v0.3.0
pivot exists to fix.

Adds stage2_model.particle_type.n_classes (default 0 = inherit
conditioning.particle.emb_dim, preserving today's behavior and every
existing checkpoint) and a single resolve_type_n_classes helper used
everywhere the coupling used to be implicit. Splits the checkpoint's shared
pdg_topn_map into a conditioning-only pdg_topn_map and a new
sec_type_topn_map, built independently through the existing
(axis, n_classes)-keyed setup cache (no extra scan when they still resolve
to the same N) and threaded through giant predict/giant rollout's decode
path. A checkpoint with no sec_type_topn_map key (pre-#29) falls back to
reusing pdg_topn_map, reproducing the old shared behavior exactly.

Decided with the user during planning: commit directly on this branch;
represent the split as an additive sec_type_topn_map checkpoint key rather
than conditionally reusing pdg_topn_map; build the two top-N maps
independently rather than the issue's proposed build-at-max-and-slice, since
the setup cache already avoids redundant scans across runs.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

tests/test_checkpoint_io.py

@@ -73,6 +73,16 @@ def _write_checkpoint(tmp_path, model_cfg=None, ema: bool = False, **ckpt_overri
     if ema:
         ckpt["model_ema"] = stage1.state_dict() if stage1 is not None else {}
         ckpt["sec_decoder_ema"] = stage2.state_dict() if stage2 is not None else {}
+    # DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
+    # "onehot", and giant train's pipeline (gitea #29) now always writes a
+    # sec_type_topn_map in that case — default one in here too, unless a
+    # test explicitly overrides it, so fixtures represent a real, loadable
+    # checkpoint by default rather than exercising the "missing" guard by
+    # accident.
+    particle_type_target = cfg.get("stage2_model", {}).get("particle_type", {}).get("target", "onehot")
+    if particle_type_target == "onehot" and "sec_type_topn_map" not in ckpt_overrides:
+        default_sec_type_topn = TopNMap(class_map=dict(zip(PDG_MAP, range(len(PDG_MAP)))), other_members={})
+        ckpt["sec_type_topn_map"] = topnmap_to_json(default_sec_type_topn)
     ckpt.update(ckpt_overrides)
     path = tmp_path / "ckpt.pt"
     torch.save(ckpt, path)
@@ -179,6 +189,40 @@ def test_onehot_particle_conditioning_with_topn_map_succeeds(tmp_path):
     assert ctx.pdg_topn_map.class_map == {11: 0, 22: 1}
 
 
+def test_onehot_particle_type_target_without_sec_type_topn_map_raises(tmp_path):
+    """DEFAULT_CONFIG's stage2_model.particle_type.target="onehot" needs a
+    sec_type_topn_map (gitea #29) — a checkpoint with neither key at all
+    (not even the pre-#29 pdg_topn_map to fall back to) must fail loudly."""
+    checkpoint = _write_checkpoint(tmp_path, sec_type_topn_map=None)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["sec_type_topn_map"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="sec_type_topn_map"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_pre_gitea_29_checkpoint_falls_back_to_pdg_topn_map_for_sec_type(tmp_path):
+    """A checkpoint written before gitea #29 has no sec_type_topn_map key at
+    all — conditioning and secondary-type onehot maps were always the same
+    map, saved once under pdg_topn_map. load_for_inference must reproduce
+    that exact pre-#29 behavior for such a checkpoint."""
+    topn = TopNMap(class_map={11: 0, 22: 1, -11: 2}, other_members={})
+    checkpoint = _write_checkpoint(
+        tmp_path,
+        model_cfg=_onehot_model_cfg(),
+        pdg_topn_map=topnmap_to_json(topn),
+        sec_type_topn_map=None,
+    )
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["sec_type_topn_map"]
+    torch.save(ckpt, checkpoint)
+
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+    assert ctx.sec_type_topn_map is not None
+    assert ctx.sec_type_topn_map.class_map == {11: 0, 22: 1, -11: 2}
+
+
 def test_ema_weights_requested_but_missing_raises(tmp_path):
     checkpoint = _write_checkpoint(tmp_path, ema=False)
 

tests/test_config.py

@@ -75,6 +75,16 @@ def test_stage2_model_config_defaults_match_documented_v030_intent():
     assert spec.particle_type.target == "onehot"
 
 
+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)

tests/test_network.py

@@ -288,6 +288,33 @@ def test_stage2_oneshot_forward_shape_onehot_flow_excludes_type():
     assert out.shape == (B, k_max * CONT_SLOT_DIM)
 
 
+def test_stage2_oneshot_particle_type_n_classes_overrides_conditioning_emb_dim():
+    """gitea #29: stage2_model.particle_type.n_classes, not
+    conditioning.particle.emb_dim, sizes the onehot type_head/type_dim when
+    explicitly set — the two used to be silently the same number."""
+    k_max = 5
+    particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1}
+    particle_type_cfg = {"target": "onehot", "lambda": 1.0, "n_classes": 20}
+    sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", k_max, 20)
+    model = Stage2OneShot(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=particle_cfg,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        cond_out_dim=16,
+        context_dim=8,
+        sec_dim=sec_dim,
+        generator="flow",
+        k_max=k_max,
+        particle_type_cfg=particle_type_cfg,
+    )
+    assert model.type_dim == 20  # not particle_cfg["emb_dim"] == 6
+    assert model.type_head is not None
+    assert model.type_head[-1].out_features == k_max * 20
+
+
 # --- MarkovHistory -----------------------------------------------------------
 
 
@@ -423,6 +450,29 @@ def test_stage2_autoregressive_history_invalid_raises():
         _build_stage2_ar("onehot", "wgan", history="bogus")
 
 
+def test_stage2_autoregressive_particle_type_n_classes_overrides_conditioning_emb_dim():
+    """gitea #29, Stage2Autoregressive side — see the Stage2OneShot version
+    of this test for the full rationale."""
+    particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1}
+    particle_type_cfg = {"target": "onehot", "lambda": 1.0, "n_classes": 20}
+    model = Stage2Autoregressive(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=particle_cfg,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        cond_out_dim=16,
+        context_dim=8,
+        generator="flow",
+        k_max=5,
+        particle_type_cfg=particle_type_cfg,
+    )
+    assert model.type_dim == 20  # not particle_cfg["emb_dim"] == 6
+    assert model.type_head is not None
+    assert model.type_head[-1].out_features == 20
+
+
 @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
 @pytest.mark.parametrize("generator", ["wgan", "flow"])
 @pytest.mark.parametrize("history", ["markov", "attention"])
@@ -655,6 +705,33 @@ def test_build_models_share_stages_true_shared_params_are_in_both_stage_paramete
     assert shared_ids <= {id(p) for p in stage2.parameters()}
 
 
+def test_build_models_particle_type_n_classes_overrides_conditioning_emb_dim():
+    """gitea #29 end-to-end through build_models: setting
+    stage2_model.particle_type.n_classes independently of
+    conditioning.particle.emb_dim actually resizes the built stage2 model,
+    not just the two lower-level unit tests above."""
+    cfg = _minimal_model_config(share_stages=False)  # conditioning.particle.emb_dim = 4
+    cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 11}
+    built = build_models(cfg)
+    assert built["stage2"] is not None
+    assert built["stage2"].type_dim == 11
+
+
+def test_build_critics_particle_type_n_classes_overrides_conditioning_emb_dim():
+    cfg = _minimal_model_config(share_stages=False)  # conditioning.particle.emb_dim = 4
+    cfg["stage2_model"]["generator"] = "wgan"
+    cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 4}
+    default_n_classes_critic = build_critics(cfg)["stage2"]
+    assert default_n_classes_critic is not None
+
+    cfg["stage2_model"]["particle_type"]["n_classes"] = 11
+    wider_critic = build_critics(cfg)["stage2"]
+    assert wider_critic is not None
+    # k_max=3 slots, each CONT_SLOT_DIM + n_classes wide under wgan folding —
+    # widening n_classes alone (emb_dim stays 4) must widen the critic input.
+    assert wider_critic.input_proj.in_features > default_n_classes_critic.input_proj.in_features
+
+
 # ── build_models/build_critics: DEFAULT_CONFIG fallback drift (issues.md #1) ─
 
 

tests/test_pipeline.py

@@ -152,28 +152,60 @@ def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch):
     assert "normalizer: cache hit" in joined
 
 
-def test_run_train_job_builds_caches_and_persists_pdg_topn_map(tmp_path, data):
+def test_run_train_job_builds_caches_and_persists_sec_type_topn_map(tmp_path, data):
     """DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
-    "onehot" — a plain _tiny_cfg() run must
-    build the shared pdg top-N map, cache it in the setup-cache sidecar, and
-    persist it into the checkpoint, with no extra config needed."""
+    "onehot" while conditioning.particle.type stays "physical" — a plain
+    _tiny_cfg() run must build the secondary-type-only pdg top-N map (gitea
+    #29: no longer shared with any conditioning-side onehot map), cache it in
+    the setup-cache sidecar, and persist it into the checkpoint's
+    sec_type_topn_map key, with no extra config needed. pdg_topn_map
+    (conditioning-only) stays unbuilt since conditioning.particle.type is
+    "physical" here."""
     echo1 = _run(data, tmp_path / "out1")
     assert any("building pdg top-N map" in m for m in echo1)
 
     loaded = setup_cache.load(data, [data])
     assert loaded is not None
-    key = setup_cache.topn_key("pdg", 4)  # conditioning.particle.emb_dim = 4
+    # stage2_model.particle_type.n_classes = 0 -> conditioning.particle.emb_dim = 4
+    key = setup_cache.topn_key("pdg", 4)
     assert key in loaded.topn_maps
     assert set(loaded.topn_maps[key].class_map.keys()) >= {11, 22}
 
     ckpt = torch.load(tmp_path / "out1" / "last.pt", weights_only=False)
-    assert "pdg_topn_map" in ckpt
-    assert set(ckpt["pdg_topn_map"]["class_map"].keys()) >= {"11", "22"}
+    assert ckpt.get("pdg_topn_map") is None
+    assert "sec_type_topn_map" in ckpt
+    assert set(ckpt["sec_type_topn_map"]["class_map"].keys()) >= {"11", "22"}
 
     echo2 = _run(data, tmp_path / "out2")
     assert any("pdg top-N map: cache hit" in m for m in echo2)
 
 
+def test_run_train_job_independent_cond_and_sec_type_topn_maps(tmp_path, data):
+    """conditioning.particle.type="onehot" and
+    stage2_model.particle_type.target="onehot" with different class counts
+    (gitea #29's fix: stage2_model.particle_type.n_classes decouples the two)
+    build two distinct top-N maps, cached under their own (axis, n_classes)
+    key and persisted under two distinct checkpoint keys — no longer forced
+    to share conditioning.particle.emb_dim."""
+    cfg = _tiny_cfg()
+    cfg["conditioning"]["particle"]["type"] = "onehot"  # emb_dim = 4, from _tiny_cfg
+    cfg["stage2_model"]["particle_type"]["n_classes"] = 3
+    echo = _run(data, tmp_path / "out", cfg=cfg)
+    assert any("mapped to 4 classes" in m for m in echo)
+    assert any("mapped to 3 classes" in m for m in echo)
+
+    loaded = setup_cache.load(data, [data])
+    assert loaded is not None
+    cond_key = setup_cache.topn_key("pdg", 4)
+    type_key = setup_cache.topn_key("pdg", 3)
+    assert cond_key in loaded.topn_maps
+    assert type_key in loaded.topn_maps
+
+    ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
+    assert ckpt.get("pdg_topn_map") is not None
+    assert ckpt.get("sec_type_topn_map") is not None
+
+
 def test_run_train_job_builds_caches_and_persists_material_topn_map(tmp_path, data):
     """conditioning.material.type="onehot" is an independent axis from the
     pdg one above, with its own build/cache-hit branch in run_setup_stage —

tests/test_rollout.py

@@ -430,7 +430,7 @@ def _run_v3(
     max_tracks_per_event=100,
     seeds=None,
     conditioning="physical",
-    pdg_topn_map=None,
+    sec_type_topn_map=None,
     other_policy="sample",
     seed=0,
     stage1_ddpm_steps=1000,
@@ -457,7 +457,7 @@ def _run_v3(
         escape_threshold=escape_threshold,
         particle_conditioning=conditioning,
         material_conditioning=conditioning,
-        pdg_topn_map=pdg_topn_map,
+        sec_type_topn_map=sec_type_topn_map,
         other_policy=other_policy,
         seed=seed,
         stage1_ddpm_steps=stage1_ddpm_steps,
@@ -532,7 +532,7 @@ def test_rollout_onehot_target_end_to_end(fake_material_props, decoder):
     (giant.particles.particle_phys_array) become the secondary's identity —
     unlike "physical", not just a reporting label."""
     s1, s2 = _models_v3(decoder=decoder, target="onehot", emb_dim=3)
-    rec = _run_v3(s1, s2, pdg_topn_map=PDG_TOPN_MAP, other_policy="modal")
+    rec = _run_v3(s1, s2, sec_type_topn_map=PDG_TOPN_MAP, other_policy="modal")
     assert len(rec["event_id"]) > 0
     # Every spawned secondary's nominal pdg must be one decode_topn_class can
     # actually produce (the topn map's known classes + its "other" members).
@@ -543,8 +543,8 @@ def test_rollout_onehot_target_end_to_end(fake_material_props, decoder):
 
 def test_rollout_onehot_target_missing_topn_map_raises(fake_material_props):
     s1, s2 = _models_v3(target="onehot", emb_dim=3)
-    with pytest.raises(RuntimeError, match="pdg_topn_map"):
-        _run_v3(s1, s2, pdg_topn_map=None)
+    with pytest.raises(RuntimeError, match="sec_type_topn_map"):
+        _run_v3(s1, s2, sec_type_topn_map=None)
 
 
 # --- conditioning.{particle,material}.type = "onehot" — a separate axis from
@@ -627,6 +627,86 @@ def test_rollout_conditioning_onehot_material_missing_topn_map_raises(
         _run_onehot_conditioning(mat_topn_map=None)
 
 
+SEC_TYPE_TOPN_MAP_DIFFERENT_N = TopNMap(class_map={22: 0, 11: 1, -11: 2, 13: 3}, other_members={2112: 3, 2212: 1})
+
+
+def _run_conditioning_and_type_onehot_different_n_classes():
+    """Both conditioning.particle.type="onehot" and
+    stage2_model.particle_type.target="onehot" active at once, with
+    stage2_model.particle_type.n_classes deliberately different from
+    conditioning.particle.emb_dim (gitea #29)."""
+    cond_emb_dim = len(PDG_MAP)  # 3
+    type_n_classes = 5  # deliberately different from cond_emb_dim
+    particle_cfg = {"type": "onehot", "emb_dim": cond_emb_dim, "n_layers": 1}
+    material_cfg = {"type": "onehot", "emb_dim": len(MAT_MAP), "n_layers": 1}
+    particle_type_cfg = {"target": "onehot", "n_classes": type_n_classes}
+
+    s1 = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+    ).eval()
+    sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", K_MAX, type_n_classes)
+    s2 = Stage2OneShot(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+        sec_dim=sec_dim,
+        generator="flow",
+        time_dim=16,
+        k_max=K_MAX,
+        particle_type_cfg=particle_type_cfg,
+    ).eval()
+    # Sanity: the model's own type_dim followed n_classes, not cond_emb_dim.
+    assert s2.type_dim == type_n_classes
+
+    cond, tgt, sec_phys = _norms()
+    return rollout(
+        s1,
+        s2,
+        _oracle(),
+        _seeds(),
+        cond,
+        tgt,
+        sec_phys,
+        PDG_MAP,
+        MAT_MAP,
+        energy_cutoff=1.0,
+        max_steps=15,
+        steps=3,
+        batch_size=128,
+        max_tracks_per_event=100,
+        escape_threshold=1e9,
+        particle_conditioning="onehot",
+        material_conditioning="onehot",
+        pdg_topn_map=COND_PDG_TOPN_MAP,
+        mat_topn_map=COND_MAT_TOPN_MAP,
+        sec_type_topn_map=SEC_TYPE_TOPN_MAP_DIFFERENT_N,
+        other_policy="modal",
+    )
+
+
+def test_rollout_conditioning_and_type_onehot_with_different_n_classes(fake_material_props):
+    """gitea #29 end-to-end: conditioning.particle.type="onehot" and
+    stage2_model.particle_type.target="onehot" now use independently sized
+    top-N maps (stage2_model.particle_type.n_classes != conditioning.particle
+    .emb_dim), and rollout must decode secondaries using the type-side map,
+    not silently reuse the conditioning-side one (the pre-#29 bug)."""
+    rec = _run_conditioning_and_type_onehot_different_n_classes()
+    assert len(rec["event_id"]) > 0
+    possible = set(SEC_TYPE_TOPN_MAP_DIFFERENT_N.class_map.keys()) | set(
+        SEC_TYPE_TOPN_MAP_DIFFERENT_N.other_members.keys()
+    )
+    secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
+    assert secondary_pdgs <= possible
+
+
 @pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
 def test_rollout_embedding_target_end_to_end(decoder):
     """particle_type.target="embedding" L1-snaps to the nearest row of the