lars
06c9ad8e5f
- make_seed_frontier only resolves particle mass/charge in "physical" mode, so "embedding"-mode rollouts no longer crash on a seed PDG code giant.particles can't resolve (the TERM_UNKNOWN_PDG gate now handles it). - nearest_known_pdg skips unresolvable candidate PDG codes instead of raising and killing the whole rollout/predict run. - predict/rollout fail with a clear message when a checkpoint predates the sec_phys normalizer, instead of a bare KeyError. - validate_marginals' phys_kl degrades to NaN (matching the energy_fraction_kl pattern) instead of crashing when a validated batch has zero secondaries on either side. - Correct CLAUDE.md's stale claim that the materials table is unfilled. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
1059 lines
36 KiB
Python
1059 lines
36 KiB
Python
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
|
|
|
|
import numpy as np
|
|
import yaml
|
|
import torch
|
|
import typer
|
|
from typing_extensions import Annotated
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
from tqdm import tqdm
|
|
|
|
from giant import config as gconfig
|
|
from giant.constants import (
|
|
LOCAL_TARGET_NAMES,
|
|
PREDICT_COORD_METADATA_KEY,
|
|
PREDICT_SCHEMA_VERSION,
|
|
PREDICT_SCHEMA_VERSION_KEY,
|
|
ROLLOUT_COORD_VALUE,
|
|
)
|
|
from giant.data.loader import (
|
|
find_parquet_files,
|
|
iter_file_chunks,
|
|
iter_cond_chunks,
|
|
)
|
|
from giant.data.transforms import (
|
|
build_features,
|
|
build_cond_features,
|
|
decode_secondaries,
|
|
energy_simplex_decode,
|
|
inv_local_frame_rotation,
|
|
inv_log_transform,
|
|
reconstruct_post_pos,
|
|
Normalizer,
|
|
)
|
|
from giant.geometry import GeometryOracle
|
|
from giant.model.network import build_models
|
|
from giant.particles import nearest_known_pdg
|
|
from giant.pipeline import run_train_job
|
|
from giant.rollout import rollout as run_rollout
|
|
from giant.sample import sample_flow, sample_secondaries
|
|
|
|
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")
|
|
|
|
|
|
def _resolve_prediction_output(data: Path, out: Path | None) -> tuple[Path, Path, str]:
|
|
"""Return (out_path, resolved_dataset_path, pred_uuid).
|
|
|
|
When *out* is None the output path is derived from *data*:
|
|
- under /ceph/ → fixed central store with a UUID filename
|
|
- elsewhere → sibling of *data* with a UUID filename
|
|
"""
|
|
dataset_path = data.resolve()
|
|
pred_uuid = str(uuid_mod.uuid4())
|
|
if out is None:
|
|
if str(dataset_path).startswith("/ceph/"):
|
|
out = _CEPH_PREDICTIONS / f"{pred_uuid}.parquet"
|
|
else:
|
|
out = data.parent / f"{pred_uuid}.parquet"
|
|
return out, dataset_path, pred_uuid
|
|
|
|
|
|
def _write_prediction_ref(
|
|
checkpoint: Path,
|
|
pred_uuid: str,
|
|
out: Path,
|
|
dataset_path: Path,
|
|
comment: str | None = None,
|
|
) -> Path:
|
|
"""Write a YAML sidecar in the checkpoint directory and return its path."""
|
|
ref = {
|
|
"prediction_id": pred_uuid,
|
|
"output": str(out),
|
|
"dataset": str(dataset_path),
|
|
"checkpoint": str(checkpoint.resolve()),
|
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
|
}
|
|
if comment is not None:
|
|
ref["comment"] = comment
|
|
ref_path = checkpoint.parent / f"{pred_uuid}.yaml"
|
|
ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False))
|
|
return ref_path
|
|
|
|
|
|
@app.callback()
|
|
def _main() -> None:
|
|
"""GIANT — Geant4 step-function surrogate."""
|
|
|
|
|
|
class Mode(str, Enum):
|
|
flow = "flow"
|
|
ddpm = "ddpm"
|
|
|
|
|
|
class Conditioning(str, Enum):
|
|
physical = "physical"
|
|
embedding = "embedding"
|
|
|
|
|
|
class Coord(str, Enum):
|
|
global_ = "global"
|
|
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[
|
|
Path, typer.Argument(help="Parquet file or directory of parquet files")
|
|
],
|
|
config: Annotated[
|
|
Optional[Path],
|
|
typer.Option(
|
|
"--config", "-c", help="TOML config file (overridden by explicit flags)"
|
|
),
|
|
] = None,
|
|
mode: Annotated[
|
|
Optional[Mode],
|
|
typer.Option("--mode", "-m", help="Generative model: flow matching or DDPM"),
|
|
] = None,
|
|
epochs: Annotated[Optional[int], typer.Option("--epochs", "-e")] = None,
|
|
batch_size: Annotated[
|
|
Optional[str],
|
|
typer.Option(
|
|
"--batch-size",
|
|
"-b",
|
|
help="Integer, or 'auto' to estimate from free GPU memory "
|
|
"(cuda devices only)",
|
|
),
|
|
] = 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,
|
|
hidden_dim: Annotated[Optional[int], typer.Option("--hidden-dim", "-H")] = None,
|
|
n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None,
|
|
emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None,
|
|
dropout: Annotated[
|
|
Optional[float],
|
|
typer.Option(
|
|
"--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"
|
|
),
|
|
] = None,
|
|
conditioning: Annotated[
|
|
Optional[Conditioning],
|
|
typer.Option(
|
|
"--conditioning",
|
|
help="Input conditioning: continuous physical properties "
|
|
"(mass/charge/Z_eff/A_eff/density/X0/lambda_int, default) or the "
|
|
"original learned PDG/material embeddings",
|
|
),
|
|
] = 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,
|
|
seed: Annotated[
|
|
Optional[int],
|
|
typer.Option("--seed", "-s", help="Random seed for reproducibility"),
|
|
] = None,
|
|
validate_every: Annotated[
|
|
Optional[int],
|
|
typer.Option(
|
|
"--validate-every",
|
|
"-v",
|
|
help="Run marginal+KL validation every N epochs (0 disables)",
|
|
),
|
|
] = None,
|
|
validate_steps: Annotated[
|
|
Optional[int],
|
|
typer.Option(
|
|
"--validate-steps",
|
|
"-t",
|
|
help="Flow matching ODE steps used during marginal validation "
|
|
"(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(
|
|
"--shuffle-buffer", "-B", help="Rows held in RAM per worker for shuffling"
|
|
),
|
|
] = 65536,
|
|
out: Annotated[
|
|
Optional[Path],
|
|
typer.Option(
|
|
"--out", "-o", help="Checkpoint dir (default: auto from hyperparams)"
|
|
),
|
|
] = None,
|
|
device: Annotated[
|
|
Optional[str],
|
|
typer.Option("--device", "-D", help="cpu | cuda | mps (default: auto)"),
|
|
] = None,
|
|
num_workers: Annotated[Optional[int], typer.Option("--num-workers", "-j")] = None,
|
|
resume: Annotated[
|
|
Optional[Path],
|
|
typer.Option("--resume", "-r", help="Checkpoint .pt to resume training from"),
|
|
] = None,
|
|
) -> None:
|
|
"""Train the GIANT surrogate model."""
|
|
batch_size_auto = False
|
|
batch_size_value: Optional[int] = None
|
|
if batch_size is not None:
|
|
if batch_size.strip().lower() == "auto":
|
|
batch_size_auto = True
|
|
else:
|
|
try:
|
|
batch_size_value = int(batch_size)
|
|
except ValueError:
|
|
typer.echo(
|
|
f"error: --batch-size must be an integer or 'auto', "
|
|
f"got {batch_size!r}",
|
|
err=True,
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
cli_train = {
|
|
k: v
|
|
for k, v in {
|
|
"mode": mode.value if mode is not None else None,
|
|
"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: dict[str, object] = {
|
|
k: v
|
|
for k, v in {
|
|
"hidden_dim": hidden_dim,
|
|
"n_blocks": n_blocks,
|
|
"emb_dim": emb_dim,
|
|
"dropout": dropout,
|
|
"conditioning": conditioning.value if conditioning is not None else None,
|
|
}.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
|
|
)
|
|
t, m = cfg["train"], cfg["model"]
|
|
|
|
_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(
|
|
est_hidden_dim, est_n_blocks, _device
|
|
)
|
|
except ValueError as exc:
|
|
typer.echo(f"error: {exc}", err=True)
|
|
raise typer.Exit(1)
|
|
typer.echo(
|
|
f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)"
|
|
)
|
|
|
|
out_dir = out or Path(
|
|
f"checkpoints/{date.today().strftime('%Y%m%d')}"
|
|
f"_{t['mode']}"
|
|
f"_h{m['hidden_dim']}"
|
|
f"_b{m['n_blocks']}"
|
|
f"_e{m['emb_dim']}"
|
|
f"_c{m['conditioning']}"
|
|
f"_lr{t['lr']}"
|
|
f"_bs{t['batch_size']}"
|
|
)
|
|
|
|
typer.echo(f"device: {_device}")
|
|
typer.echo(f"out_dir: {out_dir}")
|
|
|
|
run_train_job(
|
|
data=data,
|
|
cfg=cfg,
|
|
out_dir=out_dir,
|
|
device=_device,
|
|
shuffle_buffer=shuffle_buffer,
|
|
num_workers=t["num_workers"],
|
|
resume=resume,
|
|
echo=typer.echo,
|
|
)
|
|
|
|
|
|
@app.command()
|
|
def predict(
|
|
data: Annotated[
|
|
Path, typer.Argument(help="Parquet file or directory of parquet files")
|
|
],
|
|
checkpoint: Annotated[
|
|
Path,
|
|
typer.Option(
|
|
"--checkpoint",
|
|
"-c",
|
|
help="Path to checkpoint .pt file (best.pt or last.pt)",
|
|
),
|
|
],
|
|
coord: Annotated[
|
|
Coord,
|
|
typer.Option(
|
|
"--coord",
|
|
"-C",
|
|
help="global: full physical units, world frame (default). "
|
|
"local: raw 9D model output (denormalised only, local frame, "
|
|
"log-scaled scalars) alongside the matching ground-truth target "
|
|
"for the same input file — requires post-step columns.",
|
|
),
|
|
] = Coord.global_,
|
|
out: Annotated[
|
|
Optional[Path],
|
|
typer.Option(
|
|
"--out",
|
|
"-o",
|
|
help="Output parquet path (default: <data>_predicted[_local].parquet)",
|
|
),
|
|
] = None,
|
|
batch_size: Annotated[
|
|
str,
|
|
typer.Option(
|
|
"--batch-size",
|
|
"-b",
|
|
help="Inference batch size, or 'auto' to estimate from free GPU "
|
|
"memory (cuda devices only)",
|
|
),
|
|
] = "4096",
|
|
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)"),
|
|
] = None,
|
|
comment: Annotated[
|
|
Optional[str],
|
|
typer.Option(
|
|
"--comment",
|
|
"-m",
|
|
help="Free-text note recorded in the prediction's YAML sidecar",
|
|
),
|
|
] = None,
|
|
) -> None:
|
|
"""Run trained model on a parquet file and save predictions."""
|
|
batch_size_auto = False
|
|
batch_size_value: Optional[int] = None
|
|
if batch_size.strip().lower() == "auto":
|
|
batch_size_auto = True
|
|
else:
|
|
try:
|
|
batch_size_value = int(batch_size)
|
|
except ValueError:
|
|
typer.echo(
|
|
f"error: --batch-size must be an integer or 'auto', got {batch_size!r}",
|
|
err=True,
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
_device = torch.device(device) if device else gconfig.auto_device()
|
|
typer.echo(f"device: {_device}")
|
|
|
|
# --- Load checkpoint ---
|
|
ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
|
|
if "model_config" not in ckpt:
|
|
typer.echo(
|
|
"error: checkpoint has no model_config — retrain with the current code",
|
|
err=True,
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
if "sec_decoder" not in ckpt:
|
|
typer.echo(
|
|
"error: checkpoint has no sec_decoder — retrain with the current code",
|
|
err=True,
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
if "sec_phys" not in ckpt.get("normalizer", {}):
|
|
typer.echo(
|
|
"error: checkpoint has no normalizer.sec_phys — retrain with the "
|
|
"current code",
|
|
err=True,
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
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(
|
|
est_hidden_dim,
|
|
est_n_blocks,
|
|
_device,
|
|
training=False,
|
|
)
|
|
except ValueError as exc:
|
|
typer.echo(f"error: {exc}", err=True)
|
|
raise typer.Exit(1)
|
|
typer.echo(
|
|
f"batch_size: {batch_size_value} (auto-estimated from free GPU memory)"
|
|
)
|
|
|
|
assert batch_size_value is not None
|
|
bs = batch_size_value
|
|
conditioning = model_cfg.get("conditioning", "embedding")
|
|
pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
|
|
mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
|
|
cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
|
|
tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
|
|
sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"])
|
|
|
|
model, sec_decoder = build_models(model_cfg)
|
|
_load_model_weights(model, sec_decoder, ckpt, weights, checkpoint)
|
|
model.to(_device).eval()
|
|
sec_decoder.to(_device).eval()
|
|
|
|
typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})")
|
|
gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
|
|
|
|
# --- Output path ---
|
|
out, dataset_path, pred_uuid = _resolve_prediction_output(data, out)
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
typer.echo(f"output: {out}")
|
|
|
|
# --- Stream input, generate predictions, write output ---
|
|
files = find_parquet_files(data)
|
|
typer.echo(f"found {len(files)} parquet file(s)")
|
|
|
|
writer: pq.ParquetWriter | None = None
|
|
total = 0
|
|
skipped = 0
|
|
unknown_pdg_counts: Counter[int] = Counter()
|
|
total_rows = sum(pq.ParquetFile(path).metadata.num_rows for path in files)
|
|
chunk_iter = iter_file_chunks if coord == Coord.local else iter_cond_chunks
|
|
|
|
def _concat(
|
|
a: dict[str, np.ndarray], b: dict[str, np.ndarray]
|
|
) -> dict[str, np.ndarray]:
|
|
return {k: np.concatenate([a[k], b[k]], axis=0) for k in a}
|
|
|
|
def _process(piece: dict[str, np.ndarray]) -> None:
|
|
nonlocal writer, total
|
|
|
|
if coord == Coord.local:
|
|
cond_cont, cond_cat, target_raw, _, _, _, _, _ = build_features(
|
|
piece, pdg_map, mat_map, conditioning=conditioning
|
|
)
|
|
cond_cont = cond_norm.transform(cond_cont)
|
|
else:
|
|
cond_cont, cond_cat = build_cond_features(
|
|
piece, pdg_map, mat_map, cond_norm, conditioning=conditioning
|
|
)
|
|
|
|
cc = torch.from_numpy(cond_cont).float().to(_device)
|
|
ck = torch.from_numpy(cond_cat).long().to(_device)
|
|
stage1_norm, n_sec_pred = sample_flow(model, cc, ck, steps=steps)
|
|
|
|
if coord == Coord.global_:
|
|
sec_cont, sec_phys, _sec_valid_pred = sample_secondaries(
|
|
sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
|
|
)
|
|
sec_full_np = torch.cat([sec_cont, sec_phys], dim=-1).cpu().numpy()
|
|
|
|
n_sec_pred_np = n_sec_pred.cpu().numpy()
|
|
pred = stage1_norm.cpu().numpy() # normalised
|
|
|
|
# Inverse-normalise → local frame, log-scaled scalars
|
|
raw = tgt_norm.inverse_transform(pred)
|
|
|
|
if coord == Coord.local:
|
|
table = pa.table(
|
|
{
|
|
"event_id": piece["event_id"],
|
|
"pdg": piece["pdg"],
|
|
"pre_x": piece["pre_pos"][:, 0],
|
|
"pre_y": piece["pre_pos"][:, 1],
|
|
"pre_z": piece["pre_pos"][:, 2],
|
|
"pre_E": piece["pre_E"],
|
|
"pre_dx": piece["pre_dir"][:, 0],
|
|
"pre_dy": piece["pre_dir"][:, 1],
|
|
"pre_dz": piece["pre_dir"][:, 2],
|
|
"material": piece["material"],
|
|
"layer_id": piece["layer_id"],
|
|
"n_sec": piece["n_sec"],
|
|
**{
|
|
f"pred_{name}": raw[:, j]
|
|
for j, name in enumerate(LOCAL_TARGET_NAMES)
|
|
},
|
|
**{
|
|
f"true_{name}": target_raw[:, j]
|
|
for j, name in enumerate(LOCAL_TARGET_NAMES)
|
|
},
|
|
}
|
|
)
|
|
else:
|
|
step_length = inv_log_transform(raw[:, 0])
|
|
# Columns 1:3 are ALR coords of the deposit/secondary/post energy
|
|
# simplex; decode them against pre_E so edep + e_sec + post_E == pre_E
|
|
# (hence delta_e == edep + e_sec) holds by construction. e_sec_pred
|
|
# doubles as the stick-breaking energy budget for the Stage-2 decode
|
|
# below, since the model has no other source for it at inference.
|
|
edep, e_sec_pred, _post_E, delta_e = energy_simplex_decode(
|
|
raw[:, 1:3], piece["pre_E"]
|
|
)
|
|
|
|
# Normalise predicted direction then rotate back to world frame
|
|
post_dir_local = raw[:, 3:6].copy()
|
|
norms = np.linalg.norm(post_dir_local, axis=1, keepdims=True)
|
|
post_dir_local /= np.where(norms < 1e-8, 1.0, norms)
|
|
post_dir_world = inv_local_frame_rotation(piece["pre_dir"], post_dir_local)
|
|
|
|
# Same for the travel direction, then reconstruct post_pos from
|
|
# the single shared step_length so the two stay consistent.
|
|
travel_dir_local = raw[:, 6:9].copy()
|
|
norms = np.linalg.norm(travel_dir_local, axis=1, keepdims=True)
|
|
travel_dir_local /= np.where(norms < 1e-8, 1.0, norms)
|
|
post_pos_world = reconstruct_post_pos(
|
|
piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local
|
|
)
|
|
|
|
sec_E, sec_dir_world, sec_mass, sec_charge, _sec_valid = decode_secondaries(
|
|
sec_full_np,
|
|
n_sec_pred_np,
|
|
e_sec_pred,
|
|
piece["pre_dir"],
|
|
sec_phys_normalizer=sec_phys_norm,
|
|
)
|
|
# Reporting-only nearest-known-PDG label (never fed back into the
|
|
# model) for the sec_pdg_list output column — see
|
|
# giant/particles.py and the "no snapping at inference" design.
|
|
sec_pdg_code = nearest_known_pdg(
|
|
sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()
|
|
).reshape(sec_mass.shape)
|
|
sec_pdg_list = [
|
|
sec_pdg_code[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)
|
|
]
|
|
sec_E_list = [sec_E[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)]
|
|
sec_dx_list = [
|
|
sec_dir_world[i, :n, 0].tolist() for i, n in enumerate(n_sec_pred_np)
|
|
]
|
|
sec_dy_list = [
|
|
sec_dir_world[i, :n, 1].tolist() for i, n in enumerate(n_sec_pred_np)
|
|
]
|
|
sec_dz_list = [
|
|
sec_dir_world[i, :n, 2].tolist() for i, n in enumerate(n_sec_pred_np)
|
|
]
|
|
|
|
table = pa.table(
|
|
{
|
|
"event_id": piece["event_id"],
|
|
"pdg": piece["pdg"],
|
|
"pre_x": piece["pre_pos"][:, 0],
|
|
"pre_y": piece["pre_pos"][:, 1],
|
|
"pre_z": piece["pre_pos"][:, 2],
|
|
"pre_E": piece["pre_E"],
|
|
"pre_dx": piece["pre_dir"][:, 0],
|
|
"pre_dy": piece["pre_dir"][:, 1],
|
|
"pre_dz": piece["pre_dir"][:, 2],
|
|
"material": piece["material"],
|
|
"layer_id": piece["layer_id"],
|
|
"n_sec": piece["n_sec"],
|
|
"n_sec_pred": n_sec_pred_np,
|
|
"step_length": step_length,
|
|
"delta_e": delta_e,
|
|
"edep": edep,
|
|
"post_dx": post_dir_world[:, 0],
|
|
"post_dy": post_dir_world[:, 1],
|
|
"post_dz": post_dir_world[:, 2],
|
|
"post_x": post_pos_world[:, 0],
|
|
"post_y": post_pos_world[:, 1],
|
|
"post_z": post_pos_world[:, 2],
|
|
"sec_pdg_list": sec_pdg_list,
|
|
"sec_E_list": sec_E_list,
|
|
"sec_dx_list": sec_dx_list,
|
|
"sec_dy_list": sec_dy_list,
|
|
"sec_dz_list": sec_dz_list,
|
|
}
|
|
)
|
|
|
|
table = table.replace_schema_metadata(
|
|
{
|
|
PREDICT_COORD_METADATA_KEY: coord.value,
|
|
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
|
|
}
|
|
)
|
|
|
|
if writer is None:
|
|
writer = pq.ParquetWriter(out, table.schema)
|
|
writer.write_table(table)
|
|
total += len(piece["event_id"])
|
|
|
|
# Buffer rows across row-group boundaries so the inference batch size
|
|
# isn't capped by however the source file happens to be chunked.
|
|
buffer: dict[str, np.ndarray] | None = None
|
|
|
|
bar = tqdm(total=total_rows, desc="predict", unit="row", dynamic_ncols=True)
|
|
for path in files:
|
|
for chunk in chunk_iter(path):
|
|
N_in = len(chunk["event_id"])
|
|
|
|
pdg_mask = np.array([int(p) in pdg_map for p in chunk["pdg"]])
|
|
if not pdg_mask.all():
|
|
unknown_pdg_counts.update(int(p) for p in chunk["pdg"][~pdg_mask])
|
|
chunk = {k: v[pdg_mask] for k, v in chunk.items()}
|
|
|
|
skipped += N_in - len(chunk["event_id"])
|
|
bar.update(N_in)
|
|
if len(chunk["event_id"]) == 0:
|
|
continue
|
|
|
|
buffer = chunk if buffer is None else _concat(buffer, chunk)
|
|
while len(buffer["event_id"]) >= bs:
|
|
piece = {k: v[:bs] for k, v in buffer.items()}
|
|
buffer = {k: v[bs:] for k, v in buffer.items()}
|
|
_process(piece)
|
|
|
|
if buffer is not None and len(buffer["event_id"]) > 0:
|
|
_process(buffer)
|
|
|
|
bar.close()
|
|
if writer is not None:
|
|
writer.close()
|
|
|
|
ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path, comment)
|
|
typer.echo(f"reference: {ref_path}")
|
|
|
|
if skipped:
|
|
codes = ", ".join(
|
|
f"{pdg} ({count})" for pdg, count in sorted(unknown_pdg_counts.items())
|
|
)
|
|
typer.echo(
|
|
f"warning: skipped {skipped:,} row(s) with unknown PDG code(s): {codes}",
|
|
err=True,
|
|
)
|
|
typer.echo(f"wrote {total:,} rows → {out}")
|
|
|
|
|
|
def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.ndarray]:
|
|
"""Pick each event's primary entry state (argmax-pre_E row) as a shower seed.
|
|
|
|
Streams conditioning columns and keeps the highest-pre_E step per event_id —
|
|
the codebase's convention for the primary (a secondary always carries less
|
|
energy than its parent). See giant/analysis.py:_entry_axis_and_bin_edges.
|
|
"""
|
|
best_E: dict[int, float] = {}
|
|
best: dict[int, tuple] = {}
|
|
for path in files:
|
|
for chunk in iter_cond_chunks(path):
|
|
ev = chunk["event_id"]
|
|
pe = chunk["pre_E"]
|
|
for i in range(len(ev)):
|
|
e = int(ev[i])
|
|
if pe[i] > best_E.get(e, -np.inf):
|
|
best_E[e] = float(pe[i])
|
|
best[e] = (
|
|
int(chunk["pdg"][i]),
|
|
chunk["pre_pos"][i].astype(np.float64),
|
|
float(pe[i]),
|
|
chunk["pre_dir"][i].astype(np.float64),
|
|
)
|
|
event_ids = sorted(best)
|
|
if n_events is not None:
|
|
event_ids = event_ids[:n_events]
|
|
if not event_ids:
|
|
raise ValueError("no events found to seed from")
|
|
|
|
return {
|
|
"event_id": np.array(event_ids, dtype=np.int64),
|
|
"pdg": np.array([best[e][0] for e in event_ids], dtype=np.int64),
|
|
"pre_pos": np.stack([best[e][1] for e in event_ids]),
|
|
"pre_E": np.array([best[e][2] for e in event_ids], dtype=np.float64),
|
|
"pre_dir": np.stack([best[e][3] for e in event_ids]),
|
|
}
|
|
|
|
|
|
@app.command()
|
|
def rollout(
|
|
data: Annotated[
|
|
Path, typer.Argument(help="Parquet file/dir to seed showers from (real events)")
|
|
],
|
|
checkpoint: Annotated[
|
|
Path,
|
|
typer.Option("--checkpoint", "-c", help="Checkpoint .pt (best.pt/last.pt)"),
|
|
],
|
|
geometry: Annotated[
|
|
Path,
|
|
typer.Option(
|
|
"--geometry",
|
|
"-g",
|
|
help="Geometry oracle .pkl (dwarf build-geometry-oracle)",
|
|
),
|
|
],
|
|
energy_cutoff: Annotated[
|
|
float,
|
|
typer.Option(
|
|
"--energy-cutoff",
|
|
help="Stop a track when its energy drops below this [MeV]",
|
|
),
|
|
] = 0.1,
|
|
max_steps: Annotated[
|
|
int, typer.Option("--max-steps", help="Max steps per individual track")
|
|
] = 1000,
|
|
steps: Annotated[
|
|
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,
|
|
max_tracks_per_event: Annotated[
|
|
Optional[int],
|
|
typer.Option(
|
|
"--max-tracks-per-event",
|
|
help="Safety cap on tracks per shower (sub-cap secondaries deposit in place)",
|
|
),
|
|
] = None,
|
|
escape_threshold: Annotated[
|
|
Optional[float],
|
|
typer.Option(
|
|
"--escape-threshold",
|
|
help="Override the oracle's NN-distance escape threshold [mm]",
|
|
),
|
|
] = None,
|
|
n_events: Annotated[
|
|
Optional[int], typer.Option("--n-events", help="Cap number of seed events")
|
|
] = None,
|
|
device: Annotated[
|
|
Optional[str], typer.Option("--device", "-d", help="cpu | cuda | mps (auto)")
|
|
] = None,
|
|
out: Annotated[
|
|
Optional[Path], typer.Option("--out", "-o", help="Output steps parquet")
|
|
] = None,
|
|
seed: Annotated[
|
|
Optional[int],
|
|
typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
|
|
] = None,
|
|
) -> None:
|
|
"""Roll the surrogate forward into full showers (autoregressive)."""
|
|
if seed is not None:
|
|
torch.manual_seed(seed)
|
|
np.random.seed(seed)
|
|
|
|
_device = torch.device(device) if device else gconfig.auto_device()
|
|
typer.echo(f"device: {_device}")
|
|
|
|
ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
|
|
for key in ("model_config", "sec_decoder"):
|
|
if key not in ckpt:
|
|
typer.echo(
|
|
f"error: checkpoint has no {key} — retrain with the current code",
|
|
err=True,
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
if "sec_phys" not in ckpt.get("normalizer", {}):
|
|
typer.echo(
|
|
"error: checkpoint has no normalizer.sec_phys — retrain with the "
|
|
"current code",
|
|
err=True,
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
model_cfg = ckpt["model_config"]
|
|
conditioning = model_cfg.get("conditioning", "embedding")
|
|
pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
|
|
mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
|
|
cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
|
|
tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
|
|
sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"])
|
|
|
|
model, sec_decoder = build_models(model_cfg)
|
|
_load_model_weights(model, sec_decoder, ckpt, weights, checkpoint)
|
|
model.to(_device).eval()
|
|
sec_decoder.to(_device).eval()
|
|
typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})")
|
|
|
|
oracle = GeometryOracle.load(geometry)
|
|
typer.echo(
|
|
f"loaded geometry oracle: {geometry} "
|
|
f"(escape_threshold={oracle.escape_threshold:.3f})"
|
|
)
|
|
|
|
files = find_parquet_files(data)
|
|
seeds = _seed_from_data(files, n_events)
|
|
typer.echo(f"seeded {len(seeds['event_id']):,} shower(s)")
|
|
|
|
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,
|
|
seeds,
|
|
cond_norm,
|
|
tgt_norm,
|
|
sec_phys_norm,
|
|
pdg_map,
|
|
mat_map,
|
|
energy_cutoff=energy_cutoff,
|
|
max_steps=max_steps,
|
|
steps=steps,
|
|
batch_size=batch_size,
|
|
device=_device,
|
|
max_tracks_per_event=max_tracks_per_event,
|
|
escape_threshold=escape_threshold,
|
|
on_chunk=_write_chunk,
|
|
conditioning=conditioning,
|
|
)
|
|
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())
|
|
ref.update(
|
|
{
|
|
"kind": "rollout",
|
|
"geometry_oracle": str(geometry.resolve()),
|
|
"energy_cutoff": energy_cutoff,
|
|
"max_steps": max_steps,
|
|
"steps": steps,
|
|
"max_tracks_per_event": max_tracks_per_event,
|
|
"n_seed_events": int(len(seeds["event_id"])),
|
|
}
|
|
)
|
|
ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False))
|
|
|
|
typer.echo(f"wrote {summary['n_rows']:,} step rows → {out}")
|
|
typer.echo(f"terminations: {summary['termination_reason_counts']}")
|
|
typer.echo(f"reference: {ref_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app()
|