diff --git a/giant/cli.py b/giant/cli.py index d212bea..935bc98 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -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}") diff --git a/giant/config.py b/giant/config.py index 7001151..4834f3c 100644 --- a/giant/config.py +++ b/giant/config.py @@ -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)) diff --git a/giant/data/dataset.py b/giant/data/dataset.py index bbac493..f319d93 100644 --- a/giant/data/dataset.py +++ b/giant/data/dataset.py @@ -36,7 +36,7 @@ class StreamingStepsDataset(IterableDataset): numpy slicing instead of a per-row Python loop in the default collate. Each batch is a tuple: - (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx) + (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx) where: cond_cont: (B, COND_DIM) float32 cond_cat: (B, 2) int64 @@ -44,6 +44,8 @@ class StreamingStepsDataset(IterableDataset): n_sec: (B,) int64 — true secondary count per step sec_cont: (B, K_MAX, 4) float32 — [stick_logit, local_dir] per slot sec_pdg_idx: (B, K_MAX) int64 — PDG model-index per secondary slot + proc_idx: (B,) int64 — process-class label (ProcessRouter supervision + only; zeros when `proc_map` is None) """ def __init__( @@ -57,6 +59,7 @@ class StreamingStepsDataset(IterableDataset): batch_size: int, shuffle_buffer: int = 65536, shuffle: bool = True, + proc_map: dict[str, int] | None = None, ) -> None: self.files = list(files) self.split_events = split_events @@ -68,6 +71,7 @@ class StreamingStepsDataset(IterableDataset): self.batch_size = batch_size self.shuffle_buffer = max(shuffle_buffer, batch_size) self.shuffle = shuffle + self.proc_map = proc_map def __iter__(self): worker_info = torch.utils.data.get_worker_info() @@ -85,6 +89,7 @@ class StreamingStepsDataset(IterableDataset): buf_nsec: list[np.ndarray] = [] buf_sec: list[np.ndarray] = [] buf_spdg: list[np.ndarray] = [] + buf_proc: list[np.ndarray] = [] buf_n = 0 for path in files: @@ -94,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, ) diff --git a/giant/data/loader.py b/giant/data/loader.py index b51fd69..4ddbf7f 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -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 diff --git a/giant/data/transforms.py b/giant/data/transforms.py index ee859eb..f7d647e 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -453,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, @@ -461,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 @@ -570,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, @@ -577,6 +588,7 @@ def build_features( n_sec, sec_cont, sec_pdg_idx, + proc_idx, cond_normalizer, target_normalizer, ) diff --git a/giant/model/network.py b/giant/model/network.py index d341630..57d25df 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -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 diff --git a/giant/pipeline.py b/giant/pipeline.py index c542584..d6522c7 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -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, ) diff --git a/giant/rollout.py b/giant/rollout.py index 109faa0..228c3c5 100644 --- a/giant/rollout.py +++ b/giant/rollout.py @@ -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() diff --git a/giant/train.py b/giant/train.py index c21deb7..d796956 100644 --- a/giant/train.py +++ b/giant/train.py @@ -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 diff --git a/giant/validate.py b/giant/validate.py index 959e375..3372a06 100644 --- a/giant/validate.py +++ b/giant/validate.py @@ -85,8 +85,8 @@ def validate_marginals( for i, batch in enumerate(val_loader): if n_batches is not None and i >= n_batches: break - # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx). - cond_cont, cond_cat, x1, n_sec, sec_cont, sec_pdg_idx = batch + # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx). + cond_cont, cond_cat, x1, n_sec, sec_cont, sec_pdg_idx, _proc_idx = batch cond_cont = cond_cont.to(device) cond_cat = cond_cat.to(device) diff --git a/scripts/steps_to_parquet_parallel.py b/scripts/steps_to_parquet_parallel.py index dc1c5fe..b2ee41e 100644 --- a/scripts/steps_to_parquet_parallel.py +++ b/scripts/steps_to_parquet_parallel.py @@ -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.") diff --git a/tests/test_loader.py b/tests/test_loader.py index 6c562f3..af609dc 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -1,6 +1,7 @@ +import pandas as pd import pytest -from giant.data.loader import find_parquet_files +from giant.data.loader import build_process_map_from_files, find_parquet_files def _touch(path): @@ -58,3 +59,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 diff --git a/tests/test_rollout.py b/tests/test_rollout.py index f744755..4d7c747 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -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() diff --git a/tests/test_router.py b/tests/test_router.py new file mode 100644 index 0000000..31f6b54 --- /dev/null +++ b/tests/test_router.py @@ -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) diff --git a/tests/test_steps_to_parquet.py b/tests/test_steps_to_parquet.py index 9810d6e..8e839c1 100644 --- a/tests/test_steps_to_parquet.py +++ b/tests/test_steps_to_parquet.py @@ -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]] diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 45cfc04..c95a36b 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -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 )