diff --git a/giant/cli.py b/giant/cli.py index 42f8ad0..d38cdc0 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -472,7 +472,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 ) cond_cont = cond_norm.transform(cond_cont) diff --git a/giant/config.py b/giant/config.py index b5008a6..d0ba445 100644 --- a/giant/config.py +++ b/giant/config.py @@ -37,6 +37,12 @@ DEFAULT_CONFIG: dict = { "temperature": 0.5, # energy-router kwarg "learn_centers": True, # energy-router kwarg "lambda_balance": 0.0, # optional load-balance aux loss weight + "emb_dim": 8, # process-router kwarg: its own pdg/mat embedding width + "hidden_dim": 64, # process-router kwarg: its classifier's hidden width + "lambda_proc": 0.0, # process-router kwarg: supervised process-CE weight + # (0.0 still trains a working router — the gate gets gradient + # through the downstream flow loss like EnergyRouter's centers — + # but only lambda_proc > 0 grounds it in the true `process` label) }, }, } diff --git a/giant/data/dataset.py b/giant/data/dataset.py index 433c5f3..ce05905 100644 --- a/giant/data/dataset.py +++ b/giant/data/dataset.py @@ -36,7 +36,7 @@ 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, sec_pdg_idx) + (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx) where: cond_cont: (B, COND_DIM) float32 cond_cat: (B, 2) int64 @@ -44,6 +44,8 @@ class StreamingStepsDataset(IterableDataset): n_sec: (B,) int64 — true secondary count per step sec_cont: (B, K_MAX, 4) float32 — [stick_logit, local_dir] per slot sec_pdg_idx: (B, K_MAX) int64 — PDG model-index per secondary slot + proc_idx: (B,) int64 — process-class label (ProcessRouter supervision + only; zeros when `proc_map` is None) """ def __init__( @@ -57,6 +59,7 @@ class StreamingStepsDataset(IterableDataset): batch_size: int, shuffle_buffer: int = 65536, shuffle: bool = True, + proc_map: dict[str, int] | None = None, ) -> None: self.files = list(files) self.split_events = split_events @@ -68,6 +71,7 @@ class StreamingStepsDataset(IterableDataset): self.batch_size = batch_size self.shuffle_buffer = max(shuffle_buffer, batch_size) self.shuffle = shuffle + self.proc_map = proc_map def __iter__(self): worker_info = torch.utils.data.get_worker_info() @@ -85,6 +89,7 @@ class StreamingStepsDataset(IterableDataset): buf_nsec: list[np.ndarray] = [] buf_sec: list[np.ndarray] = [] buf_spdg: list[np.ndarray] = [] + buf_proc: list[np.ndarray] = [] buf_n = 0 for path in files: @@ -94,14 +99,16 @@ class StreamingStepsDataset(IterableDataset): continue chunk = {k: v[mask] for k, v in chunk.items()} - cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, _, _ = ( - build_features( - chunk, - self.pdg_map, - self.mat_map, - cond_normalizer=self.cond_normalizer, - target_normalizer=self.target_normalizer, - ) + ( + cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, + proc_idx, _, _, + ) = build_features( + chunk, + self.pdg_map, + self.mat_map, + cond_normalizer=self.cond_normalizer, + target_normalizer=self.target_normalizer, + proc_map=self.proc_map, ) buf_cont.append(cond_cont) buf_cat.append(cond_cat) @@ -109,19 +116,24 @@ class StreamingStepsDataset(IterableDataset): buf_nsec.append(n_sec) buf_sec.append(sec_cont) buf_spdg.append(sec_pdg_idx) + buf_proc.append(proc_idx) buf_n += len(cond_cont) if buf_n >= self.shuffle_buffer: - buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, buf_n = ( + ( + buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, + buf_proc, buf_n, + ) = ( yield from self._flush( buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, - final=False, + buf_proc, final=False, ) ) if buf_n > 0: yield from self._flush( - buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, final=True + buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, buf_proc, + final=True, ) def _flush( @@ -132,6 +144,7 @@ class StreamingStepsDataset(IterableDataset): buf_nsec: list[np.ndarray], buf_sec: list[np.ndarray], buf_spdg: list[np.ndarray], + buf_proc: list[np.ndarray], final: bool, ): cont = np.concatenate(buf_cont) @@ -140,11 +153,12 @@ class StreamingStepsDataset(IterableDataset): nsec = np.concatenate(buf_nsec) sec = np.concatenate(buf_sec) spdg = np.concatenate(buf_spdg) + proc = np.concatenate(buf_proc) if self.shuffle: idx = np.random.permutation(len(cont)) cont, cat, tgt = cont[idx], cat[idx], tgt[idx] - nsec, sec, spdg = nsec[idx], sec[idx], spdg[idx] + nsec, sec, spdg, proc = nsec[idx], sec[idx], spdg[idx], proc[idx] bs = self.batch_size n = len(cont) @@ -158,13 +172,14 @@ class StreamingStepsDataset(IterableDataset): torch.from_numpy(nsec[start:end]).long(), torch.from_numpy(sec[start:end]).float(), torch.from_numpy(spdg[start:end]).long(), + torch.from_numpy(proc[start:end]).long(), ) if final: - return [], [], [], [], [], [], 0 + return [], [], [], [], [], [], [], 0 rem = n_full * bs return ( [cont[rem:]], [cat[rem:]], [tgt[rem:]], - [nsec[rem:]], [sec[rem:]], [spdg[rem:]], + [nsec[rem:]], [sec[rem:]], [spdg[rem:]], [proc[rem:]], n - rem, ) diff --git a/giant/data/loader.py b/giant/data/loader.py index fa0ba79..a215157 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -95,6 +95,16 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]: "layer_id": df["layer_id"].to_numpy(dtype=np.int32), "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32), "e_sec": df["e_sec"].to_numpy(dtype=np.float32), + # The physics process that ended the step (e.g. "compt", "phot", + # "eBrem") — a post-step outcome, so it's a router/classifier + # supervision label only, never conditioning (see build_process_map* + # / ProcessRouter). Guarded like has_sec_lists: older parquet + # conversions predating this column still load fine. + "process": ( + df["process"].to_numpy(dtype=object) + if "process" in df.columns + else np.full(len(df), "", dtype=object) + ), "step_length": df["step_length"].to_numpy(dtype=np.float32), "post_E": df["post_E"].to_numpy(dtype=np.float32), "delta_e": (df["pre_E"] - df["post_E"]).to_numpy(dtype=np.float32), @@ -192,3 +202,30 @@ def build_index_maps_from_files( {v: i for i, v in enumerate(sorted(pdg_vals))}, {v: i for i, v in enumerate(sorted(mat_vals))}, ) + + +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. + + Physics processes have a long tail (rare nuclear captures, decays, ...) + while `ProcessRouter` needs a fixed number of expert slots, so only the + `n_experts - 1` most frequent processes get their own index; every rarer + process is bucketed into a shared "other" index (`n_experts - 1`). This + mirrors how `build_features` clamps the n_sec label to K_MAX for the + fixed-width n_sec_head classifier. + """ + counts: dict[str, int] = {} + 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 diff --git a/giant/data/transforms.py b/giant/data/transforms.py index f5df8db..e8df166 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -433,6 +433,7 @@ def build_features( cond_normalizer: Normalizer | None = None, target_normalizer: Normalizer | None = None, fit: bool = False, + proc_map: dict[str, int] | None = None, ) -> tuple[ np.ndarray, np.ndarray, @@ -440,16 +441,20 @@ 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, sec_pdg_idx) arrays. + """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_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) sec_cont: (N, K_MAX, 4) continuous secondary targets [stick_logit, dir_local] sec_pdg_idx: (N, K_MAX) integer PDG model-indices; used to look up embedding targets in the training loop + 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). """ from giant.constants import K_MAX @@ -523,6 +528,12 @@ def build_features( if target_normalizer is not None: target_s1 = target_normalizer.transform(target_s1) + process = data.get("process") + if proc_map is not None and process is not None: + proc_idx = np.array([proc_map[str(p)] for p in process], dtype=np.int64) + else: + proc_idx = np.zeros(len(cond_cat), dtype=np.int64) + return ( cond_cont, cond_cat, @@ -530,6 +541,7 @@ def build_features( n_sec, sec_cont, sec_pdg_idx, + proc_idx, cond_normalizer, target_normalizer, ) diff --git a/giant/model/network.py b/giant/model/network.py index 91c06a1..749e950 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -3,6 +3,7 @@ import math import torch import torch.nn as nn +import torch.nn.functional as F from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_DIM, X_DIM @@ -267,6 +268,19 @@ class Router(nn.Module): importance = self.gate(cond_cont, cond_cat).sum(dim=0) # (n_experts,) return (importance.std() / (importance.mean() + 1e-8)) ** 2 + def classify_loss( + self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor + ) -> torch.Tensor: + """Optional supervised auxiliary loss shaping the router's own belief. + + Default: none (a scalar 0), for routers like EnergyRouter that read a + quantity directly off cond_cont/cond_cat and need no label. Routers + gating on an unobservable pre-step quantity (e.g. ProcessRouter, + which predicts the physics process that will end the step) override + this to supervise their internal classifier against the true label. + """ + return torch.zeros((), device=cond_cont.device) + ROUTER_REGISTRY: dict[str, type[Router]] = {} @@ -330,6 +344,64 @@ class EnergyRouter(Router): return torch.softmax(-d2 / self.temperature, dim=-1) +@register_router("process") +class ProcessRouter(Router): + """Routes on the physics process expected to end the step. + + Unlike EnergyRouter (which reads a quantity that's already known at + pre-step time), the process — Compton, photoelectric, brems, ... — is a + *post-step outcome*: it can't be read off cond_cont/cond_cat directly. + Instead this router runs a small classifier over pre-step conditioning + (its own pdg/material embeddings, kept separate from the trunk's + ConditionEncoder) that predicts it, one class per expert slot + (`n_experts` doubles as the number of process classes — see + `build_process_map_from_files`, which caps the process vocabulary to + exactly this many classes, bucketing rare processes into a shared + "other" slot). + + The classifier is supervised by `classify_loss` against the true + `process` label (see `giant/train.py`) — a *training-time* signal only; + `gate`/`top1` never see it, so eval-time dispatch (rollout, predict) + needs no ground truth, same as every other Router. This sidesteps the + gradient/differentiability problem that sank the earlier + process-conditioned-flow proposal (see the archived decision doc): the + hard categorical choice only ever feeds a non-differentiable expert + *dispatch*, never the flow's own conditioning path. + """ + + def __init__( + self, + n_experts: int, + pdg_vocab: int, + mat_vocab: int, + emb_dim: int = 8, + hidden_dim: int = 64, + ) -> None: + super().__init__(n_experts) + self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim) + self.mat_emb = nn.Embedding(mat_vocab, emb_dim) + self.classifier = nn.Sequential( + nn.Linear(COND_DIM + 2 * emb_dim, hidden_dim), + nn.SiLU(), + nn.Linear(hidden_dim, n_experts), + ) + + def logits(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + """(B, n_experts) raw process-classifier logits, one class per expert.""" + pdg_e = self.pdg_emb(cond_cat[:, 0]) + mat_e = self.mat_emb(cond_cat[:, 1]) + h = torch.cat([cond_cont, pdg_e, mat_e], dim=-1) + return self.classifier(h) + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1) + + def classify_loss( + self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor + ) -> torch.Tensor: + return F.cross_entropy(self.logits(cond_cont, cond_cat), labels) + + class ExpertTrunk(nn.Module): """One small expert: `input_proj -> ResBlock stack -> out_proj`. @@ -570,6 +642,12 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts") } + # Not every router needs these (EnergyRouter doesn't declare them, so + # build_router's kwarg filtering drops them silently) but + # ProcessRouter needs its own pdg/material embeddings sized to match + # the checkpoint's vocab, same as the trunk's ConditionEncoder. + router_kwargs.setdefault("pdg_vocab", model_config["pdg_vocab"]) + router_kwargs.setdefault("mat_vocab", model_config["mat_vocab"]) shared = dict( pdg_vocab=model_config["pdg_vocab"], mat_vocab=model_config["mat_vocab"], diff --git a/giant/pipeline.py b/giant/pipeline.py index 16d655f..931c5fd 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -11,6 +11,7 @@ from giant.data.loader import ( load_event_ids, iter_file_chunks, build_index_maps_from_files, + build_process_map_from_files, ) from giant.data.transforms import build_features, _WelfordAccumulator from giant.data.dataset import make_event_split, StreamingStepsDataset @@ -54,6 +55,13 @@ def run_train_job( pdg_map, mat_map = build_index_maps_from_files(files) echo(f" {len(pdg_map)} PDG codes | {len(mat_map)} materials") + router_cfg = m["router"] + proc_map: dict[str, int] | None = None + if router_cfg.get("enabled") and router_cfg.get("type") == "process": + echo("building process vocabulary …") + proc_map = build_process_map_from_files(files, n_experts=router_cfg["n_experts"]) + echo(f" {len(proc_map)} process labels mapped to {router_cfg['n_experts']} experts") + echo("fitting normalizer (streaming) …") cond_acc = _WelfordAccumulator(COND_DIM) tgt_acc = _WelfordAccumulator(X_DIM) @@ -63,8 +71,8 @@ def run_train_job( if not mask.any(): continue chunk_tr = {k: v[mask] for k, v in chunk.items()} - cond_cont, _, target_s1, _n_sec, _sec_cont, _sec_pdg, _, _ = build_features( - chunk_tr, pdg_map, mat_map + cond_cont, _, target_s1, _n_sec, _sec_cont, _sec_pdg, _proc, _, _ = ( + build_features(chunk_tr, pdg_map, mat_map, proc_map=proc_map) ) cond_acc.update(cond_cont) tgt_acc.update(target_s1) @@ -81,6 +89,7 @@ def run_train_job( batch_size=t["batch_size"], shuffle_buffer=shuffle_buffer, shuffle=True, + proc_map=proc_map, ) val_ds = StreamingStepsDataset( files=files, @@ -91,6 +100,7 @@ def run_train_job( target_normalizer=tgt_norm, batch_size=t["batch_size"], shuffle=False, + proc_map=proc_map, ) pin = device.type == "cuda" @@ -114,7 +124,6 @@ def run_train_job( "update giant/constants.py if emb_dim changed" ) - router_cfg = m["router"] model_config = { "pdg_vocab": len(pdg_map), "mat_vocab": len(mat_map), @@ -161,9 +170,11 @@ def run_train_job( lambda_nsec=t.get("lambda_nsec", 0.1), lambda_s2=t.get("lambda_s2", 1.0), lambda_balance=router_cfg.get("lambda_balance", 0.0), + lambda_proc=router_cfg.get("lambda_proc", 0.0), normalizer_dict={"cond": cond_norm.to_dict(), "target": tgt_norm.to_dict()}, 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, model_config=model_config, resume_path=resume, validate_every=t["validate_every"], diff --git a/giant/train.py b/giant/train.py index b30a616..35f3d8c 100644 --- a/giant/train.py +++ b/giant/train.py @@ -27,11 +27,13 @@ _METRICS_FIELDS = [ "train_loss_nsec", "train_loss_s2", "train_loss_balance", + "train_loss_proc", "val_loss", "val_loss_s1", "val_loss_nsec", "val_loss_s2", "val_loss_balance", + "val_loss_proc", "lr", "epoch_time_s", ] @@ -111,15 +113,19 @@ def _compute_losses( lambda_nsec: float, lambda_s2: float, lambda_balance: float = 0.0, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance) for one batch.""" - cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_idx = batch + lambda_proc: float = 0.0, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance, L_proc) for one batch.""" + cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx = batch cond_cont = cond_cont.to(device) cond_cat = cond_cat.to(device) x1_s1 = x1_s1.to(device) n_sec = n_sec.to(device) sec_cont = sec_cont.to(device) sec_pdg_idx = sec_pdg_idx.to(device) + proc_idx = proc_idx.to(device) # Stage-1 flow loss if mode == "flow": @@ -161,13 +167,21 @@ def _compute_losses( l_balance = stage1_model.router.balance_loss( cond_cont, cond_cat ) + sec_decoder.router.balance_loss(cond_cont, cond_cat) + # Supervised router auxiliary loss (e.g. ProcessRouter's process + # classifier); a scalar 0 for routers with no such loss (EnergyRouter). + l_proc = stage1_model.router.classify_loss( + cond_cont, cond_cat, proc_idx + ) + sec_decoder.router.classify_loss(cond_cont, cond_cat, proc_idx) else: l_balance = torch.zeros((), device=device) + l_proc = torch.zeros((), device=device) total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2 if lambda_balance > 0: total = total + lambda_balance * l_balance - return total, l_s1, l_nsec, l_s2, l_balance + if lambda_proc > 0: + total = total + lambda_proc * l_proc + return total, l_s1, l_nsec, l_s2, l_balance, l_proc def train( @@ -184,9 +198,11 @@ def train( lambda_nsec: float = 0.1, lambda_s2: float = 1.0, lambda_balance: float = 0.0, + lambda_proc: float = 0.0, normalizer_dict: dict | None = None, pdg_map: dict | None = None, mat_map: dict | None = None, + proc_map: dict | None = None, model_config: dict | None = None, resume_path: str | Path | None = None, validate_every: int = 0, @@ -263,6 +279,7 @@ def train( train_nsec_sum = 0.0 train_s2_sum = 0.0 train_balance_sum = 0.0 + train_proc_sum = 0.0 train_n = 0 ema_loss = 0.0 bar = tqdm( @@ -274,7 +291,7 @@ def train( dynamic_ncols=True, ) for batch in bar: - loss, l_s1, l_nsec, l_s2, l_balance = _compute_losses( + loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses( stage1_model, sec_decoder, batch, @@ -284,6 +301,7 @@ def train( lambda_nsec, lambda_s2, lambda_balance, + lambda_proc, ) optimizer.zero_grad() loss.backward() @@ -297,6 +315,7 @@ def train( train_nsec_sum += l_nsec.item() * B train_s2_sum += l_s2.item() * B train_balance_sum += l_balance.item() * B + train_proc_sum += l_proc.item() * B train_n += B ema_loss = ( batch_loss if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss @@ -320,10 +339,11 @@ def train( val_nsec_sum = 0.0 val_s2_sum = 0.0 val_balance_sum = 0.0 + val_proc_sum = 0.0 val_n = 0 with torch.no_grad(): for batch in val_loader: - loss, l_s1, l_nsec, l_s2, l_balance = _compute_losses( + loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses( stage1_model, sec_decoder, batch, @@ -333,6 +353,7 @@ def train( lambda_nsec, lambda_s2, lambda_balance, + lambda_proc, ) B = batch[0].size(0) val_loss_sum += loss.item() * B @@ -340,6 +361,7 @@ def train( val_nsec_sum += l_nsec.item() * B val_s2_sum += l_s2.item() * B val_balance_sum += l_balance.item() * B + val_proc_sum += l_proc.item() * B val_n += B val_loss = val_loss_sum / max(val_n, 1) epoch_time = time.monotonic() - epoch_start @@ -352,7 +374,8 @@ def train( f" (s1={train_s1_sum / max(train_n, 1):.3f}" f" nsec={train_nsec_sum / max(train_n, 1):.3f}" f" s2={train_s2_sum / max(train_n, 1):.3f}" - f" bal={train_balance_sum / max(train_n, 1):.3f})" + f" bal={train_balance_sum / max(train_n, 1):.3f}" + f" proc={train_proc_sum / max(train_n, 1):.3f})" f" val {val_loss:.4f}" f" lr {current_lr:.2e} {epoch_time:.1f}s{marker}" ) @@ -364,11 +387,13 @@ def train( "train_loss_nsec": train_nsec_sum / max(train_n, 1), "train_loss_s2": train_s2_sum / max(train_n, 1), "train_loss_balance": train_balance_sum / max(train_n, 1), + "train_loss_proc": train_proc_sum / max(train_n, 1), "val_loss": val_loss, "val_loss_s1": val_s1_sum / max(val_n, 1), "val_loss_nsec": val_nsec_sum / max(val_n, 1), "val_loss_s2": val_s2_sum / max(val_n, 1), "val_loss_balance": val_balance_sum / max(val_n, 1), + "val_loss_proc": val_proc_sum / max(val_n, 1), "lr": current_lr, "epoch_time_s": epoch_time, } @@ -401,6 +426,8 @@ def train( ckpt["pdg_map"] = pdg_map if mat_map is not None: ckpt["mat_map"] = mat_map + if proc_map is not None: + ckpt["proc_map"] = proc_map if model_config is not None: ckpt["model_config"] = model_config diff --git a/giant/validate.py b/giant/validate.py index 959e375..3372a06 100644 --- a/giant/validate.py +++ b/giant/validate.py @@ -85,8 +85,8 @@ def validate_marginals( for i, batch in enumerate(val_loader): if n_batches is not None and i >= n_batches: break - # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx). - cond_cont, cond_cat, x1, n_sec, sec_cont, sec_pdg_idx = batch + # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx). + cond_cont, cond_cat, x1, n_sec, sec_cont, sec_pdg_idx, _proc_idx = batch cond_cont = cond_cont.to(device) cond_cat = cond_cat.to(device) diff --git a/tests/test_loader.py b/tests/test_loader.py index 6c562f3..703d9d3 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -1,6 +1,7 @@ +import pandas as pd import pytest -from giant.data.loader import find_parquet_files +from giant.data.loader import build_process_map_from_files, find_parquet_files def _touch(path): @@ -58,3 +59,36 @@ def test_manifest_with_no_entries_raises(tmp_path): with pytest.raises(FileNotFoundError): find_parquet_files(manifest) + + +def test_build_process_map_from_files_keeps_most_frequent(tmp_path): + """process counts: eIoni=5, phot=3, compt=2, Rayl=1 — with n_experts=3, only + the top 2 (eIoni, phot) get their own index; compt/Rayl share the "other" + (last) index.""" + process = ( + ["eIoni"] * 5 + ["phot"] * 3 + ["compt"] * 2 + ["Rayl"] * 1 + ) + path = tmp_path / "shard-000.parquet" + pd.DataFrame({"process": process}).to_parquet(path) + + proc_map = build_process_map_from_files([path], n_experts=3) + + assert proc_map["eIoni"] == 0 + assert proc_map["phot"] == 1 + assert proc_map["compt"] == 2 + assert proc_map["Rayl"] == 2 + assert set(proc_map.values()) <= {0, 1, 2} + + +def test_build_process_map_from_files_spans_multiple_files(tmp_path): + path_a = tmp_path / "a.parquet" + path_b = tmp_path / "b.parquet" + pd.DataFrame({"process": ["eIoni"] * 3 + ["phot"] * 1}).to_parquet(path_a) + pd.DataFrame({"process": ["phot"] * 4 + ["compt"] * 1}).to_parquet(path_b) + + # phot: 1+4=5 total > eIoni: 3 > compt: 1 + proc_map = build_process_map_from_files([path_a, path_b], n_experts=3) + + assert proc_map["phot"] == 0 + assert proc_map["eIoni"] == 1 + assert proc_map["compt"] == 2 diff --git a/tests/test_router.py b/tests/test_router.py index 36b7e73..0755ebb 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -6,6 +6,7 @@ from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM from giant.model.network import ( DenoisingMLP, EnergyRouter, + ProcessRouter, ROUTER_REGISTRY, RoutedDenoisingMLP, RoutedSecondaryDecoder, @@ -101,6 +102,98 @@ def test_build_router_unknown_type_raises(): raise AssertionError("expected ValueError for unknown router type") +# ── ProcessRouter ──────────────────────────────────────────────────────────── + + +def test_process_router_registered(): + assert ROUTER_REGISTRY["process"] is ProcessRouter + + +def test_process_router_gate_partition_of_unity(): + router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2) + cond_cont, cond_cat = _cond(16) + g = router.gate(cond_cont, cond_cat) + assert g.shape == (16, 4) + torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0) + + +def test_process_router_top1_matches_gate_argmax(): + router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2) + cond_cont, cond_cat = _cond(16) + assert torch.equal( + router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1) + ) + + +def test_process_router_balance_loss_is_nonnegative_scalar(): + router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2) + cond_cont, cond_cat = _cond(16) + loss = router.balance_loss(cond_cont, cond_cat) + assert loss.shape == () + assert loss.item() >= 0.0 + + +def test_process_router_classify_loss_decreases_with_training(): + """The classifier should be able to fit an arbitrary label assignment — + a sanity check that gradients actually flow to the process classifier.""" + torch.manual_seed(0) + router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2) + cond_cont, cond_cat = _cond(32) + labels = torch.randint(0, 4, (32,)) + + opt = torch.optim.Adam(router.parameters(), lr=0.05) + first = router.classify_loss(cond_cont, cond_cat, labels).item() + for _ in range(50): + opt.zero_grad() + loss = router.classify_loss(cond_cont, cond_cat, labels) + loss.backward() + opt.step() + last = loss.item() + assert last < first + + +def test_energy_router_classify_loss_defaults_to_zero(): + """Routers with no supervised signal (EnergyRouter) fall back to the + Router base class's zero-loss default.""" + router = EnergyRouter(n_experts=4) + cond_cont, cond_cat = _cond(16) + labels = torch.randint(0, 4, (16,)) + loss = router.classify_loss(cond_cont, cond_cat, labels) + assert loss.shape == () + assert loss.item() == 0.0 + + +def test_build_router_process_type_uses_pdg_mat_vocab(): + router = build_router("process", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8) + assert isinstance(router, ProcessRouter) + assert router.pdg_emb.num_embeddings == 5 + assert router.mat_emb.num_embeddings == 3 + + +def test_build_models_routed_with_process_router(): + model_config = dict( + pdg_vocab=4, + mat_vocab=2, + emb_dim=16, + dropout=0.1, + k_max=K_MAX, + expert_hidden_dim=16, + expert_n_blocks=2, + router={ + "enabled": True, + "type": "process", + "n_experts": 3, + "lambda_proc": 1.0, + }, + ) + stage1, sec_decoder = build_models(model_config) + assert isinstance(stage1, RoutedDenoisingMLP) + assert isinstance(stage1.router, ProcessRouter) + assert len(stage1.experts) == 3 + assert stage1.router.pdg_emb.num_embeddings == 4 + assert stage1.router.mat_emb.num_embeddings == 2 + + # ── RoutedDenoisingMLP ─────────────────────────────────────────────────────── diff --git a/tests/test_transforms.py b/tests/test_transforms.py index d92f347..ad8ddfb 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -235,7 +235,48 @@ 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]) + + +def _minimal_step_data(N: int, process: np.ndarray | None = None) -> dict: + rng = np.random.default_rng(0) + data = { + "pdg": np.full(N, 11, dtype=np.int32), + "material": np.full(N, "PbWO4", dtype=object), + "pre_pos": rng.standard_normal((N, 3)).astype(np.float32), + "pre_E": np.full(N, 10.0, dtype=np.float32), + "pre_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)), + "layer_id": np.zeros(N, dtype=np.int32), + "n_sec": np.zeros(N, dtype=np.int32), + "e_sec": np.full(N, 1.0, dtype=np.float32), + "step_length": np.full(N, 1.0, dtype=np.float32), + "post_E": np.full(N, 9.0, dtype=np.float32), + "edep": np.full(N, 1.0, dtype=np.float32), + "post_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)), + "post_pos": rng.standard_normal((N, 3)).astype(np.float32), + } + if process is not None: + data["process"] = process + return data + + +def test_build_features_proc_idx_zero_without_proc_map(): + data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object)) + pdg_map, mat_map = {11: 0}, {"PbWO4": 0} + + *_, proc_idx, _, _ = build_features(data, pdg_map, mat_map) + + np.testing.assert_array_equal(proc_idx, [0, 0, 0]) + + +def test_build_features_proc_idx_looks_up_proc_map(): + data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object)) + 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) + + np.testing.assert_array_equal(proc_idx, [0, 1, 2])