Merge branch '4-prototype-a-mixture-of-experts-routing-tree-architecture' into 'master'

Add mixture-of-experts routing prototype for Stage 1 and Stage 2

Closes #4

See merge request lbogner/giant!3

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}")
 
 

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))
 

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,
         )

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

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,
     )

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

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,
     )

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()
 
 

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
 

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)
 

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.")

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

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()

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)

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]]

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
     )