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

Merge branch '4-prototype-a-mixture-of-experts-routing-tree-architecture' into analysis-streaming-rewrite

Browse Source 
# Conflicts:
#	giant/analysis.py
#	tests/test_analysis.py
This commit is contained in:
Lars Bogner
2026-07-17 12:06:10 +02:00
parent 55c676fb9b 5a0e98c5d1
commit 6744d05948
21 changed files with 2854 additions and 203 deletions
Download Patch File Download Diff File Expand all files Collapse all files
analysis/rollout_validation.ipynb
+454
View File
File diff suppressed because one or more lines are too long
giant/cli.py
+238 -59
View File
@@ -1,7 +1,9 @@
from collections import Counter
from datetime import date, datetime, timezone
from enum import Enum
import math
from pathlib import Path
import re
from typing import Optional
import uuid as uuid_mod
@@ -39,31 +41,91 @@ from giant.data.transforms import (
Normalizer,
)
from giant.geometry import GeometryOracle
from giant.model.network import DenoisingMLP, SecondaryDecoder
from giant.model.network import build_models
from giant.pipeline import run_train_job
from giant.rollout import rollout as run_rollout
from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx
_STAGE1_MODEL_KEYS = {
"pdg_vocab",
"mat_vocab",
"hidden_dim",
"n_blocks",
"emb_dim",
"dropout",
"k_max",
}
_SEC_DECODER_MODEL_KEYS = {
"pdg_vocab",
"mat_vocab",
"hidden_dim",
"n_blocks",
"emb_dim",
"dropout",
}
app = typer.Typer(no_args_is_help=True)
def _router_total_experts(router_cfg: dict) -> int:
"""Total expert count for a router config, single-axis or composed.
A composed router runs one expert per *joint* cell, so its count is the
product of the per-axis `axis{i}_n_experts` (mirrors
`ComposedRouter.__init__` in giant.model.network); a single-axis router
just reports its own `n_experts`.
"""
if router_cfg.get("type") == "composed":
axis_counts = {
m.group(1): int(v)
for k, v in router_cfg.items()
if (m := re.match(r"^axis(\d+)_n_experts$", k))
}
return math.prod(axis_counts.values()) if axis_counts else 1
return int(router_cfg.get("n_experts", 1))
def _batch_size_estimate_dims(model_cfg: dict, training: bool) -> tuple[int, int]:
"""Pick the (hidden_dim, n_blocks) that dominate per-call activation memory.
Routed models spend their FLOPs in the (smaller) expert trunks, not the
monolith's hidden_dim/n_blocks, so estimate_batch_size needs the expert
dims instead when routing is enabled. Training runs the full soft mixture
(every expert on the whole batch), so its activation memory scales with
the expert count; inference does top-1 dispatch (each row hits one
expert), so the batch just partitions across experts and one expert's
dims already bound it. estimate_batch_size scales memory linearly with
hidden_dim * n_blocks, so the training multiplier folds into n_blocks.
"""
router_cfg = model_cfg.get("router")
if router_cfg and router_cfg.get("enabled"):
hidden_dim = model_cfg.get("expert_hidden_dim", 128)
n_blocks = model_cfg.get("expert_n_blocks", 3)
if training:
n_blocks *= _router_total_experts(router_cfg)
return hidden_dim, n_blocks
return model_cfg["hidden_dim"], model_cfg["n_blocks"]
def _coerce_scalar(value: str) -> object:
"""Best-effort str -> bool/int/float, else leave as str.
CLI flag values always arrive as strings; router kwargs like
`n_experts` (int) or `temperature` (float) need to come out typed the
same way a TOML file's native types would, since they're merged into
the same `model.router` dict as file-sourced config.
"""
if value.lower() in ("true", "false"):
return value.lower() == "true"
try:
return int(value)
except ValueError:
pass
try:
return float(value)
except ValueError:
pass
return value
def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]:
"""Parse repeated `--router-axis "type:key=val,key=val"` flags into
`axis{i}_{field}` flat keys (see `_parse_composed_axes` in
giant.model.network), indexed by flag order — the Nth `--router-axis`
becomes axis N.
"""
out: dict[str, object] = {}
for i, spec in enumerate(specs):
axis_type, _, rest = spec.partition(":")
out[f"axis{i}_type"] = axis_type
for pair in filter(None, rest.split(",")):
key, _, val = pair.partition("=")
out[f"axis{i}_{key}"] = _coerce_scalar(val)
return out
_CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions")
@@ -121,6 +183,39 @@ class Coord(str, Enum):
local = "local"
class Weights(str, Enum):
raw = "raw"
ema = "ema"
def _load_model_weights(
model: torch.nn.Module,
sec_decoder: torch.nn.Module,
ckpt: dict,
weights: "Weights",
checkpoint_path: Path,
) -> None:
"""Load either the raw or EMA state dicts from a training checkpoint.
EMA weights (giant.train's shadow copy, see --ema-decay) only exist in
checkpoints written after that feature landed, so `ema` fails loudly
rather than silently falling back to raw weights a caller didn't ask for.
"""
if weights == Weights.raw:
model_key, sec_key = "model", "sec_decoder"
else:
model_key, sec_key = "model_ema", "sec_decoder_ema"
if model_key not in ckpt or sec_key not in ckpt:
typer.echo(
f"error: {checkpoint_path} has no EMA weights (trained before "
"--ema-decay, or with --ema-decay 0) — use --weights raw",
err=True,
)
raise typer.Exit(1)
model.load_state_dict(ckpt[model_key])
sec_decoder.load_state_dict(ckpt[sec_key])
@app.command()
def train(
data: Annotated[
@@ -147,6 +242,18 @@ def train(
),
] = None,
lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None,
weight_decay: Annotated[
Optional[float],
typer.Option("--weight-decay", "-W", help="AdamW weight decay (default: 0.01)"),
] = None,
ema_decay: Annotated[
Optional[float],
typer.Option(
"--ema-decay",
help="EMA decay for a shadow copy of the model weights, saved "
"alongside the raw weights in checkpoints (0 disables; default: 0.9999)",
),
] = None,
warmup_epochs: Annotated[
Optional[int], typer.Option("--warmup-epochs", "-w")
] = None,
@@ -159,6 +266,33 @@ def train(
"--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"
),
] = None,
router: Annotated[
Optional[bool],
typer.Option(
"--router/--no-router",
help="Route both stages through a mixture of small experts "
"instead of one monolithic trunk (see model.router in config.toml)",
),
] = None,
router_type: Annotated[
Optional[str],
typer.Option(
"--router-type", help="Router implementation name (see ROUTER_REGISTRY)"
),
] = None,
n_experts: Annotated[
Optional[int], typer.Option("--n-experts", help="Number of routed experts")
] = None,
router_axis: Annotated[
Optional[list[str]],
typer.Option(
"--router-axis",
help="Composed-router axis spec 'type:key=val,key=val' (repeatable; "
"Nth flag = axis N). Use with --router-type composed instead of "
"--n-experts, e.g. --router-axis 'energy:n_experts=4' "
"--router-axis 'pdg:n_experts=3,emb_dim=8'",
),
] = None,
val_fraction: Annotated[
Optional[float], typer.Option("--val-fraction", "-f")
] = None,
@@ -183,6 +317,14 @@ def train(
"(ignored in ddpm mode, which always runs the full schedule)",
),
] = None,
max_val_batches: Annotated[
Optional[int],
typer.Option(
"--max-val-batches",
help="Cap the per-epoch val-loss pass to N batches (0 = full "
"val set every epoch; default: 200)",
),
] = None,
shuffle_buffer: Annotated[
int,
typer.Option(
@@ -229,16 +371,19 @@ def train(
"epochs": epochs,
"batch_size": batch_size_value,
"lr": lr,
"weight_decay": weight_decay,
"ema_decay": ema_decay,
"warmup_epochs": warmup_epochs,
"val_fraction": val_fraction,
"num_workers": num_workers,
"seed": seed,
"validate_every": validate_every,
"validate_steps": validate_steps,
"max_val_batches": max_val_batches,
}.items()
if v is not None
}
cli_model = {
cli_model: dict[str, object] = {
k: v
for k, v in {
"hidden_dim": hidden_dim,
@@ -248,6 +393,19 @@ def train(
}.items()
if v is not None
}
cli_router: dict[str, object] = {
k: v
for k, v in {
"enabled": router,
"type": router_type,
"n_experts": n_experts,
}.items()
if v is not None
}
if router_axis:
cli_router.update(_parse_router_axis_flags(router_axis))
if cli_router:
cli_model["router"] = cli_router
cfg = gconfig.merge_cli_overrides(
gconfig.DEFAULT_CONFIG, config, cli_train, cli_model
)
@@ -256,9 +414,10 @@ def train(
_device = torch.device(device) if device else gconfig.auto_device()
if batch_size_auto:
est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(m, training=True)
try:
t["batch_size"] = gconfig.estimate_batch_size(
m["hidden_dim"], m["n_blocks"], _device
est_hidden_dim, est_n_blocks, _device
)
except ValueError as exc:
typer.echo(f"error: {exc}", err=True)
@@ -336,6 +495,15 @@ def predict(
steps: Annotated[
int, typer.Option("--steps", "-s", help="Flow matching ODE steps")
] = 10,
weights: Annotated[
Weights,
typer.Option(
"--weights",
help="raw: the live training weights. ema: the EMA shadow copy "
"(see --ema-decay in `giant train`) — usually cleaner samples, "
"requires a checkpoint trained with EMA enabled.",
),
] = Weights.raw,
device: Annotated[
Optional[str],
typer.Option("--device", "-d", help="cpu | cuda | mps (default: auto)"),
@@ -386,10 +554,13 @@ def predict(
model_cfg = ckpt["model_config"]
if batch_size_auto:
est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(
model_cfg, training=False
)
try:
batch_size_value = gconfig.estimate_batch_size(
model_cfg["hidden_dim"],
model_cfg["n_blocks"],
est_hidden_dim,
est_n_blocks,
_device,
training=False,
)
@@ -408,19 +579,12 @@ def predict(
cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
model = DenoisingMLP(
**{k: v for k, v in model_cfg.items() if k in _STAGE1_MODEL_KEYS}
)
model.load_state_dict(ckpt["model"])
model, sec_decoder = build_models(model_cfg)
_load_model_weights(model, sec_decoder, ckpt, weights, checkpoint)
model.to(_device).eval()
sec_decoder = SecondaryDecoder(
**{k: v for k, v in model_cfg.items() if k in _SEC_DECODER_MODEL_KEYS}
)
sec_decoder.load_state_dict(ckpt["sec_decoder"])
sec_decoder.to(_device).eval()
typer.echo(f"loaded checkpoint: {checkpoint}")
typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})")
gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
# --- Output path ---
@@ -448,7 +612,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)
@@ -709,6 +873,15 @@ def rollout(
int,
typer.Option("--steps", "-s", help="Flow matching ODE steps per model call"),
] = 10,
weights: Annotated[
Weights,
typer.Option(
"--weights",
help="raw: the live training weights. ema: the EMA shadow copy "
"(see --ema-decay in `giant train`) — usually cleaner samples, "
"requires a checkpoint trained with EMA enabled.",
),
] = Weights.raw,
batch_size: Annotated[
int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward")
] = 4096,
@@ -763,17 +936,11 @@ def rollout(
cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
model = DenoisingMLP(
**{k: v for k, v in model_cfg.items() if k in _STAGE1_MODEL_KEYS}
)
model.load_state_dict(ckpt["model"])
model, sec_decoder = build_models(model_cfg)
_load_model_weights(model, sec_decoder, ckpt, weights, checkpoint)
model.to(_device).eval()
sec_decoder = SecondaryDecoder(
**{k: v for k, v in model_cfg.items() if k in _SEC_DECODER_MODEL_KEYS}
)
sec_decoder.load_state_dict(ckpt["sec_decoder"])
sec_decoder.to(_device).eval()
typer.echo(f"loaded checkpoint: {checkpoint}")
typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})")
oracle = GeometryOracle.load(geometry)
typer.echo(
@@ -785,7 +952,29 @@ def rollout(
seeds = _seed_from_data(files, n_events)
typer.echo(f"seeded {len(seeds['event_id']):,} shower(s)")
records = run_rollout(
out, dataset_path, pred_uuid = _resolve_prediction_output(data, out)
out.parent.mkdir(parents=True, exist_ok=True)
# Written incrementally as each batch of steps is produced, rather than
# buffering the whole run (which scales with n_events * max_steps *
# avg_tracks_per_event) — mirrors the row-group streaming `giant predict`
# already does on its input side.
writer: pq.ParquetWriter | None = None
def _write_chunk(row: dict[str, np.ndarray]) -> None:
nonlocal writer
table = pa.table(row)
if writer is None:
table = table.replace_schema_metadata(
{
PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE,
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
}
)
writer = pq.ParquetWriter(out, table.schema)
writer.write_table(table)
summary = run_rollout(
model,
sec_decoder,
oracle,
@@ -801,18 +990,10 @@ def rollout(
device=_device,
max_tracks_per_event=max_tracks_per_event,
escape_threshold=escape_threshold,
on_chunk=_write_chunk,
)
out, dataset_path, pred_uuid = _resolve_prediction_output(data, out)
out.parent.mkdir(parents=True, exist_ok=True)
table = pa.table(records).replace_schema_metadata(
{
PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE,
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
}
)
pq.write_table(table, out)
if writer is not None:
writer.close()
ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path)
ref = yaml.safe_load(ref_path.read_text())
@@ -829,10 +1010,8 @@ def rollout(
)
ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False))
n_rows = len(records["event_id"])
reasons = Counter(r for r in records["termination_reason"].tolist() if r)
typer.echo(f"wrote {n_rows:,} step rows → {out}")
typer.echo(f"terminations: {dict(reasons)}")
typer.echo(f"wrote {summary['n_rows']:,} step rows → {out}")
typer.echo(f"terminations: {summary['termination_reason_counts']}")
typer.echo(f"reference: {ref_path}")
giant/config.py
+68 -5
View File
@@ -14,6 +14,11 @@ DEFAULT_CONFIG: dict = {
"epochs": 100,
"batch_size": 4096,
"lr": 3e-4,
"weight_decay": 0.01, # AdamW default — exposed so it can be tuned
"ema_decay": 0.9999, # EMA of model weights for sampling; 0 disables
# per-epoch val loss (not the marginal/KL validate_every pass) is
# capped to this many batches; 0 = full val set every epoch
"max_val_batches": 200,
"val_fraction": 0.1,
"num_workers": 4,
"seed": 0,
@@ -28,6 +33,29 @@ DEFAULT_CONFIG: dict = {
"n_blocks": 6,
"emb_dim": 16,
"dropout": 0.1,
"router": {
"enabled": False,
"type": "energy", # selects the Router impl from ROUTER_REGISTRY
"n_experts": 4,
"expert_hidden_dim": 128,
"expert_n_blocks": 3,
"temperature": 0.5, # energy/pdg-router kwarg
"learn_centers": True, # energy/pdg-router kwarg
"lambda_balance": 0.0, # optional load-balance aux loss weight
"emb_dim": 8, # process/pdg-router kwarg: 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)
# type = "composed" routes on multiple axes at once (e.g. energy x
# pdg), each with its own expert count/hyperparameters. Axes are
# NOT in these defaults (there's no meaningful default axis list)
# — set them as flat axis{i}_{field} keys instead of "n_experts",
# e.g. axis0_type = "energy", axis0_n_experts = 4, axis1_type =
# "pdg", axis1_n_experts = 3, axis1_emb_dim = 8. See
# giant.model.network._parse_composed_axes / `--router-axis`.
},
},
}
@@ -164,15 +192,29 @@ def merge_cli_overrides(
train_overrides: dict,
model_overrides: dict,
) -> dict:
"""Resolve config as defaults -> TOML file -> explicit CLI flags."""
"""Resolve config as defaults -> TOML file -> explicit CLI flags.
`model.router` is deep-merged one level (rather than replaced wholesale)
at each stage, so a TOML file or CLI flag only overriding e.g.
`router.enabled` doesn't drop the rest of the router defaults.
"""
cfg = {"train": dict(defaults["train"]), "model": dict(defaults["model"])}
cfg["model"]["router"] = dict(defaults["model"]["router"])
if config_path is not None:
file_cfg = load_toml(config_path)
for section in ("train", "model"):
cfg[section].update(file_cfg.get(section, {}))
cfg["train"].update(file_cfg.get("train", {}))
file_model = dict(file_cfg.get("model", {}))
file_router = file_model.pop("router", None)
cfg["model"].update(file_model)
if file_router:
cfg["model"]["router"].update(file_router)
warn_if_git_hash_mismatch(file_cfg, config_path)
model_overrides = dict(model_overrides)
router_overrides = model_overrides.pop("router", None)
cfg["train"].update(train_overrides)
cfg["model"].update(model_overrides)
if router_overrides:
cfg["model"]["router"].update(router_overrides)
return cfg
@@ -184,17 +226,38 @@ def seed_everything(seed: int) -> None:
torch.cuda.manual_seed_all(seed)
def _toml_value(v) -> str:
if isinstance(v, bool):
return "true" if v else "false"
if isinstance(v, str):
return repr(v)
return str(v)
def save_config(cfg: dict, out_dir: Path, meta: dict) -> None:
lines = []
# One-level-nested dict values (e.g. model.router) are rendered as their
# own [section.subsection] table after the parent section, since TOML
# doesn't accept a bare dict as a `key = value` scalar line.
nested_sections: list[tuple[str, dict]] = []
for section, values in cfg.items():
lines.append(f"[{section}]")
for k, v in values.items():
lines.append(f"{k:<14} = {repr(v) if isinstance(v, str) else v}")
if isinstance(v, dict):
nested_sections.append((f"{section}.{k}", v))
continue
lines.append(f"{k:<14} = {_toml_value(v)}")
lines.append("")
for name, values in nested_sections:
lines.append(f"[{name}]")
for k, v in values.items():
lines.append(f"{k:<14} = {_toml_value(v)}")
lines.append("")
lines.append("[meta]")
for k, v in meta.items():
lines.append(f"{k:<14} = {repr(v) if isinstance(v, str) else v}")
lines.append(f"{k:<14} = {_toml_value(v)}")
(out_dir / "config.toml").write_text("\n".join(lines))
giant/data/dataset.py
+41 -13
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,15 +99,24 @@ 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,
require_secondaries=True,
)
(
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,
require_secondaries=True,
)
buf_cont.append(cond_cont)
buf_cat.append(cond_cat)
@@ -110,6 +124,7 @@ 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:
@@ -120,6 +135,7 @@ class StreamingStepsDataset(IterableDataset):
buf_nsec,
buf_sec,
buf_spdg,
buf_proc,
buf_n,
) = yield from self._flush(
buf_cont,
@@ -128,12 +144,20 @@ class StreamingStepsDataset(IterableDataset):
buf_nsec,
buf_sec,
buf_spdg,
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(
@@ -144,6 +168,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)
@@ -152,11 +177,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)
@@ -170,10 +196,11 @@ 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:]],
@@ -182,5 +209,6 @@ class StreamingStepsDataset(IterableDataset):
[nsec[rem:]],
[sec[rem:]],
[spdg[rem:]],
[proc[rem:]],
n - rem,
)
giant/data/loader.py
+35
View File
@@ -93,6 +93,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),
@@ -190,3 +200,28 @@ 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
+36 -4
View File
@@ -360,6 +360,8 @@ def decode_secondaries(
pdg_map_inv: maps model index → PDG code
Returns (sec_E, sec_dir_world, sec_pdg_code, sec_valid) each shape (N, K_MAX).
The valid slots' energies (`sec_E[sec_valid]`, per row) always sum to
exactly `e_sec` — see the rescaling below.
"""
N, K, _ = sec_cont.shape
stick_logits = sec_cont[:, :, 0] # (N, K)
@@ -372,15 +374,33 @@ def decode_secondaries(
fractions = 1.0 / (1.0 + np.exp(-stick_logits.astype(np.float64)))
sec_E = np.zeros((N, K), dtype=np.float32)
sec_E = np.zeros((N, K), dtype=np.float64)
e_sec = np.asarray(e_sec, dtype=np.float64)
remaining = e_sec.copy()
for i in range(K):
sec_E[:, i] = (fractions[:, i] * remaining).astype(np.float32)
remaining = np.maximum(remaining - sec_E[:, i].astype(np.float64), 0.0)
sec_E[:, i] = fractions[:, i] * remaining
remaining = np.maximum(remaining - sec_E[:, i], 0.0)
sec_valid = np.arange(K)[None, :] < n_sec[:, None] # (N, K)
# Stick-breaking guarantees sum(sec_E[valid]) <= e_sec (each fraction is in
# [0,1] of an already-shrinking remainder) but rarely hits it exactly, so
# rescale the valid slots by one common per-row factor to close that gap —
# rather than dumping the shortfall into whichever slot happens to be last
# by energy rank, which would let one low-energy secondary balloon and
# distort the shower's topology. This preserves each row's relative split
# across its secondaries and only ever scales up (valid_sum <= e_sec).
# Rows where every valid slot decoded to ~zero (scale undefined) fall back
# to an even split of e_sec across the n_sec valid slots.
sec_E = sec_E * sec_valid
valid_sum = sec_E.sum(axis=1)
degenerate = (valid_sum <= _EPS) & (n_sec > 0)
scale = np.where(valid_sum > _EPS, e_sec / np.maximum(valid_sum, _EPS), 0.0)
sec_E = sec_E * scale[:, None]
even_share = e_sec / np.maximum(n_sec, 1).astype(np.float64)
sec_E = np.where(degenerate[:, None] & sec_valid, even_share[:, None], sec_E)
sec_E = sec_E.astype(np.float32)
sec_dir_world = np.zeros((N, K, 3), dtype=np.float32)
for i in range(K):
valid = sec_valid[:, i]
@@ -433,6 +453,7 @@ def build_features(
cond_normalizer: Normalizer | None = None,
target_normalizer: Normalizer | None = None,
fit: bool = False,
proc_map: dict[str, int] | None = None,
require_secondaries: bool = False,
) -> tuple[
np.ndarray,
@@ -441,16 +462,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).
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
@@ -550,6 +575,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,
@@ -557,6 +588,7 @@ def build_features(
n_sec,
sec_cont,
sec_pdg_idx,
proc_idx,
cond_normalizer,
target_normalizer,
)
giant/model/network.py
+594 -1
View File
@@ -1,9 +1,12 @@
import inspect
import math
import re
import torch
import torch.nn as nn
import torch.nn.functional as F
from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_DIM, X_DIM
class SinusoidalEmbedding(nn.Module):
@@ -236,3 +239,593 @@ class SecondaryDecoder(nn.Module):
for block in self.blocks:
x = block(x, cond)
return self.out_proj(x)
class Router(nn.Module):
"""Contract for a pluggable mixture-of-experts routing axis.
Subclasses implement `gate` (soft partition-of-unity weights over
experts, used in train mode for a fully differentiable mixture);
`top1` and `balance_loss` have working defaults so a new routing axis
is usually a one-method add. See `ROUTER_REGISTRY` / `build_router`.
"""
def __init__(self, n_experts: int) -> None:
super().__init__()
self.n_experts = n_experts
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
"""(B, n_experts) soft weights, rows summing to 1."""
raise NotImplementedError
def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
"""(B,) hard expert index, used for eval-time grouped dispatch."""
return self.gate(cond_cont, cond_cat).argmax(dim=-1)
def balance_loss(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
) -> torch.Tensor:
"""Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017)."""
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]] = {}
def register_router(name: str):
def decorator(cls: type[Router]) -> type[Router]:
ROUTER_REGISTRY[name] = cls
return cls
return decorator
def build_router(name: str, n_experts: int, **kwargs) -> Router:
"""Factory: look up a `Router` subclass by name from the registry.
Every registered router type is fed the same `model.router` config
dict; kwargs not declared by that type's constructor are silently
dropped, so per-type hyperparameters (e.g. EnergyRouter's
`temperature`) can coexist in one config without special-casing.
"""
if name not in ROUTER_REGISTRY:
raise ValueError(
f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}"
)
cls = ROUTER_REGISTRY[name]
accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
filtered = {k: v for k, v in kwargs.items() if k in accepted}
return cls(n_experts=n_experts, **filtered)
@register_router("energy")
class EnergyRouter(Router):
"""Soft turn-on gate over normalized pre-step log-energy.
Reads `cond_cont[:, energy_idx]` (ignores cond_cat). Learnable (or
fixed) 1-D centers, initialized spread across [-2, 2] — roughly the
z-normalized energy range. `gate(e) = softmax_i(-(e - c_i)^2 / tau)`,
differentiable in e; as tau -> 0 this hardens to nearest-center
(Voronoi) selection, which is exactly what `top1` uses at eval.
"""
def __init__(
self,
n_experts: int = 4,
temperature: float = 0.5,
learn_centers: bool = True,
energy_idx: int = 3,
) -> None:
super().__init__(n_experts)
self.temperature = temperature
self.energy_idx = energy_idx
centers = torch.linspace(-2.0, 2.0, n_experts)
if learn_centers:
self.centers = nn.Parameter(centers)
else:
self.register_buffer("centers", centers)
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
e = cond_cont[:, self.energy_idx].unsqueeze(-1) # (B, 1)
d2 = (e - self.centers.unsqueeze(0)) ** 2 # (B, n_experts)
return torch.softmax(-d2 / self.temperature, dim=-1)
@register_router("pdg")
class PdgRouter(Router):
"""Soft turn-on gate over a learned PDG embedding.
Unlike ProcessRouter's process label, PDG code is already known at
pre-step time (it's a conditioning input, `cond_cat[:, 0]`), so no
supervision is needed — `classify_loss` falls back to the Router base
class's zero-loss default, same as EnergyRouter. Because PDG is
categorical rather than a scalar, this generalizes EnergyRouter's
soft-turn-on-then-Voronoi trick from a 1-D distance to a distance in a
small embedding space: its own embedding table (kept separate from the
trunk's ConditionEncoder, same reasoning as ProcessRouter's own
pdg/mat embeddings) maps each PDG code to a point, and `n_experts`
learnable (or fixed) centers partition that space.
`gate(pdg) = softmax_i(-||emb(pdg) - c_i||^2 / tau)`.
"""
def __init__(
self,
n_experts: int,
pdg_vocab: int,
emb_dim: int = 8,
temperature: float = 0.5,
learn_centers: bool = True,
) -> None:
super().__init__(n_experts)
self.temperature = temperature
self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
centers = torch.randn(n_experts, emb_dim) * 0.1
if learn_centers:
self.centers = nn.Parameter(centers)
else:
self.register_buffer("centers", centers)
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
e = self.pdg_emb(cond_cat[:, 0]) # (B, emb_dim)
d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(
-1
) # (B, n_experts)
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 ComposedRouter(Router):
"""Joint router over independent axes (e.g. energy x pdg), outer-product gated.
Wraps N already-built sub-routers, each free to have its own
`n_experts` and hyperparameters (an `EnergyRouter(n_experts=4, ...)`
composed with a `PdgRouter(n_experts=3, ...)` needs no axis to match
the other's expert count). The joint gate is the outer product of the
per-axis softmax gates, flattened to `(B, prod(n_experts_i))` — still a
partition of unity, since each factor is one. Because the axes are
routed independently, the joint argmax factors into the per-axis
argmaxes, so `top1` (inherited from `Router`) costs no more than
routing each axis alone despite the multiplicative expert count; the
same is true of `balance_loss` (inherited, computed on the flattened
joint gate — now one importance term per *joint* expert cell).
Not registered in `ROUTER_REGISTRY` / buildable via `build_router`,
since those assume one `n_experts` int shared by a single router type;
use `build_composed_router` instead, which resolves a list of per-axis
specs (each independently typed and sized) through `build_router`.
"""
def __init__(self, routers: list[Router]) -> None:
if not routers:
raise ValueError("ComposedRouter needs at least one sub-router")
n_experts = 1
for r in routers:
n_experts *= r.n_experts
super().__init__(n_experts)
self.routers = nn.ModuleList(routers)
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
joint = self.routers[0].gate(cond_cont, cond_cat) # (B, n_0)
for router in self.routers[1:]:
g = router.gate(cond_cont, cond_cat) # (B, n_i)
joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(
1
) # (B, prod so far)
return joint
def classify_loss(
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
) -> torch.Tensor:
"""Sum of each sub-router's own classify_loss (0 for unsupervised axes)."""
total = torch.zeros((), device=cond_cont.device)
for router in self.routers:
total = total + router.classify_loss(cond_cont, cond_cat, labels)
return total
def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter:
"""Build a `ComposedRouter` from a list of per-axis router specs.
Each spec is a `{"type": ..., "n_experts": ..., ...per-axis kwargs}`
dict resolved through `build_router` exactly like a single-axis router
config, so axes can differ in both expert count and hyperparameters
(e.g. an energy axis's `temperature` vs a pdg axis's `emb_dim`).
`shared_kwargs` (`pdg_vocab`, `mat_vocab`, ...) are merged under each
spec, with the spec's own keys taking precedence.
"""
routers = [
build_router(
spec["type"],
spec["n_experts"],
**{
**shared_kwargs,
**{k: v for k, v in spec.items() if k not in ("type", "n_experts")},
},
)
for spec in specs
]
return ComposedRouter(routers)
class ExpertTrunk(nn.Module):
"""One small expert: `input_proj -> ResBlock stack -> out_proj`.
Same shape as the monolithic DenoisingMLP/SecondaryDecoder trunk, but
intended to be narrower/shallower (per-call cost is the whole point).
"""
def __init__(
self,
in_dim: int,
hidden_dim: int,
n_blocks: int,
merged_cond_dim: int,
dropout: float = 0.1,
) -> None:
super().__init__()
self.input_proj = nn.Linear(in_dim, hidden_dim)
self.blocks = nn.ModuleList(
[
ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
for _ in range(n_blocks)
]
)
self.out_proj = nn.Linear(hidden_dim, in_dim)
def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
x = self.input_proj(x)
for block in self.blocks:
x = block(x, cond)
return self.out_proj(x)
def _route_forward(
experts: nn.ModuleList,
router: Router,
x: torch.Tensor,
cond: torch.Tensor,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
training: bool,
) -> torch.Tensor:
"""Shared dispatch for both Routed* trunks.
Train mode: full soft mixture `sum_i gate_i * expert_i(x)` — fully
differentiable, N-expert compute. Eval mode: grouped top-1 dispatch —
each row runs exactly one (small) expert, which is the actual source
of the per-call speedup this architecture is for.
"""
if training:
weights = router.gate(cond_cont, cond_cat) # (B, n_experts)
out = torch.zeros_like(x)
for i, expert in enumerate(experts):
out = out + weights[:, i : i + 1] * expert(x, cond)
return out
idx = router.top1(cond_cont, cond_cat) # (B,)
out = torch.zeros_like(x)
for i, expert in enumerate(experts):
mask = idx == i
if mask.any():
out[mask] = expert(x[mask], cond[mask])
return out
class RoutedDenoisingMLP(nn.Module):
"""Routed drop-in for `DenoisingMLP`.
Shares the time embedding, `ConditionEncoder`, and `n_sec_head` (all
tiny) across experts and routes only the trunk (where the FLOPs are).
Same `forward`/`predict_n_sec`/`pdg_embedding_weight` signatures as
`DenoisingMLP`, so sample.py/rollout.py/validate.py need no changes.
"""
def __init__(
self,
pdg_vocab: int,
mat_vocab: int,
router: Router,
expert_hidden_dim: int = 128,
expert_n_blocks: int = 3,
emb_dim: int = EMB_DIM,
time_dim: int = 64,
cond_out_dim: int = 128,
x_dim: int = X_DIM,
dropout: float = 0.1,
k_max: int = K_MAX,
) -> None:
super().__init__()
self.router = router
self.time_emb = SinusoidalEmbedding(time_dim)
self.cond_enc = ConditionEncoder(
pdg_vocab=pdg_vocab,
mat_vocab=mat_vocab,
emb_dim=emb_dim,
out_dim=cond_out_dim,
)
merged_cond_dim = time_dim + cond_out_dim
self.experts = nn.ModuleList(
[
ExpertTrunk(
x_dim,
expert_hidden_dim,
expert_n_blocks,
merged_cond_dim,
dropout=dropout,
)
for _ in range(router.n_experts)
]
)
self.n_sec_head = nn.Sequential(
nn.Linear(cond_out_dim, cond_out_dim),
nn.SiLU(),
nn.Linear(cond_out_dim, k_max + 1),
)
def forward(
self,
x_t: torch.Tensor,
t: torch.Tensor,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
) -> torch.Tensor:
t_emb = self.time_emb(t)
c_emb = self.cond_enc(cond_cont, cond_cat)
cond = torch.cat([t_emb, c_emb], dim=-1)
return _route_forward(
self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
)
def predict_n_sec(
self,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
) -> torch.Tensor:
"""Return n_sec logits (B, K_MAX+1) from conditioning alone."""
c_emb = self.cond_enc(cond_cont, cond_cat)
return self.n_sec_head(c_emb)
def pdg_embedding_weight(self) -> torch.Tensor:
"""Return the PDG embedding table weights for secondary type targets."""
return self.cond_enc.pdg_emb.weight
class RoutedSecondaryDecoder(nn.Module):
"""Routed drop-in for `SecondaryDecoder`.
Shares the time embedding and `SecondaryConditionEncoder` across
experts and routes only the trunk. Same `forward` signature as
`SecondaryDecoder`.
"""
def __init__(
self,
pdg_vocab: int,
mat_vocab: int,
router: Router,
expert_hidden_dim: int = 128,
expert_n_blocks: int = 3,
emb_dim: int = EMB_DIM,
time_dim: int = 64,
cond_out_dim: int = 128,
stage1_proj_dim: int = 64,
sec_dim: int = SEC_DIM,
dropout: float = 0.1,
) -> None:
super().__init__()
self.router = router
self.time_emb = SinusoidalEmbedding(time_dim)
self.cond_enc = SecondaryConditionEncoder(
pdg_vocab=pdg_vocab,
mat_vocab=mat_vocab,
emb_dim=emb_dim,
cond_out_dim=cond_out_dim,
stage1_proj_dim=stage1_proj_dim,
out_dim=cond_out_dim,
)
merged_cond_dim = time_dim + cond_out_dim
self.experts = nn.ModuleList(
[
ExpertTrunk(
sec_dim,
expert_hidden_dim,
expert_n_blocks,
merged_cond_dim,
dropout=dropout,
)
for _ in range(router.n_experts)
]
)
def forward(
self,
x_t: torch.Tensor,
t: torch.Tensor,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
stage1_out: torch.Tensor,
) -> torch.Tensor:
t_emb = self.time_emb(t)
c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
cond = torch.cat([t_emb, c_emb], dim=-1)
return _route_forward(
self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
)
_STAGE1_MODEL_KEYS = {
"pdg_vocab",
"mat_vocab",
"hidden_dim",
"n_blocks",
"emb_dim",
"dropout",
"k_max",
}
_SEC_DECODER_MODEL_KEYS = {
"pdg_vocab",
"mat_vocab",
"hidden_dim",
"n_blocks",
"emb_dim",
"dropout",
}
_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$")
def _parse_composed_axes(router_cfg: dict) -> list[dict]:
"""Regroup `axis{i}_{field}` flat keys into a list of per-axis spec dicts.
Flat keys (rather than a nested list-of-dicts) keep composed-router
config expressible in the same one-level-of-nesting TOML/CLI shape as
every other router option (`model.router` stays a flat table of
scalars) — e.g. `axis0_type = "energy"`, `axis0_n_experts = 4`,
`axis1_type = "pdg"`, `axis1_n_experts = 3`, `axis1_emb_dim = 8`.
Axis indices must be contiguous from 0; order follows the index, not
dict insertion order (TOML/CLI merging doesn't preserve it reliably).
"""
axes: dict[int, dict] = {}
for key, value in router_cfg.items():
m = _AXIS_KEY_RE.match(key)
if m is None:
continue
idx, field = int(m.group(1)), m.group(2)
axes.setdefault(idx, {})[field] = value
missing = set(range(len(axes))) - axes.keys()
if missing:
raise ValueError(f"composed router config has gaps at axis indices {missing}")
return [axes[i] for i in range(len(axes))]
def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) -> Router:
"""Resolve one `model.router` config into a `Router`, single-axis or composed.
`router_cfg["type"] == "composed"` reads `axis{i}_{field}` flat keys
(see `_parse_composed_axes`) instead of a single `type`/`n_experts` pair.
"""
shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
if router_cfg["type"] == "composed":
return build_composed_router(_parse_composed_axes(router_cfg), **shared_vocab)
router_kwargs = {
k: v 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", pdg_vocab)
router_kwargs.setdefault("mat_vocab", mat_vocab)
return build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
"""Construct (stage1, sec_decoder) from a persisted/CLI model_config dict.
Dispatches to the routed pair when `model_config["router"]["enabled"]`
is truthy; a missing/absent "router" key (pre-routing checkpoints)
falls back to the monolithic pair unchanged, so this is a drop-in
replacement for the ad-hoc constructions it replaces.
"""
router_cfg = model_config.get("router")
if router_cfg and router_cfg.get("enabled"):
pdg_vocab = model_config["pdg_vocab"]
mat_vocab = model_config["mat_vocab"]
shared = dict(
pdg_vocab=pdg_vocab,
mat_vocab=mat_vocab,
expert_hidden_dim=model_config.get("expert_hidden_dim", 128),
expert_n_blocks=model_config.get("expert_n_blocks", 3),
emb_dim=model_config.get("emb_dim", EMB_DIM),
dropout=model_config.get("dropout", 0.1),
)
stage1 = RoutedDenoisingMLP(
router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab),
k_max=model_config.get("k_max", K_MAX),
**shared,
)
sec_decoder = RoutedSecondaryDecoder(
router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab),
**shared,
)
return stage1, sec_decoder
stage1 = DenoisingMLP(
**{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS}
)
sec_decoder = SecondaryDecoder(
**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
)
return stage1, sec_decoder
giant/pipeline.py
+44 -31
View File
@@ -11,10 +11,11 @@ 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
from giant.model.network import DenoisingMLP, SecondaryDecoder
from giant.model.network import build_models
from giant.train import train as run_training
@@ -54,6 +55,17 @@ 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 +75,14 @@ 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, require_secondaries=True
cond_cont, _, target_s1, _n_sec, _sec_cont, _sec_pdg, _proc, _, _ = (
build_features(
chunk_tr,
pdg_map,
mat_map,
proc_map=proc_map,
require_secondaries=True,
)
)
cond_acc.update(cond_cont)
tgt_acc.update(target_s1)
@@ -81,6 +99,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 +110,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,23 +134,21 @@ def run_train_job(
"update giant/constants.py if emb_dim changed"
)
stage1_model = DenoisingMLP(
pdg_vocab=len(pdg_map),
mat_vocab=len(mat_map),
hidden_dim=m["hidden_dim"],
n_blocks=m["n_blocks"],
emb_dim=emb_dim,
dropout=m["dropout"],
k_max=K_MAX,
)
sec_decoder = SecondaryDecoder(
pdg_vocab=len(pdg_map),
mat_vocab=len(mat_map),
hidden_dim=m["hidden_dim"],
n_blocks=m["n_blocks"],
emb_dim=emb_dim,
dropout=m["dropout"],
)
model_config = {
"pdg_vocab": len(pdg_map),
"mat_vocab": len(mat_map),
"hidden_dim": m["hidden_dim"],
"n_blocks": m["n_blocks"],
"emb_dim": emb_dim,
"dropout": m["dropout"],
"k_max": K_MAX,
"sec_slot_dim": SEC_SLOT_DIM,
"router": dict(router_cfg),
"expert_hidden_dim": router_cfg["expert_hidden_dim"],
"expert_n_blocks": router_cfg["expert_n_blocks"],
}
stage1_model, sec_decoder = build_models(model_config)
echo(
f"stage1: {sum(p.numel() for p in stage1_model.parameters()):,} parameters | "
f"sec_decoder: {sum(p.numel() for p in sec_decoder.parameters()):,} parameters"
@@ -148,17 +166,6 @@ def run_train_job(
)
config.save_config(cfg, out_dir, meta)
model_config = {
"pdg_vocab": len(pdg_map),
"mat_vocab": len(mat_map),
"hidden_dim": m["hidden_dim"],
"n_blocks": m["n_blocks"],
"emb_dim": emb_dim,
"dropout": m["dropout"],
"k_max": K_MAX,
"sec_slot_dim": SEC_SLOT_DIM,
}
run_training(
stage1_model=stage1_model,
sec_decoder=sec_decoder,
@@ -167,17 +174,23 @@ def run_train_job(
mode=t["mode"],
epochs=t["epochs"],
lr=t["lr"],
weight_decay=t["weight_decay"],
ema_decay=t["ema_decay"],
warmup_epochs=t["warmup_epochs"],
device=device,
out_dir=out_dir,
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"],
validate_steps=t["validate_steps"],
max_val_batches=t["max_val_batches"],
total_train_batches=total_train_batches,
)
giant/rollout.py
+103 -18
View File
@@ -16,6 +16,9 @@ treated as detector leakage and not deposited.
from __future__ import annotations
from collections import Counter
from typing import Callable
import numpy as np
import torch
@@ -89,32 +92,95 @@ def _concat_frontiers(parts: list[dict[str, np.ndarray]]) -> dict[str, np.ndarra
return {k: np.concatenate([p[k] for p in parts], axis=0) for k in parts[0]}
class _Recorder:
"""Accumulates per-step rows into column lists, materialised at the end."""
# Fixed per-key dtype, so every chunk table has an identical schema — needed
# for `giant rollout --on_chunk` to stream chunks straight into one
# pq.ParquetWriter (which requires matching schemas across writes), and a
# side benefit even in the buffered path since np.concatenate would otherwise
# silently upcast any stray int32/float32 chunk to the majority dtype.
_RECORD_DTYPES: dict[str, type] = {
"event_id": np.int64,
"track_id": np.int64,
"parent_id": np.int64,
"generation": np.int64,
"step_no": np.int64,
"pdg": np.int64,
"pre_x": np.float64,
"pre_y": np.float64,
"pre_z": np.float64,
"pre_E": np.float64,
"pre_dx": np.float64,
"pre_dy": np.float64,
"pre_dz": np.float64,
"post_x": np.float64,
"post_y": np.float64,
"post_z": np.float64,
"post_E": np.float64,
"post_dx": np.float64,
"post_dy": np.float64,
"post_dz": np.float64,
"edep": np.float64,
"step_length": np.float64,
"material": object,
"layer_id": np.int64,
"n_sec_pred": np.int64,
"termination_reason": object,
}
def __init__(self) -> None:
self._cols: dict[str, list] = {k: [] for k in _RECORD_KEYS}
class _Recorder:
"""Accumulates per-step rows into column lists, materialised at the end —
or, when `sink` is given, streams each non-empty chunk to it immediately
instead, keeping only row-count / termination-reason summaries in memory.
The streaming path is what lets `giant rollout` write output incrementally
(see `rollout`'s `on_chunk` parameter): without it, a whole run's steps —
scaling with `n_events * max_steps * avg_tracks_per_event` — would sit in
RAM until the very end.
"""
def __init__(
self, sink: Callable[[dict[str, np.ndarray]], None] | None = None
) -> None:
self._sink = sink
self._cols: dict[str, list] | None = (
None if sink is not None else {k: [] for k in _RECORD_KEYS}
)
self.n_rows = 0
self.termination_reason_counts: Counter[str] = Counter()
def add(self, **cols) -> None:
n = len(cols["event_id"])
if n == 0:
return
for k in _RECORD_KEYS:
v = cols[k]
self._cols[k].append(np.asarray(v).reshape(n))
row = {
k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n)
for k in _RECORD_KEYS
}
self.n_rows += n
reasons = row["termination_reason"]
nonempty = reasons[reasons != ""]
if len(nonempty):
for r, c in zip(*np.unique(nonempty, return_counts=True)):
self.termination_reason_counts[str(r)] += int(c)
if self._sink is not None:
self._sink(row)
else:
assert self._cols is not None
for k in _RECORD_KEYS:
self._cols[k].append(row[k])
def to_dict(self) -> dict[str, np.ndarray]:
assert self._cols is not None, (
"to_dict() is unavailable when streaming to a sink — use "
"n_rows/termination_reason_counts instead"
)
out = {}
for k, chunks in self._cols.items():
if chunks:
out[k] = np.concatenate(chunks, axis=0)
else:
out[k] = np.empty(
0,
dtype=object
if k in ("material", "termination_reason")
else np.float64,
)
out[k] = np.empty(0, dtype=_RECORD_DTYPES[k])
return out
@@ -209,8 +275,20 @@ def rollout(
device: torch.device | None = None,
max_tracks_per_event: int | None = None,
escape_threshold: float | None = None,
on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
) -> dict[str, np.ndarray]:
"""Run showers to completion; return a step-record dict (see _RECORD_KEYS)."""
"""Run showers to completion.
By default, returns a step-record dict (see _RECORD_KEYS) with the whole
run's rows materialised in memory.
If `on_chunk` is given, every non-empty batch of rows is streamed to it as
soon as it's produced instead — no per-run buffering — and this returns a
small summary dict instead: `{"n_rows": int, "termination_reason_counts":
dict[str, int]}`. Use this for large `--n-events`/`--max-steps` runs,
where the full record set would otherwise scale with
`n_events * max_steps * avg_tracks_per_event`.
"""
device = device or torch.device("cpu")
stage1_model.eval()
sec_decoder.eval()
@@ -227,7 +305,7 @@ def rollout(
seeds["pre_E"],
seeds["pre_dir"],
)
rec = _Recorder()
rec = _Recorder(sink=on_chunk)
while len(frontier["event_id"]) > 0:
next_parts: list[dict[str, np.ndarray]] = []
@@ -257,6 +335,11 @@ def rollout(
)
frontier = _concat_frontiers(next_parts)
if on_chunk is not None:
return {
"n_rows": rec.n_rows,
"termination_reason_counts": dict(rec.termination_reason_counts),
}
return rec.to_dict()
@@ -394,9 +477,11 @@ def _step_chunk(
max_tracks_per_event,
)
# Energy bookkeeping so each step conserves exactly (edep + carried + post_E
# == pre_E): the primary lost `e_sec` to secondaries, but the decoded
# secondaries only carry `sec_E[valid].sum()`. Deposit the unallocated
# residual locally, plus the energy of any sub-cap secondaries we dropped.
# == pre_E): `decode_secondaries` already rescales valid slots to sum to
# exactly `e_sec` whenever n_sec > 0, so `residual` here is ~0 except when
# n_sec == 0 (no secondary to carry the budget at all — the whole `e_sec`
# becomes residual). Also deposit the energy of any sub-cap secondaries
# we dropped for hitting `max_tracks_per_event`.
sec_E_valid_sum = (sec_E * sec_valid).sum(axis=1)
residual = np.maximum(e_sec - sec_E_valid_sum, 0.0)
edep = edep + residual + dropped_edep
giant/train.py
+142 -21
View File
@@ -1,3 +1,4 @@
import copy
import csv
import math
import os
@@ -26,11 +27,16 @@ _METRICS_FIELDS = [
"train_loss_s1",
"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",
"grad_norm",
"epoch_time_s",
]
@@ -99,6 +105,14 @@ def _build_sec_x1(
return x1_s2.flatten(1) # (B, SEC_DIM)
@torch.no_grad()
def _update_ema(
ema_model: torch.nn.Module, model: torch.nn.Module, decay: float
) -> None:
for ema_p, p in zip(ema_model.parameters(), model.parameters()):
ema_p.mul_(decay).add_(p, alpha=1 - decay)
def _compute_losses(
stage1_model: torch.nn.Module,
sec_decoder: torch.nn.Module,
@@ -108,15 +122,20 @@ def _compute_losses(
device: torch.device,
lambda_nsec: float,
lambda_s2: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Compute (total_loss, L_s1, L_nsec, L_s2) for one batch."""
cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_idx = batch
lambda_balance: float = 0.0,
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":
@@ -150,8 +169,29 @@ def _compute_losses(
sec_mask,
)
# Optional MoE load-balance auxiliary loss: only present when both stages
# are routed (RoutedDenoisingMLP/RoutedSecondaryDecoder carry `.router`,
# the monolith models don't), computed on cond_cont alone (cheap — no
# trunk compute) so it's reported even when lambda_balance == 0.
if hasattr(stage1_model, "router") and hasattr(sec_decoder, "router"):
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
return total, l_s1, l_nsec, l_s2
if lambda_balance > 0:
total = total + lambda_balance * 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(
@@ -165,15 +205,21 @@ def train(
warmup_epochs: int,
device: torch.device,
out_dir: str | Path,
weight_decay: float = 0.01,
ema_decay: float = 0.9999,
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,
validate_steps: int = 10,
max_val_batches: int = 0,
total_train_batches: int = 0,
) -> None:
out_dir = Path(out_dir)
@@ -182,15 +228,38 @@ def train(
stage1_model = stage1_model.to(device)
sec_decoder = sec_decoder.to(device)
all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters())
optimizer = optim.AdamW(all_params, lr=lr)
# Flow-matching/diffusion models sample noticeably better from an EMA of
# the weights than from the raw SGD-noisy ones — buffers (e.g. the fixed
# sinusoidal-embedding freqs, or non-learned router centers) never change
# after this initial copy, so only parameters need the running average.
ema_stage1_model: torch.nn.Module | None = None
ema_sec_decoder: torch.nn.Module | None = None
if ema_decay > 0:
ema_stage1_model = copy.deepcopy(stage1_model).eval()
ema_sec_decoder = copy.deepcopy(sec_decoder).eval()
for p in ema_stage1_model.parameters():
p.requires_grad_(False)
for p in ema_sec_decoder.parameters():
p.requires_grad_(False)
def _lr_lambda(epoch: int) -> float:
if warmup_epochs > 0 and epoch < warmup_epochs:
return (epoch + 1) / warmup_epochs
t = epoch - warmup_epochs
T = max(epochs - warmup_epochs, 1)
return 0.5 * (1.0 + math.cos(math.pi * t / T))
all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters())
optimizer = optim.AdamW(all_params, lr=lr, weight_decay=weight_decay)
# Warmup/decay in units of optimizer steps rather than epochs: at large
# dataset sizes a single epoch can be tens of thousands of steps, and an
# epoch-granularity schedule would leave warmup/cosine decay unable to
# move within it. Requires an accurate `total_train_batches` (steps per
# epoch); the only caller, run_train_job, always supplies one.
steps_per_epoch = max(total_train_batches, 1)
warmup_steps = warmup_epochs * steps_per_epoch
total_steps = max(epochs * steps_per_epoch, 1)
def _lr_lambda(step: int) -> float:
if warmup_steps > 0 and step < warmup_steps:
return (step + 1) / warmup_steps
t = step - warmup_steps
T = max(total_steps - warmup_steps, 1)
return 0.5 * (1.0 + math.cos(math.pi * min(t, T) / T))
lr_sched = optim.lr_scheduler.LambdaLR(optimizer, _lr_lambda)
@@ -202,6 +271,12 @@ def train(
ckpt = torch.load(resume_path, map_location=device, weights_only=False)
stage1_model.load_state_dict(ckpt["model"])
sec_decoder.load_state_dict(ckpt["sec_decoder"])
if ema_decay > 0:
assert ema_stage1_model is not None and ema_sec_decoder is not None
ema_stage1_model.load_state_dict(ckpt.get("model_ema", ckpt["model"]))
ema_sec_decoder.load_state_dict(
ckpt.get("sec_decoder_ema", ckpt["sec_decoder"])
)
optimizer.load_state_dict(ckpt["optimizer"])
lr_sched.load_state_dict(ckpt["lr_sched"])
start_epoch = ckpt.get("epoch", 0) + 1
@@ -238,15 +313,19 @@ def train(
with _GracefulShutdown() as shutdown:
for epoch in range(start_epoch, epochs + 1):
epoch_start = time.monotonic()
current_lr = optimizer.param_groups[0]["lr"]
stage1_model.train()
sec_decoder.train()
train_loss_sum = 0.0
train_s1_sum = 0.0
train_nsec_sum = 0.0
train_s2_sum = 0.0
train_balance_sum = 0.0
train_proc_sum = 0.0
train_n = 0
train_batches = 0
grad_norm_sum = 0.0
ema_loss = 0.0
ema_grad_norm = 0.0
bar = tqdm(
train_loader,
desc=f" epoch {epoch:{epoch_w}d}/{epochs}",
@@ -256,7 +335,7 @@ def train(
dynamic_ncols=True,
)
for batch in bar:
loss, l_s1, l_nsec, l_s2 = _compute_losses(
loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses(
stage1_model,
sec_decoder,
batch,
@@ -265,23 +344,42 @@ def train(
device,
lambda_nsec,
lambda_s2,
lambda_balance,
lambda_proc,
)
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(all_params, 1.0)
grad_norm = torch.nn.utils.clip_grad_norm_(all_params, 1.0)
optimizer.step()
lr_sched.step()
if ema_decay > 0:
assert ema_stage1_model is not None and ema_sec_decoder is not None
_update_ema(ema_stage1_model, stage1_model, ema_decay)
_update_ema(ema_sec_decoder, sec_decoder, ema_decay)
B = batch[0].size(0)
batch_loss = loss.item()
batch_grad_norm = grad_norm.item()
train_loss_sum += batch_loss * B
train_s1_sum += l_s1.item() * B
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
train_batches += 1
grad_norm_sum += batch_grad_norm
ema_loss = (
batch_loss if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss
)
bar.set_postfix_str(f"loss={ema_loss:.4f}", refresh=False)
ema_grad_norm = (
batch_grad_norm
if train_batches == 1
else 0.95 * ema_grad_norm + 0.05 * batch_grad_norm
)
bar.set_postfix_str(
f"loss={ema_loss:.4f} gnorm={ema_grad_norm:.3f}", refresh=False
)
if shutdown.requested:
break
@@ -291,7 +389,8 @@ def train(
break
train_loss = train_loss_sum / max(train_n, 1)
lr_sched.step()
train_grad_norm = grad_norm_sum / max(train_batches, 1)
current_lr = optimizer.param_groups[0]["lr"]
stage1_model.eval()
sec_decoder.eval()
@@ -299,10 +398,14 @@ def train(
val_s1_sum = 0.0
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 = _compute_losses(
for val_batch_idx, batch in enumerate(val_loader):
if max_val_batches > 0 and val_batch_idx >= max_val_batches:
break
loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses(
stage1_model,
sec_decoder,
batch,
@@ -311,12 +414,16 @@ def train(
device,
lambda_nsec,
lambda_s2,
lambda_balance,
lambda_proc,
)
B = batch[0].size(0)
val_loss_sum += loss.item() * B
val_s1_sum += l_s1.item() * B
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
@@ -328,9 +435,12 @@ def train(
f" train {train_loss:.4f}"
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" s2={train_s2_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}"
f" lr {current_lr:.2e} gnorm {train_grad_norm:.3f}"
f" {epoch_time:.1f}s{marker}"
)
metrics_writer.writerow(
{
@@ -339,11 +449,16 @@ def train(
"train_loss_s1": train_s1_sum / max(train_n, 1),
"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,
"grad_norm": train_grad_norm,
"epoch_time_s": epoch_time,
}
)
@@ -369,12 +484,18 @@ def train(
"epoch": epoch,
"best_val_loss": best_val_loss,
}
if ema_decay > 0:
assert ema_stage1_model is not None and ema_sec_decoder is not None
ckpt["model_ema"] = ema_stage1_model.state_dict()
ckpt["sec_decoder_ema"] = ema_sec_decoder.state_dict()
if normalizer_dict is not None:
ckpt["normalizer"] = normalizer_dict
if pdg_map is not None:
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)
scripts/create_root_files.py
+41 -23
View File
@@ -3,15 +3,17 @@ run_pbwo4, run_sampling) and filing the output into the dataset's raw/ tree:
raw/<kind>/<gen>/<detector>/shard-NNN.root
These executables take `[configName] nEvents` and always write a fixed-name
*.root file into the current directory — so running several in parallel
needs separate working directories, and the output filename has to be
discovered rather than assumed (it differs per executable: run_pbwo4 writes
pbwo4_<n>events_hits.root, run_sampling writes sampling_<config>_<n>events_hits.root,
others may differ again). This script gives each run its own scratch
directory under <dataset-root>/.sim-tmp/, requires exactly one *.root to
appear there, and moves it to the next free shard index for that detector
(existing shards are never overwritten).
These executables take `[configName] nEvents [energy_GeV]` (configName is
only accepted by executables with a config selector, e.g. run_sampling;
energy_GeV defaults to 1.0 in the executable itself if omitted here) and
always write a fixed-name *.root file into the current directory — so
running several in parallel needs separate working directories, and the
output filename has to be discovered rather than assumed (it differs per
executable: run_pbwo4 writes pbwo4_<n>events_hits.root, run_sampling writes
sampling_<config>_<n>events_hits.root, others may differ again). This script
gives each run its own scratch directory under <dataset-root>/.sim-tmp/,
requires exactly one *.root to appear there, and moves it to the next free
shard index for that detector (existing shards are never overwritten).
--gen must already exist under raw/<kind>/ — create one first with
`dwarf bump-gen`.
@@ -105,8 +107,8 @@ def plan_jobs(
return jobs
def job_seed(kind: str, gen: str, job: SimJob) -> int:
"""Deterministic RNG seed for one sim job, unique per (kind, gen, detector, config, shard).
def job_seed(kind: str, gen: str, job: SimJob, energy_gev: float | None) -> int:
"""Deterministic RNG seed for one sim job, unique per (kind, gen, detector, config, shard, energy).
Jobs run concurrently (ThreadPoolExecutor below) and can start within the
same wall-clock second; minicalosim's default seed falls back to
@@ -115,14 +117,31 @@ def job_seed(kind: str, gen: str, job: SimJob) -> int:
landing in separate shard files. Deriving the seed from the full job
identity instead keeps it both unique and reproducible.
"""
key = f"{kind}|{gen}|{job.detector}|{job.config or ''}|{job.shard_index}"
key = (
f"{kind}|{gen}|{job.detector}|{job.config or ''}|{job.shard_index}"
f"|{energy_gev if energy_gev is not None else ''}"
)
return zlib.crc32(key.encode()) & 0x7FFFFFFF
def build_cmd(
executable: Path, job: SimJob, events_per_file: int, energy_gev: float | None
) -> list[str]:
"""minicalosim executables take positional `[configName] nEvents [energy_GeV]`."""
cmd = [str(executable)]
if job.config:
cmd.append(job.config)
cmd.append(str(events_per_file))
if energy_gev is not None:
cmd.append(str(energy_gev))
return cmd
def run_job(
job: SimJob,
executable: Path,
events_per_file: int,
energy_gev: float | None,
dataset_root: Path,
kind: str,
gen: str,
@@ -134,12 +153,9 @@ def run_job(
)
workdir.mkdir(parents=True)
cmd = [str(executable)]
if job.config:
cmd.append(job.config)
cmd.append(str(events_per_file))
cmd = build_cmd(executable, job, events_per_file, energy_gev)
env = dict(os.environ, MINICALOSIM_SEED=str(job_seed(kind, gen, job)))
env = dict(os.environ, MINICALOSIM_SEED=str(job_seed(kind, gen, job, energy_gev)))
result = subprocess.run(cmd, cwd=workdir, capture_output=True, text=True, env=env)
if result.returncode != 0:
@@ -194,6 +210,7 @@ def run_all(
jobs: list[SimJob],
executable: Path,
events_per_file: int,
energy_gev: float | None,
dataset_root: Path,
kind: str,
gen: str,
@@ -208,6 +225,7 @@ def run_all(
job,
executable,
events_per_file,
energy_gev,
dataset_root,
kind,
gen,
@@ -238,6 +256,7 @@ def run_make_root(
dataset_root: str,
jobs: int,
execute: bool,
energy_gev: float | None = None,
) -> None:
if jobs < 1:
raise SystemExit("error: --jobs must be >= 1")
@@ -245,6 +264,8 @@ def run_make_root(
raise SystemExit("error: --num-files must be >= 1")
if events_per_file < 1:
raise SystemExit("error: --events-per-file must be >= 1")
if energy_gev is not None and energy_gev <= 0:
raise SystemExit("error: --energy-gev must be > 0")
if not executable.is_file() or not os.access(executable, os.X_OK):
raise SystemExit(f"error: {executable} is not an executable file")
@@ -257,11 +278,7 @@ def run_make_root(
print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
print(f"executable: {executable}")
for job in planned_jobs:
cmd = (
[str(executable)]
+ ([job.config] if job.config else [])
+ [str(events_per_file)]
)
cmd = build_cmd(executable, job, events_per_file, energy_gev)
dest = (
dataset_root_path
/ "raw"
@@ -270,7 +287,7 @@ def run_make_root(
/ job.detector
/ f"shard-{job.shard_index:03d}.root"
)
seed = job_seed(kind, gen, job)
seed = job_seed(kind, gen, job, energy_gev)
print(f" MINICALOSIM_SEED={seed} {' '.join(cmd)} -> {dest}")
if not execute:
@@ -283,6 +300,7 @@ def run_make_root(
planned_jobs,
executable,
events_per_file,
energy_gev,
dataset_root_path,
kind,
gen,
scripts/dwarf.py
+11
View File
@@ -363,6 +363,16 @@ def make_root(
gen: Annotated[
str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")
],
energy_gev: Annotated[
float | None,
typer.Option(
"--energy-gev",
help="energy_GeV passed to the executable (default: executable's own "
"default, currently 1.0). Note the dataset detector label is not "
"derived from this — e.g. use '--detector pbwo4_10gev --energy-gev 10' "
"to name the dataset accordingly.",
),
] = None,
kind: Annotated[
str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
] = "steps",
@@ -390,6 +400,7 @@ def make_root(
dataset_root=str(dataset_root),
jobs=jobs,
execute=execute,
energy_gev=energy_gev,
)
scripts/steps_to_parquet_parallel.py
+6 -2
View File
@@ -217,8 +217,12 @@ def run_parallel_job(
raise SystemExit(1)
total_orphaned = sum(
int(m.group(1)) for _, _, stdout, _ in results for m in _ORPHAN_RE.finditer(stdout)
int(m.group(1))
for _, _, stdout, _ in results
for m in _ORPHAN_RE.finditer(stdout)
)
if total_orphaned:
print(f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s).")
print(
f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s)."
)
print(f"\nAll {len(results)} conversion(s) completed.")
tests/test_create_root_files.py
+41 -14
View File
@@ -139,24 +139,29 @@ def test_plan_jobs_multiple_detectors_each_start_independently(tmp_path):
def test_job_seed_deterministic():
job = SimJob(detector="pbwo4", config=None, shard_index=3)
assert job_seed("steps", "gen1", job) == job_seed("steps", "gen1", job)
assert job_seed("steps", "gen1", job, None) == job_seed("steps", "gen1", job, None)
def test_job_seed_varies_by_shard_index():
a = SimJob(detector="pbwo4", config=None, shard_index=0)
b = SimJob(detector="pbwo4", config=None, shard_index=1)
assert job_seed("steps", "gen1", a) != job_seed("steps", "gen1", b)
assert job_seed("steps", "gen1", a, None) != job_seed("steps", "gen1", b, None)
def test_job_seed_varies_by_detector():
a = SimJob(detector="pbwo4", config=None, shard_index=0)
b = SimJob(detector="sampling_pb_scint", config="pb_scint", shard_index=0)
assert job_seed("steps", "gen1", a) != job_seed("steps", "gen1", b)
assert job_seed("steps", "gen1", a, None) != job_seed("steps", "gen1", b, None)
def test_job_seed_varies_by_gen():
job = SimJob(detector="pbwo4", config=None, shard_index=0)
assert job_seed("steps", "gen1", job) != job_seed("steps", "gen2", job)
assert job_seed("steps", "gen1", job, None) != job_seed("steps", "gen2", job, None)
def test_job_seed_varies_by_energy():
job = SimJob(detector="pbwo4", config=None, shard_index=0)
assert job_seed("steps", "gen1", job, 1.0) != job_seed("steps", "gen1", job, 10.0)
def test_run_job_passes_deterministic_seed_env_var(tmp_path):
@@ -166,11 +171,11 @@ def test_run_job_passes_deterministic_seed_env_var(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="pbwo4", config=None, shard_index=5)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert result.dest is not None
payload = json.loads(result.dest.read_text())
assert payload["seed"] == str(job_seed("steps", "gen1", job))
assert payload["seed"] == str(job_seed("steps", "gen1", job, None))
def test_run_job_moves_output_to_correct_shard_path(tmp_path):
@@ -181,7 +186,7 @@ def test_run_job_moves_output_to_correct_shard_path(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="pbwo4", config=None, shard_index=7)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert result.ok
assert result.dest == gen_dir / "pbwo4" / "shard-007.root"
@@ -197,7 +202,7 @@ def test_run_job_passes_config_arg_and_isolates_cwd(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="sampling_pb_scint", config="pb_scint", shard_index=0)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert result.ok
assert result.dest is not None
@@ -215,13 +220,27 @@ def test_run_job_omits_config_arg_when_none(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="pbwo4", config=None, shard_index=0)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert result.dest is not None
payload = json.loads(result.dest.read_text())
assert payload["argv"] == ["10000"]
def test_run_job_appends_energy_arg_when_given(tmp_path):
fake = _write_fake_executable(tmp_path / "fake_exe.py")
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
tmp_root = tmp_path / ".sim-tmp"
tmp_root.mkdir()
job = SimJob(detector="pbwo4_10gev", config=None, shard_index=0)
result = run_job(job, fake, 10000, 10.0, tmp_path, "steps", "gen1", tmp_root)
assert result.dest is not None
payload = json.loads(result.dest.read_text())
assert payload["argv"] == ["10000", "10.0"]
def test_run_job_fails_when_executable_errors(tmp_path):
fake = _write_fake_executable(tmp_path / "fake_exe.py", exit_code=1)
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
@@ -229,7 +248,7 @@ def test_run_job_fails_when_executable_errors(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="pbwo4", config=None, shard_index=0)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert not result.ok
assert "exited 1" in result.message
@@ -242,7 +261,7 @@ def test_run_job_fails_when_no_root_file_produced(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="pbwo4", config=None, shard_index=0)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert not result.ok
assert "found 0" in result.message
@@ -255,7 +274,7 @@ def test_run_job_fails_when_multiple_root_files_produced(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="pbwo4", config=None, shard_index=0)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert not result.ok
assert "found 2" in result.message
@@ -270,7 +289,7 @@ def test_run_job_refuses_to_overwrite_existing_shard(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="pbwo4", config=None, shard_index=0)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert not result.ok
assert "overwrite" in result.message
@@ -285,7 +304,15 @@ def test_run_all_caps_concurrency(tmp_path):
jobs = [SimJob(detector="pbwo4", config=None, shard_index=i) for i in range(6)]
results = run_all(
jobs, fake, 10000, tmp_path, "steps", "gen1", max_workers=2, tmp_root=tmp_root
jobs,
fake,
10000,
None,
tmp_path,
"steps",
"gen1",
max_workers=2,
tmp_root=tmp_root,
)
assert all(r.ok and r.dest is not None for r in results)
tests/test_loader.py
+33 -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,34 @@ 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_phase2.py
+127
View File
@@ -231,3 +231,130 @@ def test_encode_secondaries_direction_encoding():
local_dirs = sec_cont[:, :2, 1:4] # (N, 2, 3) — valid slots only
norms_out = np.linalg.norm(local_dirs, axis=-1)
np.testing.assert_allclose(norms_out, 1.0, atol=1e-5)
# ── decode_secondaries: exact energy conservation ────────────────────────────
def _random_sec_cont(rng, N, stick_logit_scale=1.0):
sec_cont = rng.standard_normal((N, K_MAX, 4)).astype(np.float32)
sec_cont[:, :, 0] *= stick_logit_scale
dirs = sec_cont[:, :, 1:]
dirs /= np.linalg.norm(dirs, axis=-1, keepdims=True)
return sec_cont
def test_decode_secondaries_valid_slots_sum_to_e_sec():
"""The valid slots' energies must sum to exactly e_sec, not just <= e_sec.
Rows with n_sec=0 are excluded: there's no slot to put the budget in, so
valid_sum is correctly 0 regardless of e_sec there (see
test_decode_secondaries_zero_n_sec_has_zero_energy) — the shortfall in
that case is handled downstream (e.g. rollout.py dumps it into edep).
"""
from giant.data.transforms import decode_secondaries
rng = np.random.default_rng(0)
N = 200
sec_cont = _random_sec_cont(rng, N)
sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
n_sec = rng.integers(0, K_MAX + 1, size=N)
e_sec = rng.uniform(0.0, 50.0, size=N).astype(np.float32)
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries(
sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22}
)
valid_sum = (sec_E * sec_valid).sum(axis=1)
has_secondaries = n_sec > 0
np.testing.assert_allclose(
valid_sum[has_secondaries],
e_sec[has_secondaries],
atol=1e-3,
rtol=1e-5,
)
def test_decode_secondaries_zero_n_sec_has_zero_energy():
"""n_sec=0 rows get no secondaries and no forced energy assignment."""
from giant.data.transforms import decode_secondaries
rng = np.random.default_rng(1)
N = 10
sec_cont = _random_sec_cont(rng, N)
sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
n_sec = np.zeros(N, dtype=np.int64)
e_sec = rng.uniform(1.0, 10.0, size=N).astype(np.float32)
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries(
sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22}
)
assert not sec_valid.any()
np.testing.assert_allclose(sec_E, 0.0)
def test_decode_secondaries_degenerate_row_falls_back_to_even_split():
"""All-zero stick fractions for the valid slots fall back to an even split."""
from giant.data.transforms import decode_secondaries
rng = np.random.default_rng(2)
N = 4
sec_cont = _random_sec_cont(rng, N)
# Drive every valid slot's stick-breaking fraction to ~0 (huge negative logit).
n_sec = np.array([0, 1, 3, K_MAX])
for i, k in enumerate(n_sec):
sec_cont[i, :k, 0] = -80.0
sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
e_sec = np.array([0.0, 4.0, 9.0, 30.0], dtype=np.float32)
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries(
sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22}
)
for i, k in enumerate(n_sec):
if k == 0:
continue
np.testing.assert_allclose(sec_E[i, :k], e_sec[i] / k, atol=1e-4)
np.testing.assert_allclose(sec_E[i, :k].sum(), e_sec[i], atol=1e-3)
def test_decode_secondaries_rescale_preserves_relative_shares():
"""Rescaling should keep each valid slot's *share* of the budget unchanged.
A shortfall shouldn't get dumped into whichever slot is last by energy
rank — it should be spread proportionally, i.e. sec_E[i] / sec_E[j] for
two valid slots must match before and after the e_sec rescale.
"""
from giant.data.transforms import decode_secondaries
rng = np.random.default_rng(3)
N = 1
sec_cont = _random_sec_cont(rng, N)
n_sec = np.array([4])
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
sec_E_small, _, _, sec_valid = decode_secondaries(
sec_cont,
sec_pdg_pred,
n_sec,
np.array([5.0], dtype=np.float32),
pre_dir,
{0: 22},
)
sec_E_large, _, _, _ = decode_secondaries(
sec_cont,
sec_pdg_pred,
n_sec,
np.array([50.0], dtype=np.float32),
pre_dir,
{0: 22},
)
ratio_small = sec_E_small[0, :4] / sec_E_small[0, 0]
ratio_large = sec_E_large[0, :4] / sec_E_large[0, 0]
np.testing.assert_allclose(ratio_small, ratio_large, rtol=1e-4)
tests/test_rollout.py
+92
View File
@@ -1,5 +1,6 @@
"""Tests for the autoregressive shower rollout driver."""
from collections import Counter
from pathlib import Path
from unittest.mock import patch
@@ -161,3 +162,94 @@ def test_max_tracks_cap_conserves_energy():
leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)
assert len(np.unique(rec["track_id"][m])) <= 3
# ── Streaming output (on_chunk) ──────────────────────────────────────────────
def _run_streaming(on_chunk, **kwargs):
torch.manual_seed(0)
np.random.seed(0)
s1, s2 = _models()
cond, tgt = _norms()
seeds = kwargs.pop("seeds", None) or _seeds()
return rollout(
s1,
s2,
_oracle(),
seeds,
cond,
tgt,
PDG_MAP,
MAT_MAP,
energy_cutoff=kwargs.pop("energy_cutoff", 1.0),
max_steps=kwargs.pop("max_steps", 30),
steps=4,
batch_size=128,
max_tracks_per_event=kwargs.pop("max_tracks_per_event", 300),
escape_threshold=kwargs.pop("escape_threshold", 1e9),
on_chunk=on_chunk,
)
def test_on_chunk_receives_every_row_exactly_once():
"""Concatenating the streamed chunks must reproduce the buffered result."""
from giant.rollout import _RECORD_KEYS
buffered = _run()
chunks: list[dict[str, np.ndarray]] = []
summary = _run_streaming(chunks.append)
streamed = {k: np.concatenate([c[k] for c in chunks]) for k in _RECORD_KEYS}
assert summary["n_rows"] == len(buffered["event_id"])
assert len(streamed["event_id"]) == len(buffered["event_id"])
for k in _RECORD_KEYS:
np.testing.assert_array_equal(streamed[k], buffered[k])
def test_on_chunk_summary_termination_reason_counts_match_buffered():
buffered = _run()
summary = _run_streaming(lambda row: None)
expected = Counter(r for r in buffered["termination_reason"].tolist() if r)
assert summary["termination_reason_counts"] == dict(expected)
def test_on_chunk_never_buffers_full_records():
"""Streaming mode must not accumulate rows for later to_dict() retrieval."""
from giant.rollout import _Recorder
rec = _Recorder(sink=lambda row: None)
rec.add(
event_id=np.array([0]),
track_id=np.array([0]),
parent_id=np.array([-1]),
generation=np.array([0]),
step_no=np.array([0]),
pdg=np.array([11]),
pre_x=np.array([0.0]),
pre_y=np.array([0.0]),
pre_z=np.array([0.0]),
pre_E=np.array([1.0]),
pre_dx=np.array([0.0]),
pre_dy=np.array([0.0]),
pre_dz=np.array([1.0]),
post_x=np.array([0.0]),
post_y=np.array([0.0]),
post_z=np.array([1.0]),
post_E=np.array([0.0]),
post_dx=np.array([0.0]),
post_dy=np.array([0.0]),
post_dz=np.array([1.0]),
edep=np.array([1.0]),
step_length=np.array([1.0]),
material=np.array(["G4_AIR"], dtype=object),
layer_id=np.array([0]),
n_sec_pred=np.array([0]),
termination_reason=np.array(["natural_end"], dtype=object),
)
assert rec.n_rows == 1
assert rec.termination_reason_counts == {"natural_end": 1}
with pytest.raises(AssertionError):
rec.to_dict()
tests/test_router.py
+704
View File
@@ -0,0 +1,704 @@
"""Tests for the mixture-of-experts routing prototype (giant/model/network.py)."""
import torch
from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
from giant.model.network import (
ComposedRouter,
DenoisingMLP,
EnergyRouter,
PdgRouter,
ProcessRouter,
ROUTER_REGISTRY,
RoutedDenoisingMLP,
RoutedSecondaryDecoder,
SecondaryDecoder,
build_composed_router,
build_models,
build_router,
)
def _cond(B=8, pdg=3, mat=2):
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.stack(
[torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
)
return cond_cont, cond_cat
def _routed_stage1(n_experts=4, pdg=3, mat=2, **router_kwargs):
router = build_router("energy", n_experts, **router_kwargs)
return RoutedDenoisingMLP(
pdg_vocab=pdg,
mat_vocab=mat,
router=router,
expert_hidden_dim=16,
expert_n_blocks=2,
)
def _routed_sec_decoder(n_experts=4, pdg=3, mat=2, **router_kwargs):
router = build_router("energy", n_experts, **router_kwargs)
return RoutedSecondaryDecoder(
pdg_vocab=pdg,
mat_vocab=mat,
router=router,
expert_hidden_dim=16,
expert_n_blocks=2,
)
# ── Router / EnergyRouter contract ──────────────────────────────────────────
def test_energy_router_registered():
assert ROUTER_REGISTRY["energy"] is EnergyRouter
def test_energy_router_gate_partition_of_unity():
router = EnergyRouter(n_experts=4)
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_energy_router_top1_matches_gate_argmax():
router = EnergyRouter(n_experts=4)
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_energy_router_hardens_as_temperature_shrinks():
"""As tau -> 0 the soft gate should converge to a one-hot at the argmax."""
router = EnergyRouter(n_experts=4, temperature=1e-4)
cond_cont, cond_cat = _cond(16)
g = router.gate(cond_cont, cond_cat)
top1 = router.top1(cond_cont, cond_cat)
onehot = torch.nn.functional.one_hot(top1, num_classes=4).float()
torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0)
def test_energy_router_balance_loss_is_nonnegative_scalar():
router = EnergyRouter(n_experts=4)
cond_cont, cond_cat = _cond(16)
loss = router.balance_loss(cond_cont, cond_cat)
assert loss.shape == ()
assert loss.item() >= 0.0
def test_build_router_ignores_unrecognized_kwargs():
# lambda_balance is a model_config.router key but not an EnergyRouter kwarg
router = build_router("energy", 4, temperature=0.3, lambda_balance=0.5)
assert isinstance(router, EnergyRouter)
assert router.temperature == 0.3
def test_build_router_unknown_type_raises():
try:
build_router("nonexistent", 4)
except ValueError:
return
raise AssertionError("expected ValueError for unknown router type")
# ── PdgRouter ────────────────────────────────────────────────────────────────
def test_pdg_router_registered():
assert ROUTER_REGISTRY["pdg"] is PdgRouter
def test_pdg_router_gate_partition_of_unity():
router = PdgRouter(n_experts=4, pdg_vocab=3)
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_pdg_router_top1_matches_gate_argmax():
router = PdgRouter(n_experts=4, pdg_vocab=3)
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_pdg_router_hardens_as_temperature_shrinks():
"""As tau -> 0 the soft gate should converge to a one-hot at the argmax."""
router = PdgRouter(n_experts=4, pdg_vocab=3, temperature=1e-4)
cond_cont, cond_cat = _cond(16)
g = router.gate(cond_cont, cond_cat)
top1 = router.top1(cond_cont, cond_cat)
onehot = torch.nn.functional.one_hot(top1, num_classes=4).float()
torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0)
def test_pdg_router_balance_loss_is_nonnegative_scalar():
router = PdgRouter(n_experts=4, pdg_vocab=3)
cond_cont, cond_cat = _cond(16)
loss = router.balance_loss(cond_cont, cond_cat)
assert loss.shape == ()
assert loss.item() >= 0.0
def test_pdg_router_classify_loss_defaults_to_zero():
"""PDG is already known at gate time (unlike ProcessRouter's process
label), so no supervision is needed — falls back to Router's default."""
router = PdgRouter(n_experts=4, pdg_vocab=3)
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_pdg_router_only_reads_pdg_column():
"""Gate must depend on cond_cat[:, 0] (pdg) only, not cond_cont or material."""
router = PdgRouter(n_experts=4, pdg_vocab=3)
cond_cont, cond_cat = _cond(16)
g_before = router.gate(cond_cont, cond_cat)
cond_cont_perturbed = torch.randn_like(cond_cont)
cond_cat_diff_mat = cond_cat.clone()
cond_cat_diff_mat[:, 1] = (cond_cat_diff_mat[:, 1] + 1) % 2
g_after = router.gate(cond_cont_perturbed, cond_cat_diff_mat)
torch.testing.assert_close(g_before, g_after, atol=1e-6, rtol=0)
def test_build_router_pdg_type_uses_pdg_vocab():
router = build_router("pdg", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8)
assert isinstance(router, PdgRouter)
assert router.pdg_emb.num_embeddings == 5
def test_build_models_routed_with_pdg_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": "pdg",
"n_experts": 3,
},
)
stage1, sec_decoder = build_models(model_config)
assert isinstance(stage1, RoutedDenoisingMLP)
assert isinstance(stage1.router, PdgRouter)
assert len(stage1.experts) == 3
assert stage1.router.pdg_emb.num_embeddings == 4
# ── 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
# ── ComposedRouter ───────────────────────────────────────────────────────────
def test_composed_router_n_experts_is_product():
router = ComposedRouter(
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
)
assert router.n_experts == 12
def test_composed_router_gate_partition_of_unity():
router = ComposedRouter(
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
)
cond_cont, cond_cat = _cond(16, pdg=5)
g = router.gate(cond_cont, cond_cat)
assert g.shape == (16, 12)
torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
def test_composed_router_gate_is_outer_product_of_sub_gates():
energy_router = EnergyRouter(n_experts=4)
pdg_router = PdgRouter(n_experts=3, pdg_vocab=5)
router = ComposedRouter([energy_router, pdg_router])
cond_cont, cond_cat = _cond(16, pdg=5)
g_energy = energy_router.gate(cond_cont, cond_cat) # (16, 4)
g_pdg = pdg_router.gate(cond_cont, cond_cat) # (16, 3)
expected = (g_energy.unsqueeze(-1) * g_pdg.unsqueeze(1)).flatten(1) # (16, 12)
torch.testing.assert_close(router.gate(cond_cont, cond_cat), expected)
def test_composed_router_top1_factors_into_per_axis_argmax():
"""Joint argmax over the outer product must equal the pair of per-axis
argmaxes, flattened with the same row-major index convention as gate()."""
energy_router = EnergyRouter(n_experts=4)
pdg_router = PdgRouter(n_experts=3, pdg_vocab=5)
router = ComposedRouter([energy_router, pdg_router])
cond_cont, cond_cat = _cond(16, pdg=5)
joint_idx = router.top1(cond_cont, cond_cat)
energy_idx = energy_router.top1(cond_cont, cond_cat)
pdg_idx = pdg_router.top1(cond_cont, cond_cat)
expected = energy_idx * pdg_router.n_experts + pdg_idx
assert torch.equal(joint_idx, expected)
def test_composed_router_supports_different_expert_counts_per_axis():
router = ComposedRouter(
[EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)]
)
assert router.n_experts == 10
cond_cont, cond_cat = _cond(8, pdg=5)
assert router.gate(cond_cont, cond_cat).shape == (8, 10)
def test_composed_router_classify_loss_sums_sub_router_losses():
"""energy/pdg both default to zero, so the composed loss should too."""
router = ComposedRouter(
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
)
cond_cont, cond_cat = _cond(16, pdg=5)
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_composed_router_rejects_empty_router_list():
try:
ComposedRouter([])
except ValueError:
return
raise AssertionError("expected ValueError for empty router list")
def test_composed_router_not_in_registry():
assert "composed" not in ROUTER_REGISTRY
# ── _parse_composed_axes (axis{i}_{field} flat-key config convention) ───────
def test_parse_composed_axes_groups_indexed_keys():
from giant.model.network import _parse_composed_axes
router_cfg = {
"enabled": True,
"type": "composed",
"axis0_type": "energy",
"axis0_n_experts": 4,
"axis0_temperature": 0.3,
"axis1_type": "pdg",
"axis1_n_experts": 3,
"axis1_emb_dim": 6,
}
axes = _parse_composed_axes(router_cfg)
assert axes == [
{"type": "energy", "n_experts": 4, "temperature": 0.3},
{"type": "pdg", "n_experts": 3, "emb_dim": 6},
]
def test_parse_composed_axes_ignores_unrelated_keys():
from giant.model.network import _parse_composed_axes
router_cfg = {
"enabled": True,
"type": "composed",
"lambda_balance": 0.0,
"axis0_type": "energy",
"axis0_n_experts": 4,
}
axes = _parse_composed_axes(router_cfg)
assert axes == [{"type": "energy", "n_experts": 4}]
def test_parse_composed_axes_raises_on_index_gap():
from giant.model.network import _parse_composed_axes
router_cfg = {
"type": "composed",
"axis0_type": "energy",
"axis0_n_experts": 4,
# axis1 missing entirely
"axis2_type": "pdg",
"axis2_n_experts": 3,
}
try:
_parse_composed_axes(router_cfg)
except ValueError:
return
raise AssertionError("expected ValueError for a gap in axis indices")
def test_build_composed_router_resolves_per_axis_specs():
router = build_composed_router(
[
{"type": "energy", "n_experts": 4, "temperature": 0.3},
{"type": "pdg", "n_experts": 3, "emb_dim": 6},
],
pdg_vocab=5,
mat_vocab=2,
)
assert isinstance(router, ComposedRouter)
assert router.n_experts == 12
energy_router, pdg_router = router.routers
assert isinstance(energy_router, EnergyRouter)
assert energy_router.temperature == 0.3
assert isinstance(pdg_router, PdgRouter)
assert pdg_router.pdg_emb.num_embeddings == 5
assert pdg_router.pdg_emb.embedding_dim == 6
def test_build_models_routed_with_composed_router():
model_config = dict(
pdg_vocab=5,
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": "composed",
"axis0_type": "energy",
"axis0_n_experts": 4,
"axis1_type": "pdg",
"axis1_n_experts": 3,
},
)
stage1, sec_decoder = build_models(model_config)
assert isinstance(stage1, RoutedDenoisingMLP)
assert isinstance(stage1.router, ComposedRouter)
assert len(stage1.experts) == 12
assert len(sec_decoder.experts) == 12
# stage1 and sec_decoder must not share router weights (same convention
# as the single-axis routers built by build_models).
assert stage1.router is not sec_decoder.router
def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow():
from giant.sample import sample_flow, sample_secondaries
model_config = dict(
pdg_vocab=3,
mat_vocab=2,
emb_dim=16,
dropout=0.1,
k_max=K_MAX,
expert_hidden_dim=8,
expert_n_blocks=1,
router={
"enabled": True,
"type": "composed",
"axis0_type": "energy",
"axis0_n_experts": 2,
"axis1_type": "pdg",
"axis1_n_experts": 2,
},
)
stage1, sec_decoder = build_models(model_config)
B = 5
cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
assert stage1_norm.shape == (B, X_DIM)
assert n_sec_pred.shape == (B,)
sec_cont, sec_type_emb, sec_valid = sample_secondaries(
sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
)
assert sec_cont.shape == (B, K_MAX, 4)
assert sec_valid.shape == (B, K_MAX)
# ── RoutedDenoisingMLP ───────────────────────────────────────────────────────
def test_routed_denoising_mlp_output_shape_train_and_eval():
B = 8
model = _routed_stage1()
x_t = torch.randn(B, X_DIM)
t = torch.rand(B)
cond_cont, cond_cat = _cond(B)
model.train()
out_train = model(x_t, t, cond_cont, cond_cat)
assert out_train.shape == (B, X_DIM)
model.eval()
with torch.no_grad():
out_eval = model(x_t, t, cond_cont, cond_cat)
assert out_eval.shape == (B, X_DIM)
def test_routed_denoising_mlp_gradients_flow_in_train_mode():
"""Soft mixture in train mode should touch every expert's parameters."""
B = 8
model = _routed_stage1(n_experts=3)
x_t = torch.randn(B, X_DIM)
t = torch.rand(B)
cond_cont, cond_cat = _cond(B)
model.train()
flow_loss = model(x_t, t, cond_cont, cond_cat).sum()
nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum()
(flow_loss + nsec_loss).backward()
for name, p in model.named_parameters():
assert p.grad is not None, f"no grad for {name}"
def test_routed_denoising_mlp_eval_dispatch_matches_manual_grouping():
"""Eval-mode grouped top-1 dispatch must equal running each row through
its assigned expert individually (batch order shouldn't matter)."""
B = 12
model = _routed_stage1(n_experts=4)
model.eval()
x_t = torch.randn(B, X_DIM)
t = torch.rand(B)
cond_cont, cond_cat = _cond(B)
with torch.no_grad():
batched = model(x_t, t, cond_cont, cond_cat)
t_emb = model.time_emb(t)
c_emb = model.cond_enc(cond_cont, cond_cat)
cond = torch.cat([t_emb, c_emb], dim=-1)
idx = model.router.top1(cond_cont, cond_cat)
manual = torch.zeros_like(x_t)
for i in range(B):
manual[i] = model.experts[int(idx[i])](x_t[i : i + 1], cond[i : i + 1])[0]
torch.testing.assert_close(batched, manual, atol=1e-5, rtol=1e-4)
def test_routed_denoising_mlp_predict_n_sec_shape():
B = 6
model = _routed_stage1()
cond_cont, cond_cat = _cond(B)
logits = model.predict_n_sec(cond_cont, cond_cat)
assert logits.shape == (B, K_MAX + 1)
def test_routed_denoising_mlp_pdg_embedding_weight_shape():
model = _routed_stage1(pdg=5, mat=2)
from giant.constants import EMB_DIM
assert model.pdg_embedding_weight().shape == (5, EMB_DIM)
# ── RoutedSecondaryDecoder ───────────────────────────────────────────────────
def test_routed_secondary_decoder_output_shape_train_and_eval():
B = 8
decoder = _routed_sec_decoder()
x_t = torch.randn(B, SEC_DIM)
t = torch.rand(B)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
decoder.train()
out_train = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
assert out_train.shape == (B, SEC_DIM)
decoder.eval()
with torch.no_grad():
out_eval = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
assert out_eval.shape == (B, SEC_DIM)
def test_routed_secondary_decoder_gradients_flow():
B = 4
decoder = _routed_sec_decoder(n_experts=3)
x_t = torch.randn(B, SEC_DIM)
t = torch.rand(B)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
decoder.train()
decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward()
for name, p in decoder.named_parameters():
assert p.grad is not None, f"no grad for {name}"
# ── build_models dispatch ────────────────────────────────────────────────────
def test_build_models_monolith_when_router_absent():
model_config = dict(
pdg_vocab=4,
mat_vocab=2,
hidden_dim=32,
n_blocks=2,
emb_dim=16,
dropout=0.1,
k_max=K_MAX,
)
stage1, sec_decoder = build_models(model_config)
assert isinstance(stage1, DenoisingMLP)
assert isinstance(sec_decoder, SecondaryDecoder)
def test_build_models_monolith_when_router_disabled():
model_config = dict(
pdg_vocab=4,
mat_vocab=2,
hidden_dim=32,
n_blocks=2,
emb_dim=16,
dropout=0.1,
k_max=K_MAX,
router={"enabled": False, "type": "energy", "n_experts": 4},
)
stage1, sec_decoder = build_models(model_config)
assert isinstance(stage1, DenoisingMLP)
assert isinstance(sec_decoder, SecondaryDecoder)
def test_build_models_routed_when_enabled():
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": "energy",
"n_experts": 4,
"temperature": 0.5,
"learn_centers": True,
"lambda_balance": 0.0,
},
)
stage1, sec_decoder = build_models(model_config)
assert isinstance(stage1, RoutedDenoisingMLP)
assert isinstance(sec_decoder, RoutedSecondaryDecoder)
assert len(stage1.experts) == 4
assert len(sec_decoder.experts) == 4
def test_build_models_routed_pair_is_drop_in_for_sample_flow():
"""Exercise the exact calling convention giant/sample.py uses."""
from giant.sample import sample_flow, sample_secondaries
model_config = dict(
pdg_vocab=3,
mat_vocab=2,
emb_dim=16,
dropout=0.1,
k_max=K_MAX,
expert_hidden_dim=8,
expert_n_blocks=1,
router={"enabled": True, "type": "energy", "n_experts": 2},
)
stage1, sec_decoder = build_models(model_config)
B = 5
cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
assert stage1_norm.shape == (B, X_DIM)
assert n_sec_pred.shape == (B,)
sec_cont, sec_type_emb, sec_valid = sample_secondaries(
sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
)
assert sec_cont.shape == (B, K_MAX, 4)
assert sec_valid.shape == (B, K_MAX)
tests/test_steps_to_parquet.py
+3 -1
View File
@@ -69,7 +69,9 @@ def test_orphaned_child_track_is_dropped_not_nulled():
}
)
out, n_orphaned = steps_to_parquet._add_secondary_attributes(df)
row = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0))
row = out.filter(
(pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0)
)
assert n_orphaned == 1
assert row["child_track_ids"].to_list() == [[2]]
tests/test_transforms.py
+39 -8
View File
@@ -235,25 +235,22 @@ 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 _step_data_no_sec_lists(n_sec: np.ndarray) -> dict:
"""Minimal build_features input with n_sec but no per-secondary list columns
(mimics a parquet that skipped the parent->child join)."""
N = len(n_sec)
def _minimal_step_data(N: int, process: np.ndarray | None = None) -> dict:
rng = np.random.default_rng(0)
return {
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.asarray(n_sec, 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),
@@ -261,6 +258,40 @@ def _step_data_no_sec_lists(n_sec: np.ndarray) -> dict:
"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 _step_data_no_sec_lists(n_sec: np.ndarray) -> dict:
"""Minimal build_features input with n_sec but no per-secondary list columns
(mimics a parquet that skipped the parent->child join)."""
data = _minimal_step_data(len(n_sec))
data["n_sec"] = np.asarray(n_sec, dtype=np.int32)
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])
def test_build_features_require_secondaries_raises_when_lists_missing():
@@ -281,7 +312,7 @@ def test_build_features_require_secondaries_ok_when_no_secondaries():
data = _step_data_no_sec_lists(np.zeros(3, dtype=np.int32))
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
_, _, _, _, sec_cont, sec_pdg_idx, _, _ = build_features(
_, _, _, _, sec_cont, sec_pdg_idx, *_ = build_features(
data, pdg_map, mat_map, require_secondaries=True
)