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/cli.py

@@ -86,6 +86,41 @@ def _conditioning_str(model_cfg: dict, default: str = "embedding") -> str:
     return raw
 
 
+def _check_v030_onehot_support(model_cfg: dict, command: str) -> None:
+    """`giant predict`/`giant rollout` don't yet support conditioning
+    `"onehot"` mode or `stage2_model.particle_type.target` in `("onehot",
+    "embedding")` — full decode (class index -> concrete PDG, `other_policy`
+    sampling) is v0.3.0 step 6 (docs/v0.3.0-design.md), which also lands the
+    autoregressive decoder these targets are meant to pair with. Without
+    this guard, predict/rollout would crash later on a `cond_cat`/`sec_dim`
+    shape mismatch (onehot) or silently read a meaningless raw (log_mass,
+    charge) pair (embedding) instead of failing clearly. A v0.2 (flat)
+    `model_config` never has these, so this is a no-op there.
+    """
+    conditioning = model_cfg.get("conditioning")
+    if not isinstance(conditioning, dict):
+        return
+    particle_type = conditioning.get("particle", {}).get("type")
+    material_type = conditioning.get("material", {}).get("type")
+    particle_type_target = (
+        model_cfg.get("stage2_model", {}).get("particle_type", {}).get("target")
+    )
+    if (
+        particle_type == "onehot"
+        or material_type == "onehot"
+        or particle_type_target in ("onehot", "embedding")
+    ):
+        typer.echo(
+            f"error: giant {command} does not yet support conditioning "
+            "onehot mode or stage2_model.particle_type.target in "
+            "('onehot', 'embedding') — full decode (class index -> concrete "
+            "PDG) lands in v0.3.0 step 6 alongside the autoregressive "
+            "decoder these targets are meant to pair with.",
+            err=True,
+        )
+        raise typer.Exit(1)
+
+
 def _batch_size_estimate_dims(
     model_cfg: dict, training: bool, stage: str = "stage1"
 ) -> tuple[int, int]:
@@ -930,6 +965,7 @@ def predict(
         raise typer.Exit(1)
 
     model_cfg = ckpt["model_config"]
+    _check_v030_onehot_support(model_cfg, "predict")
 
     if batch_size_auto:
         est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(
@@ -999,7 +1035,7 @@ def predict(
         nonlocal writer, total
 
         if coord == Coord.local:
-            cond_cont, cond_cat, target_raw, _, _, _, _, _ = build_features(
+            cond_cont, cond_cat, target_raw, _, _, _, _, _, _ = build_features(
                 piece, pdg_map, mat_map, conditioning=conditioning
             )
             cond_cont = cond_norm.transform(cond_cont)
@@ -1342,6 +1378,7 @@ def rollout(
     training_cfg = gconfig.load_checkpoint_config(checkpoint)
 
     model_cfg = ckpt["model_config"]
+    _check_v030_onehot_support(model_cfg, "rollout")
     conditioning = _conditioning_str(model_cfg)
     pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
     mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}

giant/data/dataset.py

@@ -39,17 +39,24 @@ class StreamingStepsDataset(IterableDataset):
     numpy slicing instead of a per-row Python loop in the default collate.
 
     Each batch is a tuple:
-        (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx)
+        (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, sec_type_idx)
     where:
         cond_cont:  (B, COND_DIM) float32
-        cond_cat:   (B, 2) int64
+        cond_cat:   (B, 2/3/4) int64 — width 2 unless conditioning="onehot"
+                    (see docs/v0.3.0-design.md decision 4)
         target_s1:  (B, 9) float32 — normalised Stage-1 primary target
         n_sec:      (B,) int64 — true secondary count per step
         sec_cont:   (B, K_MAX, SEC_SLOT_DIM) float32 — [stick_logit,
                     local_dir, log_mass, charge] per slot (mass/charge
-                    normalised iff `sec_phys_normalizer` was given)
+                    normalised iff `sec_phys_normalizer` was given); always
+                    computed the same way regardless of
+                    stage2_model.particle_type.target, only actually used
+                    downstream under target="physical"
         proc_idx:   (B,) int64 — process-class label (ProcessRouter supervision
                     only; zeros when `proc_map` is None)
+        sec_type_idx: (B, K_MAX) int64 — per-slot class index into
+                    `sec_type_class_map`, for particle_type.target in
+                    ("onehot", "embedding"); zeros (unused) otherwise
     """
 
     def __init__(
@@ -66,6 +73,9 @@ class StreamingStepsDataset(IterableDataset):
         proc_map: dict[str, int] | None = None,
         conditioning: str = "embedding",
         sec_phys_normalizer: Normalizer | None = None,
+        pdg_topn_map: dict[int, int] | None = None,
+        mat_topn_map: dict[str, int] | None = None,
+        sec_type_class_map: dict | None = None,
     ) -> None:
         self.files = list(files)
         self._offsets = {path: event_id_offset(i) for i, path in enumerate(self.files)}
@@ -81,6 +91,9 @@ class StreamingStepsDataset(IterableDataset):
         self.proc_map = proc_map
         self.conditioning = conditioning
         self.sec_phys_normalizer = sec_phys_normalizer
+        self.pdg_topn_map = pdg_topn_map
+        self.mat_topn_map = mat_topn_map
+        self.sec_type_class_map = sec_type_class_map
 
     def __iter__(self):
         worker_info = torch.utils.data.get_worker_info()
@@ -98,6 +111,7 @@ class StreamingStepsDataset(IterableDataset):
         buf_nsec: list[np.ndarray] = []
         buf_sec: list[np.ndarray] = []
         buf_proc: list[np.ndarray] = []
+        buf_type: list[np.ndarray] = []
         buf_n = 0
 
         for path in files:
@@ -114,6 +128,7 @@ class StreamingStepsDataset(IterableDataset):
                     n_sec,
                     sec_cont,
                     proc_idx,
+                    sec_type_idx,
                     _,
                     _,
                 ) = build_features(
@@ -126,6 +141,9 @@ class StreamingStepsDataset(IterableDataset):
                     proc_map=self.proc_map,
                     require_secondaries=True,
                     conditioning=self.conditioning,
+                    pdg_topn_map=self.pdg_topn_map,
+                    mat_topn_map=self.mat_topn_map,
+                    sec_type_class_map=self.sec_type_class_map,
                 )
                 buf_cont.append(cond_cont)
                 buf_cat.append(cond_cat)
@@ -133,6 +151,7 @@ class StreamingStepsDataset(IterableDataset):
                 buf_nsec.append(n_sec)
                 buf_sec.append(sec_cont)
                 buf_proc.append(proc_idx)
+                buf_type.append(sec_type_idx)
                 buf_n += len(cond_cont)
 
                 if buf_n >= self.shuffle_buffer:
@@ -143,6 +162,7 @@ class StreamingStepsDataset(IterableDataset):
                         buf_nsec,
                         buf_sec,
                         buf_proc,
+                        buf_type,
                         buf_n,
                     ) = yield from self._flush(
                         buf_cont,
@@ -151,6 +171,7 @@ class StreamingStepsDataset(IterableDataset):
                         buf_nsec,
                         buf_sec,
                         buf_proc,
+                        buf_type,
                         final=False,
                     )
 
@@ -162,6 +183,7 @@ class StreamingStepsDataset(IterableDataset):
                 buf_nsec,
                 buf_sec,
                 buf_proc,
+                buf_type,
                 final=True,
             )
 
@@ -173,6 +195,7 @@ class StreamingStepsDataset(IterableDataset):
         buf_nsec: list[np.ndarray],
         buf_sec: list[np.ndarray],
         buf_proc: list[np.ndarray],
+        buf_type: list[np.ndarray],
         final: bool,
     ):
         cont = np.concatenate(buf_cont)
@@ -181,11 +204,12 @@ class StreamingStepsDataset(IterableDataset):
         nsec = np.concatenate(buf_nsec)
         sec = np.concatenate(buf_sec)
         proc = np.concatenate(buf_proc)
+        styp = np.concatenate(buf_type)
 
         if self.shuffle:
             idx = np.random.permutation(len(cont))
             cont, cat, tgt = cont[idx], cat[idx], tgt[idx]
-            nsec, sec, proc = nsec[idx], sec[idx], proc[idx]
+            nsec, sec, proc, styp = nsec[idx], sec[idx], proc[idx], styp[idx]
 
         bs = self.batch_size
         n = len(cont)
@@ -199,10 +223,11 @@ class StreamingStepsDataset(IterableDataset):
                 torch.from_numpy(nsec[start:end]).long(),
                 torch.from_numpy(sec[start:end]).float(),
                 torch.from_numpy(proc[start:end]).long(),
+                torch.from_numpy(styp[start:end]).long(),
             )
 
         if final:
-            return [], [], [], [], [], [], 0
+            return [], [], [], [], [], [], [], 0
         rem = n_full * bs
         return (
             [cont[rem:]],
@@ -211,5 +236,6 @@ class StreamingStepsDataset(IterableDataset):
             [nsec[rem:]],
             [sec[rem:]],
             [proc[rem:]],
+            [styp[rem:]],
             n - rem,
         )

giant/data/loader.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Iterator
 
@@ -243,6 +244,45 @@ def build_index_maps_from_files(
     )
 
 
+def _accumulate_value_counts(counts: dict, series: pd.Series, cast) -> None:
+    for name, count in series.value_counts().items():
+        name = cast(name)
+        counts[name] = counts.get(name, 0) + int(count)
+
+
+def _rank_by_frequency_from_files(files: list[Path], column: str, cast) -> dict:
+    """Scan `column` across `files` and return `{cast(value): total_count}`,
+    accumulated in file order (see `fingerprint_files`'s docstring on why
+    scan order — not a normalized/sorted order — is preserved: it drives
+    tie-breaking in the frequency ranking below)."""
+    counts: dict = {}
+    for path in files:
+        df = pd.read_parquet(path, columns=[column])
+        _accumulate_value_counts(counts, df[column], cast)
+    return counts
+
+
+def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict]:
+    """Frequency-capped value->index map: the `n_classes - 1` most frequent
+    keys get their own index; every rarer key is bucketed into a shared
+    "other" index (`n_classes - 1`).
+
+    Returns `(class_map, other_members)` — `other_members` is `{key: count}`
+    for every key bucketed into "other" (the empirical within-bucket
+    distribution, for `other_policy = "sample"` at rollout — see
+    docs/v0.3.0-design.md §8).
+    """
+    ranked = sorted(counts, key=lambda k: counts[k], reverse=True)
+    keep = ranked[: max(n_classes - 1, 0)]
+    class_map = {k: i for i, k in enumerate(keep)}
+    other_idx = n_classes - 1
+    other_members: dict = {}
+    for k in ranked[len(keep) :]:
+        class_map[k] = other_idx
+        other_members[k] = counts[k]
+    return class_map, other_members
+
+
 def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str, int]:
     """Scan the `process` column and build a frequency-capped process->index map.
 
@@ -253,16 +293,68 @@ def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str,
     mirrors how `build_features` clamps the n_sec label to K_MAX for the
     fixed-width n_sec_head classifier.
     """
-    counts: dict[str, int] = {}
+    counts = _rank_by_frequency_from_files(files, "process", str)
+    class_map, _ = _topn_plus_other_map(counts, n_experts)
+    return class_map
+
+
+@dataclass
+class TopNMap:
+    """A frequency-capped value->index map for a conditioning/type axis (PDG
+    or material), plus the empirical within-bucket distribution of whatever
+    got folded into "other" — see `build_topn_map_from_files`."""
+
+    class_map: dict
+    other_members: dict
+
+
+def build_topn_map_from_files(
+    files: list[Path], column: str, n_classes: int, cast=str
+) -> TopNMap:
+    """Scan `column` and build a frequency-capped value->index map, structurally
+    identical to `build_process_map_from_files` (shares its ranking core via
+    `_topn_plus_other_map`), generalized over the source column and key type.
+
+    Used for the material axis (`column="material"`, `cast=str`, matching
+    `mat_map`'s key type) — see docs/v0.3.0-design.md §8. The PDG axis uses
+    `build_pdg_topn_map_from_files` instead (it needs to pool two columns,
+    which this single-column form can't express). Also records
+    `other_members` (the empirical within-"other" distribution), needed
+    later for `other_policy = "sample"` at rollout — computed now since it's
+    free during this same scan.
+    """
+    counts = _rank_by_frequency_from_files(files, column, cast)
+    class_map, other_members = _topn_plus_other_map(counts, n_classes)
+    return TopNMap(class_map=class_map, other_members=other_members)
+
+
+def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap:
+    """PDG top-N-plus-other map, pooling counts from BOTH roles a PDG code
+    plays in this dataset: a step's own primary particle (`pdg` column) and
+    an emitted secondary's species (`sec_pdg_list`, exploded) — shared by
+    `conditioning.particle.type = "onehot"` and
+    `stage2_model.particle_type.target = "onehot"` (docs/v0.3.0-design.md
+    §8). Pooling both is what keeps a species that's common as a secondary
+    but rare as a primary (or vice versa) from being pushed into "other"
+    just because one role's count alone looks small — the meeting's failure
+    mode (zero photon secondaries, hallucinated antineutrinos) was
+    specifically about secondary-species collapse, so the map this feeds
+    needs to reflect secondary frequency, not just primary frequency.
+
+    `sec_pdg_list` is absent from parquet files predating the parent->child
+    join (see `_df_to_dict`'s `has_sec_lists` guard) — silently skipped for
+    those, same convention as elsewhere in this module.
+    """
+    counts: dict = {}
     for path in files:
-        df = pd.read_parquet(path, columns=["process"])
-        for name, count in df["process"].value_counts().items():
-            name = str(name)
-            counts[name] = counts.get(name, 0) + int(count)
-    ranked = sorted(counts, key=lambda name: counts[name], reverse=True)
-    keep = ranked[: max(n_experts - 1, 0)]
-    proc_map = {name: i for i, name in enumerate(keep)}
-    other_idx = n_experts - 1
-    for name in ranked[len(keep) :]:
-        proc_map[name] = other_idx
-    return proc_map
+        columns = ["pdg"]
+        has_sec = "sec_pdg_list" in pq.ParquetFile(path).schema_arrow.names
+        if has_sec:
+            columns.append("sec_pdg_list")
+        df = pd.read_parquet(path, columns=columns)
+        _accumulate_value_counts(counts, df["pdg"], int)
+        if has_sec:
+            exploded = df["sec_pdg_list"].explode().dropna()
+            _accumulate_value_counts(counts, exploded, int)
+    class_map, other_members = _topn_plus_other_map(counts, n_classes)
+    return TopNMap(class_map=class_map, other_members=other_members)

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},
         )
 
 

giant/data/transforms.py

@@ -340,21 +340,22 @@ def sorted_membership(values: np.ndarray, sorted_arr: np.ndarray) -> np.ndarray:
 
 
 def _vectorized_map_lookup(
-    values: np.ndarray, mapping: dict, strict: bool = True
+    values: np.ndarray, mapping: dict, strict: bool = True, default: int = 0
 ) -> np.ndarray:
     """Vectorized equivalent of `np.array([mapping[v] for v in values], dtype=np.int64)`.
 
     Replaces a per-element Python dict lookup with one `searchsorted` call.
     Raises `KeyError` if any value in `values` isn't a key of `mapping`,
     matching the dict-comprehension it replaces (never silently misassigns)
-    — unless `strict=False`, in which case unmapped values get a dummy index
-    of 0 instead. Only pass `strict=False` where the caller has independently
-    verified the resulting index is never actually read (e.g.
+    — unless `strict=False`, in which case unmapped values get `default`
+    instead. Only pass `strict=False` where the caller has independently
+    verified the resulting index is either never actually read (e.g.
     `build_cond_features` under `conditioning="physical"`, where
-    `ConditionEncoder` ignores `cond_cat` entirely); it exists so a rollout
-    can be seeded with a species/material outside the training vocab without
-    a spurious `KeyError`, which is the entire point of physical-property
-    conditioning.
+    `ConditionEncoder` ignores `cond_cat` entirely) or where `default` is a
+    deliberate fallback class (e.g. a top-N map's "other" index for a raw
+    value outside the training vocab). It exists so a rollout can be seeded
+    with a species/material outside the training vocab without a spurious
+    `KeyError`, which is the entire point of physical-property conditioning.
     """
     keys = np.asarray(list(mapping.keys()))
     vals = np.asarray(list(mapping.values()), dtype=np.int64)
@@ -366,7 +367,7 @@ def _vectorized_map_lookup(
     found = keys_sorted[pos] == values
     if not found.all():
         if not strict:
-            out = np.zeros(values.shape, dtype=np.int64)
+            out = np.full(values.shape, default, dtype=np.int64)
             out[found] = vals_sorted[pos[found]]
             return out
         missing = np.unique(values[~found])
@@ -550,6 +551,42 @@ def encode_secondaries(
     return sec_cont.astype(np.float32)
 
 
+def encode_secondary_type_idx(
+    sec_pdg_list: np.ndarray, sec_valid: np.ndarray, class_map: dict
+) -> np.ndarray:
+    """Per-secondary-slot class index into `class_map` — (N, K_MAX) int64.
+
+    `class_map` is either a top-N-plus-other map's `class_map`
+    (`stage2_model.particle_type.target = "onehot"`, see
+    `giant.data.loader.build_pdg_topn_map_from_files`) or the dense `pdg_map`
+    (`target = "embedding"`). Not used at all for `target = "physical"`
+    (see docs/v0.3.0-design.md decision 1) — that target keeps using
+    `encode_secondaries`'s (log_mass, charge) columns unchanged.
+
+    Padding slots get index 0 (their looked-up value is discarded downstream
+    by the `sec_valid`/`n_sec` mask regardless, so any in-vocabulary dummy
+    code works). A *real, valid* secondary whose code is missing from
+    `class_map` raises `KeyError` (`strict=True`) rather than silently
+    misassigning — for `target="onehot"` this should never actually
+    trigger, since `build_pdg_topn_map_from_files` pools both primary and
+    secondary occurrences precisely so every secondary species seen in
+    these files has a key (in "other" at worst); for `target="embedding"`
+    (which reuses the dense, primary-only `pdg_map`) it's a real signal
+    that a secondary-only species exists with no primary-role counterpart.
+    """
+    N, K = sec_pdg_list.shape
+    # An arbitrary already-present key works as the padding-slot dummy code
+    # (unlike encode_secondaries' physics-derived phys lookup, this is an
+    # index into class_map's own vocabulary, so a fixed sentinel like 22
+    # isn't guaranteed to be a key — an arbitrary present one always is).
+    dummy = next(iter(class_map))
+    safe_pdg = np.where(sec_valid, sec_pdg_list, dummy)
+    idx = _vectorized_map_lookup(safe_pdg.reshape(-1), class_map, strict=True).reshape(
+        N, K
+    )
+    return np.where(sec_valid, idx, 0).astype(np.int64)
+
+
 def decode_secondaries(
     sec_cont: np.ndarray,
     n_sec: np.ndarray,
@@ -643,19 +680,20 @@ def _physical_cond_columns(
 ) -> np.ndarray:
     """(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns.
 
-    "embedding" mode zero-fills (cheap, and ConditionEncoder never reads
-    these columns in that mode — so an unfilled giant.materials table can
-    never crash an "embedding"-mode run). "physical" mode computes them for
-    real: particle columns come from `data["mass"]`/`data["charge"]` when the
-    caller already knows them directly (rollout.py, for a track descended
-    from a model-predicted secondary — see giant/rollout.py's "no snapping"
-    design), else derived from `data["pdg"]` via giant.particles; material
-    columns always come from `data["material"]` via giant.materials, since
-    material is never itself a model prediction.
+    "embedding"/"onehot" modes zero-fill (cheap, and ConditionEncoder never
+    reads these columns in either mode — so an unfilled giant.materials table
+    can never crash an "embedding"/"onehot"-mode run). "physical" mode
+    computes them for real: particle columns come from
+    `data["mass"]`/`data["charge"]` when the caller already knows them
+    directly (rollout.py, for a track descended from a model-predicted
+    secondary — see giant/rollout.py's "no snapping" design), else derived
+    from `data["pdg"]` via giant.particles; material columns always come
+    from `data["material"]` via giant.materials, since material is never
+    itself a model prediction.
     """
     from giant.constants import MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
 
-    if conditioning == "embedding":
+    if conditioning in ("embedding", "onehot"):
         n = len(next(iter(data.values())))
         return np.zeros((n, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM), dtype=np.float32)
     if conditioning != "physical":
@@ -693,8 +731,19 @@ def build_cond_features(
     mat_map: dict[str, int],
     cond_normalizer: "Normalizer | None" = None,
     conditioning: str = "embedding",
+    pdg_topn_map: dict[int, int] | None = None,
+    mat_topn_map: dict[str, int] | None = None,
 ) -> tuple[np.ndarray, np.ndarray]:
-    """Build conditioning arrays only — no target, no post-step variables."""
+    """Build conditioning arrays only — no target, no post-step variables.
+
+    `pdg_topn_map`/`mat_topn_map` (a top-N-plus-other `class_map`, see
+    `giant.data.loader.build_topn_map_from_files`) append extra `cond_cat`
+    columns read by `ConditionEncoder`'s `"onehot"` mode
+    (docs/v0.3.0-design.md decision 4): pdg topN index at column 2 (iff
+    `pdg_topn_map` given), material topN index at column 3 (iff
+    `mat_topn_map` given, after column 2 if both are). Only ever given when
+    `conditioning == "onehot"`; `cond_cat` stays `(N, 2)` otherwise.
+    """
     cond_cont = np.column_stack(
         [
             data["pre_pos"],
@@ -707,16 +756,24 @@ def build_cond_features(
         [cond_cont, _physical_cond_columns(data, conditioning)]
     ).astype(np.float32)
 
-    # In "physical" mode cond_cat is only a reporting/router convenience —
-    # ConditionEncoder never reads it (giant/model/network.py) — so a
-    # species/material outside the training vocab (the whole point of
-    # physical-property conditioning) gets a dummy index instead of raising.
-    # In "embedding" mode cond_cat IS the conditioning signal, so an unmapped
-    # value must still raise loudly rather than silently misassign.
+    # In "physical" mode cond_cat's first two columns are only a
+    # reporting/router convenience — ConditionEncoder never reads them
+    # (giant/model/network.py) — so a species/material outside the training
+    # vocab (the whole point of physical-property conditioning) gets a dummy
+    # index instead of raising. In "embedding" mode those columns ARE the
+    # conditioning signal, so an unmapped value must still raise loudly
+    # rather than silently misassign. In "onehot" mode they again go unread
+    # (the topN columns below are the real signal), so they're as permissive
+    # as "physical".
     strict = conditioning == "embedding"
     pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=strict)
     mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=strict)
-    cond_cat = np.column_stack([pdg_idx, mat_idx])
+    cat_cols = [pdg_idx, mat_idx]
+    if pdg_topn_map is not None:
+        cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map))
+    if mat_topn_map is not None:
+        cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
+    cond_cat = np.column_stack(cat_cols)
 
     if cond_normalizer is not None:
         cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, conditioning)
@@ -769,6 +826,9 @@ def build_features(
     require_secondaries: bool = False,
     conditioning: str = "embedding",
     sec_phys_only: bool = False,
+    pdg_topn_map: dict[int, int] | None = None,
+    mat_topn_map: dict[str, int] | None = None,
+    sec_type_class_map: dict | None = None,
 ) -> tuple[
     np.ndarray,
     np.ndarray,
@@ -776,10 +836,12 @@ def build_features(
     np.ndarray,
     np.ndarray,
     np.ndarray,
+    np.ndarray,
     Normalizer | None,
     Normalizer | None,
 ]:
-    """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx) arrays.
+    """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx,
+    sec_type_idx) arrays.
 
     target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1)
     n_sec: (N,) integer secondary counts (target for n_sec head)
@@ -787,9 +849,18 @@ def build_features(
               [stick_logit, dir_local, log_mass, charge] — mass/charge are
               the secondary's real physical identity (from its ground-truth
               PDG code), a fixed regression target, not a learned/snapped one.
+              Always computed the same way regardless of
+              `stage2_model.particle_type.target` (docs/v0.3.0-design.md
+              decision 1/3) — only actually used downstream under `target =
+              "physical"`.
     proc_idx: (N,) integer process-class label (ProcessRouter supervision only —
               never conditioning). Zeros when `proc_map` is None or the loaded
               data has no "process" column (e.g. pre-conversion parquet files).
+    sec_type_idx: (N, K_MAX) integer secondary class index into
+              `sec_type_class_map`, for `stage2_model.particle_type.target`
+              in `("onehot", "embedding")` — see `encode_secondary_type_idx`.
+              Zero-filled (and unused) when `sec_type_class_map` is None
+              (i.e. `target = "physical"`).
 
     require_secondaries: when True, raise if any step has n_sec > 0 but the
         per-secondary list columns are absent (a mis-converted file that would
@@ -800,6 +871,14 @@ def build_features(
         the stick-breaking/direction-rotation blocks of `sec_cont` (zero-filled
         instead) for callers (normalizer fitting) that only read
         `sec_cont[:, :, 4:6]` and would otherwise discard that work.
+
+    pdg_topn_map/mat_topn_map: appended `cond_cat` columns for
+        `ConditionEncoder`'s `"onehot"` mode — see `build_cond_features`.
+
+    sec_type_class_map: the map `sec_type_idx` is looked up against — a
+        top-N-plus-other map's `class_map` for `target = "onehot"`, or the
+        dense `pdg_map` for `target = "embedding"` (pass `pdg_map` itself).
+        `None` for `target = "physical"`.
     """
     from giant.constants import K_MAX
 
@@ -836,7 +915,12 @@ def build_features(
 
     pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map)
     mat_idx = _vectorized_map_lookup(data["material"], mat_map)
-    cond_cat = np.column_stack([pdg_idx, mat_idx])  # (N, 2)
+    cat_cols = [pdg_idx, mat_idx]
+    if pdg_topn_map is not None:
+        cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map))
+    if mat_topn_map is not None:
+        cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
+    cond_cat = np.column_stack(cat_cols)  # (N, 2/3/4) — see decision 4
 
     n_sec_raw = data["n_sec"].astype(
         np.int64
@@ -864,6 +948,11 @@ def build_features(
             sec_pdg_list=sec_pdg_list,
             phys_only=sec_phys_only,
         )  # (N, K_MAX, 6)
+        sec_type_idx = (
+            encode_secondary_type_idx(sec_pdg_list, sec_valid, sec_type_class_map)
+            if sec_type_class_map is not None
+            else np.zeros((len(n_sec), K_MAX), dtype=np.int64)
+        )
     else:
         # Guard against silently training Stage 2 on zeroed targets: if any step
         # actually spawned secondaries (n_sec > 0, from child_track_ids) but the
@@ -886,6 +975,7 @@ def build_features(
             )
         N = len(n_sec)
         sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
+        sec_type_idx = np.zeros((N, K_MAX), dtype=np.int64)
 
     if fit:
         cond_normalizer = Normalizer().fit(cond_cont)
@@ -914,6 +1004,7 @@ def build_features(
         n_sec,
         sec_cont,
         proc_idx,
+        sec_type_idx,
         cond_normalizer,
         target_normalizer,
     )

giant/model/network.py

@@ -10,6 +10,7 @@ import torch.nn.functional as F
 from giant.constants import (
     COND_DIM,
     COND_DIM_BASE,
+    CONT_SLOT_DIM,
     EMB_DIM,
     K_MAX,
     MATERIAL_PHYS_DIM,
@@ -42,6 +43,30 @@ class SinusoidalEmbedding(nn.Module):
         return torch.cat([args.sin(), args.cos()], dim=-1)  # (B, dim)
 
 
+def cat_col_layout(
+    particle_type: str, material_type: str
+) -> tuple[int | None, int | None]:
+    """`cond_cat` column indices for each axis's top-N-onehot index, or
+    `None` if that axis isn't `"onehot"` (docs/v0.3.0-design.md decision 4).
+
+    Columns 0/1 are always the dense pdg/material vocab index. The particle
+    top-N column (if any) comes next, then the material top-N column (if
+    any) — `giant.data.transforms.build_cond_features`/`build_features`
+    append columns in this same order, so the two sides must never drift
+    apart.
+    """
+    col = 2
+    particle_col = None
+    if particle_type == "onehot":
+        particle_col = col
+        col += 1
+    material_col = None
+    if material_type == "onehot":
+        material_col = col
+        col += 1
+    return particle_col, material_col
+
+
 def _make_axis_mlp(in_dim: int, emb_dim: int, n_layers: int) -> nn.Sequential:
     """`n_layers`-deep MLP producing an `emb_dim`-wide vector from `in_dim`
     physical properties (`conditioning.{particle,material}.n_layers`).
@@ -76,8 +101,11 @@ class ConditionEncoder(nn.Module):
         properties (already present in `cond_cont[:, COND_DIM_BASE:]` — see
         giant.data.transforms.build_features), computable for any PDG code /
         material name rather than only ones seen in training.
-      - "onehot": not yet implemented (v0.3.0 step 4 — the top-N map isn't
-        built yet); raises `NotImplementedError` if selected.
+      - "onehot": a fixed, unlearned one-hot vector over a top-N-plus-other
+        class map (`giant.data.loader.build_topn_map_from_files`/
+        `build_pdg_topn_map_from_files`), read from `cond_cat`'s extra
+        top-N-index column(s) — see `_cat_col_layout` and
+        docs/v0.3.0-design.md decision 4.
     """
 
     def __init__(
@@ -92,6 +120,9 @@ class ConditionEncoder(nn.Module):
         super().__init__()
         self.particle_cfg = dict(particle_cfg)
         self.material_cfg = dict(material_cfg)
+        self._particle_topn_col, self._material_topn_col = cat_col_layout(
+            particle_cfg["type"], material_cfg["type"]
+        )
 
         p_type = particle_cfg["type"]
         p_emb_dim = particle_cfg["emb_dim"]
@@ -131,10 +162,11 @@ class ConditionEncoder(nn.Module):
                 :, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM
             ]
             return self.particle_mlp(particle_phys)
-        raise NotImplementedError(
-            "conditioning.particle.type='onehot' needs the top-N PDG map "
-            "(v0.3.0 step 4, not yet implemented)"
-        )
+        assert self._particle_topn_col is not None
+        return F.one_hot(
+            cond_cat[:, self._particle_topn_col],
+            num_classes=self.particle_cfg["emb_dim"],
+        ).float()
 
     def _material_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor):
         m_type = self.material_cfg["type"]
@@ -143,10 +175,11 @@ class ConditionEncoder(nn.Module):
         if m_type == "physical":
             material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :]
             return self.material_mlp(material_phys)
-        raise NotImplementedError(
-            "conditioning.material.type='onehot' needs the top-N material "
-            "map (v0.3.0 step 4, not yet implemented)"
-        )
+        assert self._material_topn_col is not None
+        return F.one_hot(
+            cond_cat[:, self._material_topn_col],
+            num_classes=self.material_cfg["emb_dim"],
+        ).float()
 
     def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         pdg_e = self._particle_embed(cond_cont, cond_cat)
@@ -754,6 +787,40 @@ def build_trunk(
 # ---------------------------------------------------------------------------
 
 
+def stage2_type_dim(particle_type_cfg: dict, emb_dim: int) -> int:
+    """Width of a single secondary slot's type slice —
+    `PARTICLE_PHYS_DIM` (log_mass, charge) for `target = "physical"`, else
+    `emb_dim` (both `"onehot"` class logits and `"embedding"` vectors are
+    `conditioning.particle.emb_dim` wide — docs/v0.3.0-design.md §3.3)."""
+    target = particle_type_cfg.get("target", "physical")
+    return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
+
+
+def stage2_trunk_sec_dim(
+    particle_type_cfg: dict, generator: str, k_max: int, emb_dim: int
+) -> int:
+    """`Stage2OneShot`'s trunk output width (docs/v0.3.0-design.md decision 2).
+
+    `target = "physical"` is untouched from v0.2/today (decision 1):
+    `k_max * SEC_SLOT_DIM`, the type slice folded into the same
+    flow-matched/WGAN vector as the continuous stick/dir slots.
+
+    `target` in `("onehot", "embedding")`: under `generator == "wgan"` the
+    type slice is still folded in (adversarial for onehot via ST-Gumbel,
+    already-continuous for embedding — §2.1), just `emb_dim` wide instead of
+    `PARTICLE_PHYS_DIM` wide: `k_max * (CONT_SLOT_DIM + emb_dim)`. Under
+    `generator in ("flow", "ddpm")` the type slice isn't part of this vector
+    at all — it's `Stage2OneShot.type_head`'s job instead — so the trunk
+    only covers `k_max * CONT_SLOT_DIM`.
+    """
+    target = particle_type_cfg.get("target", "physical")
+    if target == "physical":
+        return k_max * SEC_SLOT_DIM
+    if generator == "wgan":
+        return k_max * (CONT_SLOT_DIM + emb_dim)
+    return k_max * CONT_SLOT_DIM
+
+
 class Stage1Model(nn.Module):
     """Predicts the 9D primary post-step vector. No `n_sec_head` — decision 1
     (docs/v0.3.0-design.md §2) moves it to stage 2, except for a migrated
@@ -839,6 +906,19 @@ class Stage2OneShot(nn.Module):
     Owns `n_sec_head` by default (decision 1) unless `build_n_sec_head=False`
     (a migrated v0.2 checkpoint, whose n_sec_head instead attaches to
     Stage1Model — see `_migrate_legacy_model_config`).
+
+    `particle_type_cfg["target"]` (default `"physical"`) selects the
+    secondary-type mechanism (docs/v0.3.0-design.md decision 2):
+    `"physical"` keeps the type slice folded into the trunk's own
+    flow-matched/WGAN output, unchanged from v0.2 (`sec_dim` — computed by
+    the caller via `stage2_trunk_sec_dim` — already reflects this). Under
+    `"onehot"`/`"embedding"` with `generator in ("flow", "ddpm")`, the type
+    slice is predicted by a separate `type_head` instead (same shape pattern
+    as `n_sec_head`) — `sec_dim` then covers only the continuous
+    stick/dir slots, `type_head` covers `k_max * emb_dim` type logits/vectors.
+    Under `generator == "wgan"` the type slice stays folded into `sec_dim`
+    (just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is
+    unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation.
     """
 
     def __init__(
@@ -860,10 +940,12 @@ class Stage2OneShot(nn.Module):
         k_max: int = K_MAX,
         router: Router | None = None,
         build_n_sec_head: bool = True,
+        particle_type_cfg: dict | None = None,
     ) -> None:
         super().__init__()
         self.generator_kind = generator
         self.noise_dim = noise_dim
+        self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"})
         self.cond_enc = ConditionEncoder(
             pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
         )
@@ -886,6 +968,17 @@ class Stage2OneShot(nn.Module):
                 nn.SiLU(),
                 nn.Linear(hidden_dim // 2, k_max + 1),
             )
+        self.type_head = None
+        target = self.particle_type_cfg.get("target", "physical")
+        if target != "physical" and generator in ("flow", "ddpm"):
+            emb_dim = particle_cfg["emb_dim"]
+            self.type_head = nn.Sequential(
+                nn.Linear(cond_out_dim, hidden_dim // 2),
+                nn.SiLU(),
+                nn.Linear(hidden_dim // 2, k_max * emb_dim),
+            )
+            self._type_k_max = k_max
+            self._type_emb_dim = emb_dim
 
     def _cond_embed(
         self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor
@@ -925,6 +1018,27 @@ class Stage2OneShot(nn.Module):
         c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
         return self.n_sec_head(c_emb)
 
+    def predict_type(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+    ) -> torch.Tensor:
+        """`(B, k_max, emb_dim)` per-slot type logits (`target="onehot"`) or
+        vectors (`target="embedding"`) — only under `generator in ("flow",
+        "ddpm")`; `generator == "wgan"` folds the type slice into `forward`'s
+        own output instead (see class docstring)."""
+        if self.type_head is None:
+            raise RuntimeError(
+                "this Stage2OneShot has no type_head — either "
+                "particle_type.target='physical' (the type slice is part of "
+                "forward()'s own output) or generator='wgan' (the WGAN "
+                "trainer reads the type slice out of forward()'s output "
+                "directly instead)"
+            )
+        c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
+        return self.type_head(c_emb).view(-1, self._type_k_max, self._type_emb_dim)
+
 
 class Stage2Autoregressive(nn.Module):
     """Not implemented until v0.3.0 steps 4-7 (docs/v0.3.0-design.md §6, §12)
@@ -1201,6 +1315,10 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
         gen_sub = s2cfg.get(generator, {}) or {}
         legacy_owner = (s2cfg.get("n_sec") or {}).get("legacy_owner")
         k_max = s2cfg.get("k_max", K_MAX)
+        particle_type_cfg = s2cfg.get("particle_type") or {"target": "physical"}
+        sec_dim = stage2_trunk_sec_dim(
+            particle_type_cfg, generator, k_max, particle_cfg["emb_dim"]
+        )
         result["stage2"] = Stage2OneShot(
             pdg_vocab=pdg_vocab,
             mat_vocab=mat_vocab,
@@ -1210,7 +1328,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
             n_res_blocks=s2cfg.get("n_res_blocks", 6),
             cond_out_dim=cond_out_dim,
             context_dim=s2cfg.get("context_dim", 64),
-            sec_dim=k_max * SEC_SLOT_DIM,
+            sec_dim=sec_dim,
             dropout=s2cfg.get("dropout", 0.0),
             generator=generator,
             time_dim=gen_sub.get("time_dim", 64),
@@ -1218,6 +1336,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
             k_max=k_max,
             router=stage2_router,
             build_n_sec_head=legacy_owner != "stage1",
+            particle_type_cfg=particle_type_cfg,
         )
 
     return result
@@ -1264,12 +1383,16 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
         and s2cfg.get("decoder", "one_shot") != "autoregressive"
     ):
         k_max = s2cfg.get("k_max", K_MAX)
+        particle_type_cfg = s2cfg.get("particle_type") or {"target": "physical"}
+        in_dim = stage2_trunk_sec_dim(
+            particle_type_cfg, "wgan", k_max, particle_cfg["emb_dim"]
+        )
         result["stage2"] = CriticModel(
             pdg_vocab=pdg_vocab,
             mat_vocab=mat_vocab,
             particle_cfg=particle_cfg,
             material_cfg=material_cfg,
-            in_dim=k_max * SEC_SLOT_DIM,
+            in_dim=in_dim,
             hidden_dim=s2cfg.get("hidden_dim", 256),
             n_res_blocks=s2cfg.get("n_res_blocks", 6),
             cond_out_dim=cond_out_dim,

giant/model/schedule.py

@@ -78,25 +78,40 @@ def flow_matching_loss_secondary(
     cond_cat: torch.Tensor,
     stage1_out: torch.Tensor,
     sec_mask: torch.Tensor,
+    type_dim: int | None = None,
 ) -> torch.Tensor:
     """Flow matching loss for the secondary decoder with per-slot masking.
 
-    x1: (B, SEC_DIM) — flattened secondary target (stick_logit, dir, log_mass, charge)
+    x1: (B, K_MAX * (CONT_SLOT_DIM + type_dim)) — flattened secondary target
+        (stick_logit, dir, then a `type_dim`-wide type slice)
     sec_mask: (B, K_MAX) bool — True for valid secondary slots
+    type_dim: width of the per-slot type slice folded into `x1` — defaults to
+        `PARTICLE_PHYS_DIM` (log_mass, charge), `particle_type.target =
+        "physical"`'s width and the only case this function handled before
+        v0.3.0 step 4. `0` means no type slice is in `x1` at all (`target`
+        in `("onehot", "embedding")` under `generator in ("flow", "ddpm")` —
+        see docs/v0.3.0-design.md decision 2, `Stage2OneShot.type_head`
+        handles the type loss separately in that case).
 
     Only valid-slot dimensions contribute to the loss; padded slots are zeroed
     before averaging, so the loss is not diluted by empty slots.
 
     Each slot packs CONT_SLOT_DIM continuous dims (stick_logit, dir) followed
-    by PARTICLE_PHYS_DIM physical-identity dims (log_mass, charge) — the
-    secondary's predicted physical identity, a fixed regression target (see
-    giant.data.transforms.encode_secondaries). Even though the two blocks are
-    the same order of magnitude now (unlike the 16-wide learned embedding
-    block this replaced), they're still on different physical scales, so
-    they're each averaged over their own width first and then combined with
-    equal weight — this stays correct if PARTICLE_PHYS_DIM/CONT_SLOT_DIM change.
+    by the `type_dim`-wide type slice — under `target = "physical"` (the
+    default) that's the secondary's predicted physical identity, a fixed
+    regression target (see giant.data.transforms.encode_secondaries); under
+    `target = "embedding"` (folded in only for `generator = "wgan"`, so
+    `type_dim > 0` here only ever means "physical") it would be the detached
+    embedding-table row. Even though the two blocks are the same order of
+    magnitude now (unlike the 16-wide learned embedding block "physical"
+    replaced), they're still on different physical scales, so they're each
+    averaged over their own width first and then combined with equal weight
+    — this stays correct if type_dim/CONT_SLOT_DIM change.
     """
-    from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
+    from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM
+
+    if type_dim is None:
+        type_dim = PARTICLE_PHYS_DIM
 
     B = x1.size(0)
     t = torch.rand(B, device=x1.device)
@@ -105,12 +120,15 @@ def flow_matching_loss_secondary(
     u_t = x1 - x0
     v_t = model(x_t, cond_cont, cond_cat, stage1_out, t=t)
 
-    err = ((v_t - u_t) ** 2).view(B, K_MAX, SEC_SLOT_DIM)
+    slot_dim = CONT_SLOT_DIM + type_dim
+    err = ((v_t - u_t) ** 2).view(B, K_MAX, slot_dim)
     cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1)  # (B, K_MAX)
-    phys_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM].mean(dim=-1)
 
     mask = sec_mask.float()
     denom = mask.sum().clamp(min=1)
     cont_loss = (cont_err * mask).sum() / denom
+    if type_dim == 0:
+        return cont_loss
+    phys_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + type_dim].mean(dim=-1)
     phys_loss = (phys_err * mask).sum() / denom
     return cont_loss + phys_loss

giant/pipeline.py

@@ -15,11 +15,14 @@ from giant.constants import (
 )
 from giant.data import setup_cache
 from giant.data.loader import (
+    TopNMap,
     event_id_offset,
     find_parquet_files,
     iter_file_chunks,
     build_index_maps_from_files,
+    build_pdg_topn_map_from_files,
     build_process_map_from_files,
+    build_topn_map_from_files,
 )
 from giant.data.transforms import (
     Normalizer,
@@ -46,6 +49,8 @@ class SetupStageResult:
     pdg_map: dict[int, int]
     mat_map: dict[str, int]
     proc_map: dict[str, int] | None
+    pdg_topn_map: TopNMap | None
+    mat_topn_map: TopNMap | None
     cond_norm: Normalizer
     tgt_norm: Normalizer
     sec_phys_norm: Normalizer
@@ -193,6 +198,58 @@ def run_setup_stage(
             if cache is not None:
                 cache.proc_maps[n_experts] = proc_map
 
+    # Top-N-plus-other maps for onehot conditioning/type axes
+    # (docs/v0.3.0-design.md §8). The PDG axis is shared by
+    # conditioning.particle.type="onehot" and
+    # stage2_model.particle_type.target="onehot" (both key off
+    # conditioning.particle.emb_dim), so at most one PDG scan is needed even
+    # if both consumers are active. The material axis is independent.
+    particle_cfg = cfg["conditioning"]["particle"]
+    material_cfg = cfg["conditioning"]["material"]
+    particle_type_target = cfg["stage2_model"].get("particle_type", {}).get("target")
+
+    pdg_topn_map: TopNMap | None = None
+    if particle_cfg["type"] == "onehot" or particle_type_target == "onehot":
+        n_classes = particle_cfg["emb_dim"]
+        cache_key = setup_cache.topn_key("pdg", n_classes)
+        cached = cache.topn_maps.get(cache_key) if cache is not None else None
+        if cached is not None:
+            pdg_topn_map = cached
+            echo(
+                f"pdg top-N map: cache hit ({len(pdg_topn_map.class_map)} codes, "
+                f"{n_classes} classes)"
+            )
+        else:
+            echo("building pdg top-N map …")
+            pdg_topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
+            echo(
+                f"  {len(pdg_topn_map.class_map)} pdg codes mapped to {n_classes} classes"
+            )
+            if cache is not None:
+                cache.topn_maps[cache_key] = pdg_topn_map
+
+    mat_topn_map: TopNMap | None = None
+    if material_cfg["type"] == "onehot":
+        n_classes = material_cfg["emb_dim"]
+        cache_key = setup_cache.topn_key("material", n_classes)
+        cached = cache.topn_maps.get(cache_key) if cache is not None else None
+        if cached is not None:
+            mat_topn_map = cached
+            echo(
+                f"material top-N map: cache hit ({len(mat_topn_map.class_map)} "
+                f"materials, {n_classes} classes)"
+            )
+        else:
+            echo("building material top-N map …")
+            mat_topn_map = build_topn_map_from_files(
+                files, "material", n_classes=n_classes, cast=str
+            )
+            echo(
+                f"  {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes"
+            )
+            if cache is not None:
+                cache.topn_maps[cache_key] = mat_topn_map
+
     energy_router_active = any(
         r.get("enabled") and r.get("type") == "energy"
         for r in (stage1_router, stage2_router)
@@ -232,14 +289,16 @@ def run_setup_stage(
                 if not mask.any():
                     continue
                 chunk_tr = {k: v[mask] for k, v in chunk.items()}
-                cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _ = build_features(
-                    chunk_tr,
-                    pdg_map,
-                    mat_map,
-                    proc_map=proc_map,
-                    require_secondaries=True,
-                    conditioning=conditioning,
-                    sec_phys_only=True,
+                cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _, _ = (
+                    build_features(
+                        chunk_tr,
+                        pdg_map,
+                        mat_map,
+                        proc_map=proc_map,
+                        require_secondaries=True,
+                        conditioning=conditioning,
+                        sec_phys_only=True,
+                    )
                 )
                 cond_acc.update(cond_cont)
                 tgt_acc.update(target_s1)
@@ -273,6 +332,8 @@ def run_setup_stage(
         pdg_map=pdg_map,
         mat_map=mat_map,
         proc_map=proc_map,
+        pdg_topn_map=pdg_topn_map,
+        mat_topn_map=mat_topn_map,
         cond_norm=cond_norm,
         tgt_norm=tgt_norm,
         sec_phys_norm=sec_phys_norm,
@@ -337,6 +398,35 @@ def run_train_job(
         setup.n_train_steps,
     )
 
+    # cond_cat's onehot columns (docs/v0.3.0-design.md decision 4) are
+    # present only under conditioning="onehot" — validate_config enforces
+    # conditioning.particle.type == conditioning.material.type, so both
+    # axes' topN columns are always present or absent together. run_setup_stage
+    # builds both maps whenever conditioning=="onehot" (see its own
+    # particle_cfg["type"] == "onehot" check), so they're guaranteed non-None
+    # here — asserted, not just assumed, so a future wiring bug fails loudly
+    # instead of silently dropping the onehot columns.
+    cond_pdg_topn = None
+    cond_mat_topn = None
+    if conditioning == "onehot":
+        assert setup.pdg_topn_map is not None and setup.mat_topn_map is not None
+        cond_pdg_topn = setup.pdg_topn_map.class_map
+        cond_mat_topn = setup.mat_topn_map.class_map
+
+    # The secondary type-index map depends on stage2_model.particle_type.target,
+    # independently of conditioning's own onehot/embedding choice above
+    # (docs/v0.3.0-design.md §3.3 — physical stays untouched/None).
+    particle_type_target = (
+        cfg["stage2_model"].get("particle_type", {}).get("target", "physical")
+    )
+    if particle_type_target == "onehot":
+        assert setup.pdg_topn_map is not None
+        sec_type_class_map = setup.pdg_topn_map.class_map
+    elif particle_type_target == "embedding":
+        sec_type_class_map = pdg_map
+    else:
+        sec_type_class_map = None
+
     total_train_batches = n_train_steps // t["batch_size"]
     echo(f"  ~{n_train_steps:,} train steps, ~{total_train_batches:,} batches")
 
@@ -353,6 +443,9 @@ def run_train_job(
         proc_map=proc_map,
         conditioning=conditioning,
         sec_phys_normalizer=sec_phys_norm,
+        pdg_topn_map=cond_pdg_topn,
+        mat_topn_map=cond_mat_topn,
+        sec_type_class_map=sec_type_class_map,
     )
     val_ds = StreamingStepsDataset(
         files=files,
@@ -366,6 +459,9 @@ def run_train_job(
         proc_map=proc_map,
         conditioning=conditioning,
         sec_phys_normalizer=sec_phys_norm,
+        pdg_topn_map=cond_pdg_topn,
+        mat_topn_map=cond_mat_topn,
+        sec_type_class_map=sec_type_class_map,
     )
 
     pin = device.type == "cuda"
@@ -424,6 +520,8 @@ def run_train_job(
         pdg_map={str(k): v for k, v in pdg_map.items()},
         mat_map={str(k): v for k, v in mat_map.items()},
         proc_map=proc_map,
+        pdg_topn_map=setup.pdg_topn_map,
+        mat_topn_map=setup.mat_topn_map,
         model_config=model_config,
         resume_path=resume,
         total_train_batches=total_train_batches,

giant/train.py

@@ -16,8 +16,10 @@ import torch.optim as optim
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-from giant.constants import K_MAX
-from giant.model.network import Router
+from giant.constants import CONT_SLOT_DIM, K_MAX
+from giant.data.loader import TopNMap
+from giant.data.setup_cache import topnmap_to_json
+from giant.model.network import Router, stage2_type_dim
 from giant.model.schedule import (
     CosineSchedule,
     flow_matching_loss,
@@ -100,6 +102,61 @@ def _batch_to_device(batch: tuple, device: torch.device) -> tuple:
     return tuple(t.to(device) for t in batch)
 
 
+def _assemble_stage2_real(
+    sec_cont: torch.Tensor,
+    sec_type_idx: torch.Tensor,
+    particle_type_cfg: dict,
+    generator: str,
+    cond_enc: torch.nn.Module,
+    emb_dim: int,
+) -> torch.Tensor:
+    """Ground-truth flattened stage-2 vector, matching whatever width
+    `Stage2OneShot`'s own trunk produces for this (target, generator)
+    combination (`giant.model.network.stage2_trunk_sec_dim`;
+    docs/v0.3.0-design.md decision 2/3):
+
+    - `target = "physical"`: unchanged from v0.2 — `sec_cont` (stick_logit,
+      dir, log_mass, charge) flattened as-is.
+    - `target` in `("onehot", "embedding")` + `generator in ("flow", "ddpm")`:
+      just the continuous stick/dir slots — the type slice isn't part of
+      this vector at all (`Stage2OneShot.type_head` handles it separately).
+    - `target` in `("onehot", "embedding")` + `generator == "wgan"`: stick/dir
+      slots concatenated with the per-slot type vector — a one-hot of the
+      true class (`"onehot"`, relaxed on the *generated* side only, by the
+      caller) or the conditioning's own detached embedding-table row
+      (`"embedding"`, already continuous — no relaxation needed either side,
+      §2.1).
+    """
+    target = particle_type_cfg.get("target", "physical")
+    if target == "physical":
+        return sec_cont.flatten(1)
+    cont = sec_cont[..., :CONT_SLOT_DIM]
+    if generator != "wgan":
+        return cont.flatten(1)
+    if target == "onehot":
+        type_vec = F.one_hot(sec_type_idx, num_classes=emb_dim).float()
+    else:
+        type_vec = cond_enc.pdg_emb(sec_type_idx).detach()
+    return torch.cat([cont, type_vec], dim=-1).flatten(1)
+
+
+def _relax_onehot_type_slice(
+    x_flat: torch.Tensor, k_max: int, cont_dim: int, type_dim: int, tau: float
+) -> torch.Tensor:
+    """Straight-through Gumbel-softmax relaxation of the per-slot type slice
+    inside a flattened `(B, k_max * (cont_dim + type_dim))` WGAN generator
+    output — decision 5 (docs/v0.3.0-design.md §2.1): the forward pass is a
+    hard one-hot (matching what the critic sees from real data), the
+    backward pass flows smooth gradient. Continuous slots (stick/dir, and
+    the type slice itself under `target = "embedding"`, which never calls
+    this) pass through unchanged."""
+    B = x_flat.size(0)
+    x = x_flat.view(B, k_max, cont_dim + type_dim)
+    cont, type_logits = x[..., :cont_dim], x[..., cont_dim:]
+    type_soft = F.gumbel_softmax(type_logits, tau=tau, hard=True, dim=-1)
+    return torch.cat([cont, type_soft], dim=-1).reshape(B, -1)
+
+
 class StageTrainer:
     """One active stage's optimizer(s), EMA, and per-batch step.
 
@@ -175,6 +232,8 @@ class FlowDDPMStageTrainer(StageTrainer):
         steps_per_epoch: int,
         ddpm_n_steps: int,
         device: torch.device,
+        particle_type_cfg: dict | None = None,
+        particle_type_emb_dim: int = 16,
     ) -> None:
         if is_stage2 and generator not in ("flow",):
             raise NotImplementedError(
@@ -197,6 +256,20 @@ class FlowDDPMStageTrainer(StageTrainer):
         self.ema_decay = ema_decay
         self.router = _stage_router(self.model)
 
+        self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"})
+        self.particle_type_emb_dim = particle_type_emb_dim
+        self.particle_type_lambda = self.particle_type_cfg.get("lambda", 1.0)
+        # Width of the type slice actually folded into x1_s2 by
+        # _assemble_stage2_real, under this trainer's generator (flow/ddpm
+        # only — see the NotImplementedError above): "physical" keeps it
+        # folded in (PARTICLE_PHYS_DIM wide, unchanged from v0.2); "onehot"/
+        # "embedding" pull it out into model.type_head instead (0 here).
+        self._flow_type_dim = (
+            None
+            if self.particle_type_cfg.get("target", "physical") == "physical"
+            else 0
+        )
+
         self.params = list(self.model.parameters())
         self.optimizer = optim.AdamW(self.params, lr=lr, weight_decay=weight_decay)
         warmup_steps = warmup_epochs * steps_per_epoch
@@ -236,16 +309,69 @@ class FlowDDPMStageTrainer(StageTrainer):
             assert self.ddpm_schedule is not None
             return self.ddpm_schedule.loss(self.model, x1_s1, cond_cont, cond_cat)
         return flow_matching_loss_secondary(
-            self.model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask
+            self.model,
+            x1_s2,
+            cond_cont,
+            cond_cat,
+            stage1_ctx,
+            sec_mask,
+            type_dim=self._flow_type_dim,
         )
 
+    def _type_loss(
+        self, cond_cont, cond_cat, stage1_ctx, sec_type_idx, sec_mask, device
+    ):
+        """CE (`target="onehot"`) or MSE (`target="embedding"`) loss for
+        `Stage2OneShot.type_head` — the non-adversarial counterpart to
+        WGANStageTrainer's ST-Gumbel-into-the-critic path (decision 2/5).
+        Zero when this stage has no `type_head` (stage 1, or
+        `particle_type.target = "physical"`)."""
+        l_type = torch.zeros((), device=device)
+        type_acc = torch.zeros((), device=device)
+        type_head = getattr(self.model, "type_head", None)
+        if not self.is_stage2 or type_head is None:
+            return l_type, type_acc
+        type_out = self.model.predict_type(cond_cont, cond_cat, stage1_ctx)
+        mask = sec_mask.float()
+        denom = mask.sum().clamp(min=1)
+        if self.particle_type_cfg.get("target") == "onehot":
+            ce = F.cross_entropy(
+                type_out.transpose(1, 2), sec_type_idx, reduction="none"
+            )
+            l_type = (ce * mask).sum() / denom
+            type_acc = (
+                (type_out.argmax(-1) == sec_type_idx).float() * mask
+            ).sum() / denom
+        else:  # "embedding"
+            target_vec = self.model.cond_enc.pdg_emb(sec_type_idx).detach()
+            se = ((type_out - target_vec) ** 2).mean(-1)
+            l_type = (se * mask).sum() / denom
+        return l_type, type_acc
+
     def _compute(self, batch: tuple, device: torch.device) -> dict:
-        cond_cont, cond_cat, x1_s1, n_sec, sec_cont, proc_idx = _batch_to_device(
-            batch, device
-        )
-        x1_s2 = sec_cont.flatten(1)
+        (
+            cond_cont,
+            cond_cat,
+            x1_s1,
+            n_sec,
+            sec_cont,
+            proc_idx,
+            sec_type_idx,
+        ) = _batch_to_device(batch, device)
         sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
         stage1_ctx = x1_s1.detach()
+        x1_s2 = (
+            _assemble_stage2_real(
+                sec_cont,
+                sec_type_idx,
+                self.particle_type_cfg,
+                self.generator,
+                self.model.cond_enc,
+                self.particle_type_emb_dim,
+            )
+            if self.is_stage2
+            else None
+        )
 
         l_gen = self._generator_loss(
             cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx
@@ -257,13 +383,21 @@ class FlowDDPMStageTrainer(StageTrainer):
             l_nsec = F.cross_entropy(n_sec_logits, n_sec)
             nsec_acc = (n_sec_logits.argmax(dim=-1) == n_sec).float().mean()
 
+        l_type, type_acc = self._type_loss(
+            cond_cont, cond_cat, stage1_ctx, sec_type_idx, sec_mask, device
+        )
+
         l_balance = l_proc = l_entropy = torch.zeros((), device=device)
         if self.router is not None:
             l_balance = self.router.balance_loss(cond_cont, cond_cat)
             l_proc = self.router.classify_loss(cond_cont, cond_cat, proc_idx)
             l_entropy = self.router.entropy_loss(cond_cont, cond_cat)
 
-        total = self.lambda_weight * l_gen + self.n_sec_lambda * l_nsec
+        total = (
+            self.lambda_weight * l_gen
+            + self.n_sec_lambda * l_nsec
+            + self.particle_type_lambda * l_type
+        )
         if self.lambda_balance > 0:
             total = total + self.lambda_balance * l_balance
         if self.lambda_proc > 0:
@@ -275,6 +409,8 @@ class FlowDDPMStageTrainer(StageTrainer):
             "total": total,
             "loss_gen": l_gen,
             "loss_nsec": l_nsec,
+            "loss_type": l_type,
+            "type_acc": type_acc,
             "loss_balance": l_balance,
             "loss_proc": l_proc,
             "loss_entropy": l_entropy,
@@ -302,6 +438,8 @@ class FlowDDPMStageTrainer(StageTrainer):
             "loss": out["total"].item(),
             "loss_gen": out["loss_gen"].item(),
             "loss_nsec": out["loss_nsec"].item(),
+            "loss_type": out["loss_type"].item(),
+            "type_acc": out["type_acc"].item(),
             "loss_balance": out["loss_balance"].item(),
             "loss_proc": out["loss_proc"].item(),
             "loss_entropy": out["loss_entropy"].item(),
@@ -318,6 +456,8 @@ class FlowDDPMStageTrainer(StageTrainer):
             "loss": out["total"].item(),
             "loss_gen": out["loss_gen"].item(),
             "loss_nsec": out["loss_nsec"].item(),
+            "loss_type": out["loss_type"].item(),
+            "type_acc": out["type_acc"].item(),
             "loss_balance": out["loss_balance"].item(),
             "loss_proc": out["loss_proc"].item(),
             "loss_entropy": out["loss_entropy"].item(),
@@ -387,6 +527,10 @@ class WGANStageTrainer(StageTrainer):
         epochs: int,
         steps_per_epoch: int,
         device: torch.device,
+        particle_type_cfg: dict | None = None,
+        particle_type_emb_dim: int = 16,
+        type_gumbel_tau_start: float = 1.0,
+        type_gumbel_tau_end: float = 0.1,
     ) -> None:
         self.name = name
         self.is_stage2 = is_stage2
@@ -400,6 +544,11 @@ class WGANStageTrainer(StageTrainer):
         self.ema_decay = ema_decay
         self.router = _stage_router(self.model)
 
+        self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"})
+        self.particle_type_emb_dim = particle_type_emb_dim
+        self.type_gumbel_tau_start = type_gumbel_tau_start
+        self.type_gumbel_tau_end = type_gumbel_tau_end
+
         self.g_params = list(self.model.parameters())
         self.d_params = list(self.critic.parameters())
         # WGAN-GP recipe (Gulrajani et al. 2017): Adam, beta1=0, no weight decay.
@@ -423,6 +572,7 @@ class WGANStageTrainer(StageTrainer):
 
         self._lr_lambda = _lr_lambda
         self.lr_sched = optim.lr_scheduler.LambdaLR(self.optimizer, _lr_lambda)
+        self.total_steps = total_steps
 
         self.ema_model: torch.nn.Module | None = None
         if ema_decay > 0:
@@ -431,9 +581,15 @@ class WGANStageTrainer(StageTrainer):
                 p.requires_grad_(False)
 
     def step(self, batch: tuple, device: torch.device, global_step: int) -> dict:
-        cond_cont, cond_cat, x1_s1, n_sec, sec_cont, _proc_idx = _batch_to_device(
-            batch, device
-        )
+        (
+            cond_cont,
+            cond_cat,
+            x1_s1,
+            n_sec,
+            sec_cont,
+            _proc_idx,
+            sec_type_idx,
+        ) = _batch_to_device(batch, device)
         B = cond_cont.size(0)
         stage1_ctx = x1_s1.detach()
 
@@ -447,24 +603,49 @@ class WGANStageTrainer(StageTrainer):
             fake = self.model(z, cond_cont, cond_cat)
             mask = None
         else:
-            from giant.constants import SEC_SLOT_DIM
+            target = self.particle_type_cfg.get("target", "physical")
+            type_dim = stage2_type_dim(
+                self.particle_type_cfg, self.particle_type_emb_dim
+            )
+            slot_width = CONT_SLOT_DIM + type_dim
 
             sec_mask = torch.arange(K_MAX, device=device).unsqueeze(
                 0
             ) < n_sec.unsqueeze(1)
             mask = (
-                sec_mask.unsqueeze(-1)
-                .expand(-1, -1, SEC_SLOT_DIM)
-                .reshape(B, -1)
-                .float()
+                sec_mask.unsqueeze(-1).expand(-1, -1, slot_width).reshape(B, -1).float()
+            )
+            real = (
+                _assemble_stage2_real(
+                    sec_cont,
+                    sec_type_idx,
+                    self.particle_type_cfg,
+                    "wgan",
+                    self.model.cond_enc,
+                    self.particle_type_emb_dim,
+                )
+                * mask
             )
-            real = sec_cont.flatten(1) * mask
 
             def critic_fn(x):
                 return self.critic(x, cond_cont, cond_cat, stage1_ctx)
 
             z = torch.randn(B, self.model.noise_dim, device=device)
             fake_raw = self.model(z, cond_cont, cond_cat, stage1_ctx)
+            if target == "onehot":
+                # Straight-through Gumbel-softmax relaxation of the type
+                # slice only (decision 5) — the critic must see a hard
+                # one-hot forward (matching what "real" data looks like)
+                # while gradient still flows smoothly to the generator.
+                tau = _gumbel_tau(
+                    global_step,
+                    self.total_steps,
+                    self.type_gumbel_tau_start,
+                    self.type_gumbel_tau_end,
+                )
+                fake_raw = _relax_onehot_type_slice(
+                    fake_raw, K_MAX, CONT_SLOT_DIM, type_dim, tau
+                )
             fake = fake_raw * mask
 
         # --- critic step (every batch) ---
@@ -635,6 +816,10 @@ def _build_stage_trainers(
         lambda_entropy = router_cfg.get("lambda_entropy", 0.0)
         gumbel_tau_start = router_cfg.get("gumbel_tau_start", 1.0)
         gumbel_tau_end = router_cfg.get("gumbel_tau_end", 0.1)
+        particle_type_cfg = cfg["stage2_model"].get("particle_type") or {
+            "target": "physical"
+        }
+        particle_type_emb_dim = cfg["conditioning"]["particle"]["emb_dim"]
 
         if generator == "wgan":
             critic = critics.get(name)
@@ -659,6 +844,10 @@ def _build_stage_trainers(
                 epochs=t["epochs"],
                 steps_per_epoch=max(total_train_batches, 1),
                 device=device,
+                particle_type_cfg=particle_type_cfg,
+                particle_type_emb_dim=particle_type_emb_dim,
+                type_gumbel_tau_start=wgan_cfg.get("gumbel_tau_start", 1.0),
+                type_gumbel_tau_end=wgan_cfg.get("gumbel_tau_end", 0.1),
             )
         else:
             ddpm_n_steps = stage_cfg.get("ddpm", {}).get("n_steps", 1000)
@@ -682,6 +871,8 @@ def _build_stage_trainers(
                 steps_per_epoch=max(total_train_batches, 1),
                 ddpm_n_steps=ddpm_n_steps,
                 device=device,
+                particle_type_cfg=particle_type_cfg,
+                particle_type_emb_dim=particle_type_emb_dim,
             )
     return trainers
 
@@ -710,11 +901,15 @@ def _metrics_fields(trainers: dict[str, StageTrainer]) -> list[str]:
                 f"{name}_train_loss_proc",
                 f"{name}_train_loss_entropy",
                 f"{name}_train_nsec_acc",
+                f"{name}_train_loss_type",
+                f"{name}_train_type_acc",
                 f"{name}_train_grad_norm",
                 f"{name}_val_loss",
                 f"{name}_val_loss_gen",
                 f"{name}_val_loss_nsec",
                 f"{name}_val_nsec_acc",
+                f"{name}_val_loss_type",
+                f"{name}_val_type_acc",
             ]
         if trainer.router is not None:
             fields += [
@@ -766,6 +961,8 @@ def train(
     pdg_map: dict | None = None,
     mat_map: dict | None = None,
     proc_map: dict | None = None,
+    pdg_topn_map: TopNMap | None = None,
+    mat_topn_map: TopNMap | None = None,
     model_config: dict | None = None,
     resume_path: str | Path | None = None,
     total_train_batches: int = 0,
@@ -842,6 +1039,10 @@ def train(
             ckpt["mat_map"] = mat_map
         if proc_map is not None:
             ckpt["proc_map"] = proc_map
+        if pdg_topn_map is not None:
+            ckpt["pdg_topn_map"] = topnmap_to_json(pdg_topn_map)
+        if mat_topn_map is not None:
+            ckpt["mat_topn_map"] = topnmap_to_json(mat_topn_map)
         if model_config is not None:
             ckpt["model_config"] = model_config
         return ckpt
@@ -1183,6 +1384,12 @@ def train(
                     metrics_row[f"{name}_train_nsec_acc"] = (
                         sums.get("nsec_acc", 0.0) / n_train
                     )
+                    metrics_row[f"{name}_train_loss_type"] = (
+                        sums.get("loss_type", 0.0) / n_train
+                    )
+                    metrics_row[f"{name}_train_type_acc"] = (
+                        sums.get("type_acc", 0.0) / n_train
+                    )
                     metrics_row[f"{name}_train_grad_norm"] = (
                         sums.get("grad_norm", 0.0) / n_train
                     )
@@ -1192,6 +1399,10 @@ def train(
                         v.get("loss_nsec", 0.0) / n_val
                     )
                     metrics_row[f"{name}_val_nsec_acc"] = v.get("nsec_acc", 0.0) / n_val
+                    metrics_row[f"{name}_val_loss_type"] = (
+                        v.get("loss_type", 0.0) / n_val
+                    )
+                    metrics_row[f"{name}_val_type_acc"] = v.get("type_acc", 0.0) / n_val
                     grad_norm_total += sums.get("grad_norm", 0.0) / n_train
                 if tr.router is not None:
                     rs = val_router_sums.get(name)

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])