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>

giant/data/setup_cache.py

@@ -23,7 +23,7 @@ import numpy as np
 
 from giant import config
 from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM
-from giant.data.loader import event_id_offset, load_event_ids
+from giant.data.loader import TopNMap, event_id_offset, load_event_ids
 from giant.data.transforms import Normalizer, sorted_membership
 
 # Bump manually on a change to the data-encoding semantics (e.g. a future
@@ -102,6 +102,40 @@ def normalizer_key(val_fraction: float, seed: int, conditioning: str) -> str:
     return f"valfrac={val_fraction:.6g}_seed={seed}_cond={conditioning}"
 
 
+# Top-N-map axes (docs/v0.3.0-design.md §8): "pdg" keys match pdg_map's int
+# keys (shared by conditioning.particle.type="onehot" and
+# stage2_model.particle_type.target="onehot" — one map for both), "material"
+# keys match mat_map's str keys.
+_TOPN_AXIS_CASTS = {"pdg": int, "material": str}
+
+
+def topn_key(axis: str, n_classes: int) -> str:
+    """JSON-safe key for `SetupCache.topn_maps` — N is part of the key so the
+    sidecar stays reusable across runs with different emb_dim (see the
+    dict[int, dict] precedent `proc_maps` sets, keyed by n_experts)."""
+    if axis not in _TOPN_AXIS_CASTS:
+        raise ValueError(
+            f"unknown top-N map axis {axis!r}, expected one of "
+            f"{sorted(_TOPN_AXIS_CASTS)}"
+        )
+    return f"{axis}:{n_classes}"
+
+
+def topnmap_to_json(m: TopNMap) -> dict:
+    return {
+        "class_map": {str(k): v for k, v in m.class_map.items()},
+        "other_members": {str(k): v for k, v in m.other_members.items()},
+    }
+
+
+def topnmap_from_json(d: dict, axis: str) -> TopNMap:
+    cast = _TOPN_AXIS_CASTS[axis]
+    return TopNMap(
+        class_map={cast(k): v for k, v in d["class_map"].items()},
+        other_members={cast(k): v for k, v in d["other_members"].items()},
+    )
+
+
 @dataclass
 class NormalizerEntry:
     cond_norm: Normalizer
@@ -143,6 +177,8 @@ class SetupCache:
     event_index: tuple[np.ndarray, np.ndarray] | None = None
     proc_maps: dict[int, dict[str, int]] = field(default_factory=dict)
     normalizers: dict[str, NormalizerEntry] = field(default_factory=dict)
+    topn_maps: dict[str, TopNMap] = field(default_factory=dict)
+    """Keyed by `topn_key(axis, n_classes)` — see docs/v0.3.0-design.md §8."""
 
     @classmethod
     def empty(cls, files: list[Path]) -> "SetupCache":
@@ -156,6 +192,7 @@ class SetupCache:
             "fingerprint": self.fingerprint,
             "proc_maps": {str(k): v for k, v in self.proc_maps.items()},
             "normalizers": {k: v.to_json() for k, v in self.normalizers.items()},
+            "topn_maps": {k: topnmap_to_json(v) for k, v in self.topn_maps.items()},
         }
         if self.vocab is not None:
             pdg_map, mat_map = self.vocab
@@ -188,6 +225,10 @@ class SetupCache:
         normalizers = {
             k: NormalizerEntry.from_json(v) for k, v in d.get("normalizers", {}).items()
         }
+        topn_maps = {
+            k: topnmap_from_json(v, axis=k.split(":", 1)[0])
+            for k, v in d.get("topn_maps", {}).items()
+        }
         return cls(
             fingerprint=d["fingerprint"],
             git_hash=d.get("git_hash", "unknown"),
@@ -195,6 +236,7 @@ class SetupCache:
             event_index=event_index,
             proc_maps=proc_maps,
             normalizers=normalizers,
+            topn_maps=topn_maps,
         )
 
     def merge(self, other: "SetupCache") -> "SetupCache":
@@ -214,6 +256,7 @@ class SetupCache:
             ),
             proc_maps={**self.proc_maps, **other.proc_maps},
             normalizers={**self.normalizers, **other.normalizers},
+            topn_maps={**self.topn_maps, **other.topn_maps},
         )