diff --git a/giant/cli.py b/giant/cli.py index 1abf63b..99dfdba 100644 --- a/giant/cli.py +++ b/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()} diff --git a/giant/data/dataset.py b/giant/data/dataset.py index 3b56a79..3736b92 100644 --- a/giant/data/dataset.py +++ b/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, ) diff --git a/giant/data/loader.py b/giant/data/loader.py index e76af49..40f6a6b 100644 --- a/giant/data/loader.py +++ b/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) diff --git a/giant/data/setup_cache.py b/giant/data/setup_cache.py index b1c69d4..a20d0b6 100644 --- a/giant/data/setup_cache.py +++ b/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}, ) diff --git a/giant/data/transforms.py b/giant/data/transforms.py index e41041e..a605cab 100644 --- a/giant/data/transforms.py +++ b/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, ) diff --git a/giant/model/network.py b/giant/model/network.py index 88e89e9..3a52b32 100644 --- a/giant/model/network.py +++ b/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, diff --git a/giant/model/schedule.py b/giant/model/schedule.py index d9f437a..0a5bf3d 100644 --- a/giant/model/schedule.py +++ b/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 diff --git a/giant/pipeline.py b/giant/pipeline.py index dfb2111..5584996 100644 --- a/giant/pipeline.py +++ b/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, diff --git a/giant/train.py b/giant/train.py index 2e1688c..53403c9 100644 --- a/giant/train.py +++ b/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) diff --git a/tests/test_cli_predict.py b/tests/test_cli_predict.py index a6f74fb..8bae93f 100644 --- a/tests/test_cli_predict.py +++ b/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 # --------------------------------------------------------------------------- diff --git a/tests/test_loader.py b/tests/test_loader.py index f0900c7..23003a8 100644 --- a/tests/test_loader.py +++ b/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) ──────────────────────────────────────────── diff --git a/tests/test_network.py b/tests/test_network.py index 0bb99b5..2117898 100644 --- a/tests/test_network.py +++ b/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) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 833b8fd..f49c38d 100644 --- a/tests/test_pipeline.py +++ b/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 ): diff --git a/tests/test_setup_cache.py b/tests/test_setup_cache.py index a243a53..9224b64 100644 --- a/tests/test_setup_cache.py +++ b/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 diff --git a/tests/test_train.py b/tests/test_train.py index 5892de7..e20223d 100644 --- a/tests/test_train.py +++ b/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): diff --git a/tests/test_transforms.py b/tests/test_transforms.py index dcda475..a2e19c2 100644 --- a/tests/test_transforms.py +++ b/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])