lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Add ProcessRouter for physics-process-based expert gating

Browse Source 
Routes on the physics process (Compton, phot, brems, ...) that ends a
step, supervised by a small classifier since process is a post-step
outcome unobservable at gate time. Threads a process label end-to-end
through the data pipeline (loader, build_features, dataset batches,
training loss/checkpointing) alongside the existing EnergyRouter.
This commit is contained in:
Lars Bogner
2026-07-08 16:15:09 +02:00
parent bac541240f
commit 0b3ece52ed
12 changed files with 385 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
giant/cli.py
+1 -1
View File
@@ -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)
giant/config.py
+6
View File
@@ -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)
},
},
}
giant/data/dataset.py
+30 -15
View File
@@ -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,
)
giant/data/loader.py
+37
View File
@@ -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
giant/data/transforms.py
+13 -1
View File
@@ -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,
)
giant/model/network.py
+78
View File
@@ -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"],
giant/pipeline.py
+14 -3
View File
@@ -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"],
giant/train.py
+34 -7
View File
@@ -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
giant/validate.py
+2 -2
View File
@@ -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)
tests/test_loader.py
+35 -1
View File
@@ -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
tests/test_router.py
+93
View File
@@ -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 ───────────────────────────────────────────────────────
tests/test_transforms.py
+42 -1
View File
@@ -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])