v0.3.0 step 4: type map + particle_type.target = "onehot"/"embedding"

Builds the shared top-N-plus-other PDG/material maps (pooling both primary
and secondary occurrences for PDG, directly targeting the meeting's
species-collapse failure mode) and wires up conditioning.{particle,material}
= "onehot" plus stage2_model.particle_type.target in ("onehot", "embedding")
end-to-end: setup-cache persistence, Stage2OneShot's type_head (flow/ddpm)
vs. folded+ST-Gumbel-relaxed adversarial slice (wgan), and the corresponding
CE/MSE training losses. particle_type.target = "physical" stays byte-for-byte
unchanged, keeping the v0.2 migration shim's bit-identical guarantee intact.
giant predict/rollout fail loudly on a onehot/embedding checkpoint until
full decode support lands in step 6.

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

tests/test_cli_predict.py

@@ -1,14 +1,72 @@
 import uuid
 
+import pytest
+import typer
 import yaml
 
 from giant.cli import (
     _CEPH_PREDICTIONS,
+    _check_v030_onehot_support,
     _resolve_prediction_output,
     _write_prediction_ref,
 )
 
 
+# ---------------------------------------------------------------------------
+# _check_v030_onehot_support
+# ---------------------------------------------------------------------------
+
+
+def _nested_model_cfg(
+    particle_type="physical", material_type="physical", target="physical"
+):
+    return {
+        "conditioning": {
+            "particle": {"type": particle_type, "emb_dim": 8},
+            "material": {"type": material_type, "emb_dim": 8},
+        },
+        "stage2_model": {"particle_type": {"target": target}},
+    }
+
+
+def test_check_v030_onehot_support_allows_physical():
+    _check_v030_onehot_support(_nested_model_cfg(), "predict")  # no raise
+
+
+def test_check_v030_onehot_support_rejects_onehot_particle_conditioning():
+    cfg = _nested_model_cfg(particle_type="onehot")
+    with pytest.raises(typer.Exit):
+        _check_v030_onehot_support(cfg, "predict")
+
+
+def test_check_v030_onehot_support_rejects_onehot_material_conditioning():
+    cfg = _nested_model_cfg(material_type="onehot")
+    with pytest.raises(typer.Exit):
+        _check_v030_onehot_support(cfg, "rollout")
+
+
+def test_check_v030_onehot_support_rejects_onehot_particle_type_target():
+    cfg = _nested_model_cfg(target="onehot")
+    with pytest.raises(typer.Exit):
+        _check_v030_onehot_support(cfg, "predict")
+
+
+def test_check_v030_onehot_support_rejects_embedding_particle_type_target():
+    cfg = _nested_model_cfg(
+        particle_type="embedding", material_type="embedding", target="embedding"
+    )
+    with pytest.raises(typer.Exit):
+        _check_v030_onehot_support(cfg, "predict")
+
+
+def test_check_v030_onehot_support_is_noop_for_v02_flat_model_config():
+    """A v0.2 checkpoint's flat model_config has conditioning as a plain
+    string, not a dict — never onehot/embedding-target, so this must be a
+    silent no-op rather than crash on `.get("particle")` against a string."""
+    cfg = {"conditioning": "embedding", "mode": "flow"}
+    _check_v030_onehot_support(cfg, "predict")  # no raise
+
+
 # ---------------------------------------------------------------------------
 # _resolve_prediction_output
 # ---------------------------------------------------------------------------

tests/test_loader.py

@@ -6,7 +6,9 @@ from giant.data.loader import (
     EVENT_ID_FILE_STRIDE,
     build_index_maps,
     build_index_maps_from_files,
+    build_pdg_topn_map_from_files,
     build_process_map_from_files,
+    build_topn_map_from_files,
     event_id_offset,
     find_parquet_files,
     iter_cond_chunks,
@@ -178,6 +180,64 @@ def test_build_process_map_from_files_three_files_partial_overlap(tmp_path):
     assert proc_map["compt"] == 2
 
 
+# ── build_topn_map_from_files / build_pdg_topn_map_from_files ──────────────
+
+
+def test_build_topn_map_from_files_keeps_most_frequent(tmp_path):
+    materials = ["G4_AIR"] * 5 + ["PbWO4"] * 3 + ["G4_Fe"] * 2 + ["G4_Pb"] * 1
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"material": materials}).to_parquet(path)
+
+    m = build_topn_map_from_files([path], "material", n_classes=3, cast=str)
+
+    assert m.class_map["G4_AIR"] == 0
+    assert m.class_map["PbWO4"] == 1
+    assert m.class_map["G4_Fe"] == 2  # "other" (n_classes - 1)
+    assert m.class_map["G4_Pb"] == 2
+    assert m.other_members == {"G4_Fe": 2, "G4_Pb": 1}
+
+
+def test_build_topn_map_from_files_fewer_values_than_n_classes(tmp_path):
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"material": ["G4_AIR", "PbWO4"]}).to_parquet(path)
+
+    m = build_topn_map_from_files([path], "material", n_classes=5, cast=str)
+
+    assert m.class_map == {"G4_AIR": 0, "PbWO4": 1}
+    assert m.other_members == {}
+
+
+def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path):
+    """A species that's rare as a primary but common as a secondary must
+    still rank by its pooled (primary + secondary) count, not just its
+    primary-role count alone — the whole point of pooling both roles
+    (docs/v0.3.0-design.md §8)."""
+    path = tmp_path / "a.parquet"
+    # primary pdg: mostly 11 (electron), one lone 22 (photon)
+    pdg = [11] * 5 + [22] * 1
+    # secondaries: 22 (photon) appears often as a secondary despite being
+    # rare as a primary above
+    sec_pdg_list = [[22, 22]] * 5 + [[]] * 1
+    pd.DataFrame({"pdg": pdg, "sec_pdg_list": sec_pdg_list}).to_parquet(path)
+
+    m = build_pdg_topn_map_from_files([path], n_classes=3)
+
+    # pooled: 11 -> 5, 22 -> 1 (primary) + 10 (secondary) = 11
+    assert m.class_map[22] == 0
+    assert m.class_map[11] == 1
+
+
+def test_build_pdg_topn_map_from_files_missing_sec_pdg_list_column(tmp_path):
+    """Files predating the parent->child join have no sec_pdg_list column —
+    must not raise, just count the primary pdg column alone."""
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"pdg": [11, 11, 22]}).to_parquet(path)
+
+    m = build_pdg_topn_map_from_files([path], n_classes=3)
+
+    assert m.class_map == {11: 0, 22: 1}
+
+
 # ── build_index_maps (in-memory) ────────────────────────────────────────────
 
 

tests/test_network.py

@@ -1,9 +1,19 @@
 import torch
-from giant.constants import COND_DIM
-from giant.model.network import SinusoidalEmbedding, Stage1Model
+from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
+from giant.model.network import (
+    ConditionEncoder,
+    SinusoidalEmbedding,
+    Stage1Model,
+    Stage2OneShot,
+    cat_col_layout,
+    stage2_trunk_sec_dim,
+    stage2_type_dim,
+)
 
 PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
 MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
+ONEHOT_PARTICLE_CFG = {"type": "onehot", "emb_dim": 6, "n_layers": 1}
+ONEHOT_MATERIAL_CFG = {"type": "onehot", "emb_dim": 4, "n_layers": 1}
 
 
 def test_sinusoidal_embedding_shape():
@@ -71,3 +81,212 @@ def test_stage1_model_no_n_sec_head_by_default():
         pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG
     )
     assert model.n_sec_head is None
+
+
+# --- cat_col_layout / stage2_type_dim / stage2_trunk_sec_dim ---------------
+
+
+def test_cat_col_layout_neither_onehot():
+    assert cat_col_layout("physical", "embedding") == (None, None)
+
+
+def test_cat_col_layout_particle_only():
+    assert cat_col_layout("onehot", "physical") == (2, None)
+
+
+def test_cat_col_layout_material_only():
+    assert cat_col_layout("physical", "onehot") == (None, 2)
+
+
+def test_cat_col_layout_both_onehot_particle_then_material():
+    assert cat_col_layout("onehot", "onehot") == (2, 3)
+
+
+def test_stage2_type_dim_physical_is_particle_phys_dim():
+    assert stage2_type_dim({"target": "physical"}, emb_dim=16) == PARTICLE_PHYS_DIM
+
+
+def test_stage2_type_dim_onehot_and_embedding_are_emb_dim():
+    assert stage2_type_dim({"target": "onehot"}, emb_dim=16) == 16
+    assert stage2_type_dim({"target": "embedding"}, emb_dim=16) == 16
+
+
+def test_stage2_trunk_sec_dim_physical_matches_v02_sec_dim():
+    k_max = 15
+    assert (
+        stage2_trunk_sec_dim({"target": "physical"}, "flow", k_max, emb_dim=16)
+        == k_max * SEC_SLOT_DIM
+    )
+    assert (
+        stage2_trunk_sec_dim({"target": "physical"}, "wgan", k_max, emb_dim=16)
+        == k_max * SEC_SLOT_DIM
+    )
+
+
+def test_stage2_trunk_sec_dim_onehot_wgan_folds_type_in():
+    k_max = 15
+    assert stage2_trunk_sec_dim(
+        {"target": "onehot"}, "wgan", k_max, emb_dim=16
+    ) == k_max * (CONT_SLOT_DIM + 16)
+
+
+def test_stage2_trunk_sec_dim_onehot_flow_excludes_type():
+    k_max = 15
+    assert (
+        stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16)
+        == k_max * CONT_SLOT_DIM
+    )
+
+
+# --- ConditionEncoder onehot mode -------------------------------------------
+
+
+def test_condition_encoder_onehot_forward_shape_and_gradients():
+    B = 8
+    particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
+    material_emb_dim = int(ONEHOT_MATERIAL_CFG["emb_dim"])
+    enc = ConditionEncoder(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=ONEHOT_PARTICLE_CFG,
+        material_cfg=ONEHOT_MATERIAL_CFG,
+        out_dim=32,
+    )
+    cond_cont = torch.randn(B, COND_DIM)
+    cond_cat = torch.stack(
+        [
+            torch.randint(0, 5, (B,)),
+            torch.randint(0, 3, (B,)),
+            torch.randint(0, particle_emb_dim, (B,)),
+            torch.randint(0, material_emb_dim, (B,)),
+        ],
+        dim=1,
+    )
+    out = enc(cond_cont, cond_cat)
+    assert out.shape == (B, 32)
+    # onehot itself is unlearned, but the fusion MLP downstream still has
+    # gradients — the encoder as a whole must still be trainable.
+    out.sum().backward()
+    assert enc.mlp[0].weight.grad is not None
+
+
+def test_condition_encoder_onehot_is_a_true_one_hot_vector():
+    """The onehot axis feeds a fixed, unlearned one-hot into the fusion MLP —
+    verify the concatenated input segment really is one-hot, not e.g. an
+    accidentally-learned embedding."""
+    B = 4
+    particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
+    enc = ConditionEncoder(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=ONEHOT_PARTICLE_CFG,
+        material_cfg={"type": "physical", "emb_dim": 4, "n_layers": 1},
+        out_dim=16,
+    )
+    cond_cont = torch.zeros(B, COND_DIM)
+    idx = torch.tensor([0, 1, 2, 5])
+    cond_cat = torch.stack(
+        [
+            torch.zeros(B, dtype=torch.long),
+            torch.zeros(B, dtype=torch.long),
+            idx.clamp(max=particle_emb_dim - 1),
+        ],
+        dim=1,
+    )
+    pdg_e = enc._particle_embed(cond_cont, cond_cat)
+    assert pdg_e.shape == (B, particle_emb_dim)
+    assert torch.all(pdg_e.sum(dim=-1) == 1.0)
+
+
+# --- Stage2OneShot particle_type architecture (docs/v0.3.0-design.md decision 2) --
+
+
+def _build_stage2(target: str, generator: str, emb_dim: int = 6) -> Stage2OneShot:
+    particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1}
+    if target != "physical":
+        particle_cfg = dict(particle_cfg)
+        if target == "embedding":
+            particle_cfg["type"] = "embedding"
+    k_max = 5
+    sec_dim = stage2_trunk_sec_dim({"target": target}, generator, k_max, emb_dim)
+    return 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=generator,
+        k_max=k_max,
+        particle_type_cfg={"target": target, "lambda": 1.0},
+    )
+
+
+def test_stage2_oneshot_physical_has_no_type_head_regardless_of_generator():
+    assert _build_stage2("physical", "flow").type_head is None
+    assert _build_stage2("physical", "wgan").type_head is None
+
+
+def test_stage2_oneshot_onehot_flow_has_type_head():
+    model = _build_stage2("onehot", "flow")
+    assert model.type_head is not None
+
+
+def test_stage2_oneshot_onehot_wgan_has_no_type_head():
+    """Under wgan the type slice is folded into forward()'s own output and
+    relaxed via ST-Gumbel by the trainer — no separate head needed."""
+    model = _build_stage2("onehot", "wgan")
+    assert model.type_head is None
+
+
+def test_stage2_oneshot_embedding_flow_has_type_head():
+    model = _build_stage2("embedding", "flow")
+    assert model.type_head is not None
+
+
+def test_stage2_oneshot_predict_type_shape():
+    B, k_max, emb_dim = 4, 5, 6
+    model = _build_stage2("onehot", "flow", emb_dim=emb_dim)
+    cond_cont = torch.randn(B, COND_DIM)
+    cond_cat = torch.zeros(B, 2, dtype=torch.long)
+    stage1_out = torch.randn(B, 9)
+    out = model.predict_type(cond_cont, cond_cat, stage1_out)
+    assert out.shape == (B, k_max, emb_dim)
+
+
+def test_stage2_oneshot_predict_type_raises_when_no_type_head():
+    model = _build_stage2("physical", "flow")
+    cond_cont = torch.randn(2, COND_DIM)
+    cond_cat = torch.zeros(2, 2, dtype=torch.long)
+    stage1_out = torch.randn(2, 9)
+    try:
+        model.predict_type(cond_cont, cond_cat, stage1_out)
+        raise AssertionError("expected RuntimeError")
+    except RuntimeError:
+        pass
+
+
+def test_stage2_oneshot_forward_shape_onehot_wgan():
+    B, k_max, emb_dim = 4, 5, 6
+    model = _build_stage2("onehot", "wgan", emb_dim=emb_dim)
+    cond_cont = torch.randn(B, COND_DIM)
+    cond_cat = torch.zeros(B, 2, dtype=torch.long)
+    stage1_out = torch.randn(B, 9)
+    z = torch.randn(B, model.noise_dim)
+    out = model(z, cond_cont, cond_cat, stage1_out)
+    assert out.shape == (B, k_max * (CONT_SLOT_DIM + emb_dim))
+
+
+def test_stage2_oneshot_forward_shape_onehot_flow_excludes_type():
+    B, k_max, emb_dim = 4, 5, 6
+    model = _build_stage2("onehot", "flow", emb_dim=emb_dim)
+    cond_cont = torch.randn(B, COND_DIM)
+    cond_cat = torch.zeros(B, 2, dtype=torch.long)
+    stage1_out = torch.randn(B, 9)
+    x_t = torch.randn(B, k_max * CONT_SLOT_DIM)
+    t = torch.rand(B)
+    out = model(x_t, cond_cont, cond_cat, stage1_out, t=t)
+    assert out.shape == (B, k_max * CONT_SLOT_DIM)

tests/test_pipeline.py

@@ -152,6 +152,39 @@ 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):
+    """DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
+    "onehot" (docs/v0.3.0-design.md §3.3/§8) — 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."""
+    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
+    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"}
+
+    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_no_topn_map_for_physical_target(tmp_path, data):
+    cfg = _tiny_cfg()
+    cfg["stage2_model"]["particle_type"] = {"target": "physical", "lambda": 1.0}
+    echo = _run(data, tmp_path / "out", cfg=cfg)
+    assert not any("top-N map" in m for m in echo)
+
+    loaded = setup_cache.load(data, [data])
+    assert loaded is not None
+    assert loaded.topn_maps == {}
+
+
 def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(
     tmp_path, data, monkeypatch
 ):

tests/test_setup_cache.py

@@ -7,6 +7,7 @@ import pandas as pd
 import pytest
 
 from giant.data import setup_cache
+from giant.data.loader import TopNMap
 from giant.data.setup_cache import NormalizerEntry, SetupCache
 from giant.data.transforms import Normalizer
 
@@ -101,6 +102,32 @@ def test_save_load_round_trip(tmp_path):
     np.testing.assert_allclose(entry.energy_quantiles, [1.0, 2.0, 3.0])
 
 
+def test_save_load_round_trip_topn_maps(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+
+    cache = SetupCache.empty(files)
+    cache.topn_maps[setup_cache.topn_key("pdg", 3)] = TopNMap(
+        class_map={22: 0, 11: 1, 2212: 2}, other_members={2212: 5}
+    )
+    cache.topn_maps[setup_cache.topn_key("material", 2)] = TopNMap(
+        class_map={"G4_AIR": 0, "PbWO4": 1}, other_members={}
+    )
+
+    setup_cache.save(data, files, cache)
+    loaded = setup_cache.load(data, files)
+
+    assert loaded is not None
+    pdg_m = loaded.topn_maps[setup_cache.topn_key("pdg", 3)]
+    assert pdg_m.class_map == {22: 0, 11: 1, 2212: 2}
+    assert pdg_m.other_members == {2212: 5}
+    # key type is int (matches pdg_map's own key type), not str
+    assert all(isinstance(k, int) for k in pdg_m.class_map)
+
+    mat_m = loaded.topn_maps[setup_cache.topn_key("material", 2)]
+    assert mat_m.class_map == {"G4_AIR": 0, "PbWO4": 1}
+
+
 def test_load_missing_sidecar_returns_none(tmp_path):
     data = _touch_parquet(tmp_path / "shard.parquet")
     assert setup_cache.load(data, [data]) is None

tests/test_train.py

@@ -153,7 +153,10 @@ def _fake_batches(n_batches, batch_size, seed=0):
         n_sec = torch.randint(0, K_MAX, (batch_size,), generator=g)
         sec_cont = torch.randn(batch_size, K_MAX, SEC_SLOT_DIM, generator=g)
         proc_idx = torch.zeros(batch_size, dtype=torch.long)
-        batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx))
+        sec_type_idx = torch.zeros(batch_size, K_MAX, dtype=torch.long)
+        batches.append(
+            (cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)
+        )
     return batches
 
 
@@ -234,6 +237,42 @@ def _run_train(cfg, out_dir, resume_path=None):
                 },
             ),
         ),
+        (
+            "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}
+                ),
+            ),
+        ),
     ],
 )
 def test_train_end_to_end(label, mutate):

tests/test_transforms.py

@@ -319,7 +319,7 @@ def test_build_features_clamps_n_sec_label_to_k_max():
     pdg_map = {11: 0}
     mat_map = {"PbWO4": 0}
 
-    _, _, _, n_sec, _, _, _, _ = build_features(data, pdg_map, mat_map)
+    _, _, _, n_sec, _, _, _, _, _ = build_features(data, pdg_map, mat_map)
 
     assert n_sec.max() <= K_MAX
     np.testing.assert_array_equal(n_sec, [0, 5, K_MAX])
@@ -361,7 +361,7 @@ def test_build_features_proc_idx_zero_without_proc_map():
     )
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
-    *_, proc_idx, _, _ = build_features(data, pdg_map, mat_map)
+    *_, proc_idx, _, _, _ = build_features(data, pdg_map, mat_map)
 
     np.testing.assert_array_equal(proc_idx, [0, 0, 0])
 
@@ -373,7 +373,7 @@ def test_build_features_proc_idx_looks_up_proc_map():
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
     proc_map = {"compt": 0, "phot": 1, "eIoni": 2}
 
-    *_, proc_idx, _, _ = build_features(data, pdg_map, mat_map, proc_map=proc_map)
+    *_, proc_idx, _, _, _ = build_features(data, pdg_map, mat_map, proc_map=proc_map)
 
     np.testing.assert_array_equal(proc_idx, [0, 1, 2])