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

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

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

giant/data/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)