lars
c3b7b2744c
The dataset yielded one row at a time, forcing DataLoader's default collate to Python-loop over every row to assemble each batch. That loop scales with batch size and was pinning a CPU core at 100% while the GPU sat idle. Now the dataset yields whole batches via vectorized numpy slicing, used with DataLoader(batch_size=None). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
330 lines
12 KiB
Python
330 lines
12 KiB
Python
import subprocess
|
|
import tomllib
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
import torch
|
|
import typer
|
|
from torch.utils.data import DataLoader
|
|
from typing_extensions import Annotated
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
|
|
from giant.data.loader import (
|
|
find_parquet_files,
|
|
load_event_ids,
|
|
iter_file_chunks,
|
|
iter_cond_chunks,
|
|
build_index_maps_from_files,
|
|
)
|
|
from giant.data.transforms import (
|
|
build_features,
|
|
build_cond_features,
|
|
inv_local_frame_rotation,
|
|
inv_log_transform,
|
|
_WelfordAccumulator,
|
|
Normalizer,
|
|
)
|
|
from giant.data.dataset import make_event_split, StreamingStepsDataset
|
|
from giant.model.network import DenoisingMLP
|
|
from giant.sample import sample_flow
|
|
from giant.train import train as run_training
|
|
|
|
app = typer.Typer(no_args_is_help=True)
|
|
|
|
|
|
@app.callback()
|
|
def _main() -> None:
|
|
"""GIANT — Geant4 step-function surrogate."""
|
|
|
|
|
|
class Mode(str, Enum):
|
|
flow = "flow"
|
|
ddpm = "ddpm"
|
|
|
|
|
|
def _auto_device() -> torch.device:
|
|
if torch.cuda.is_available():
|
|
return torch.device("cuda")
|
|
if torch.backends.mps.is_available():
|
|
return torch.device("mps")
|
|
return torch.device("cpu")
|
|
|
|
|
|
def _git_hash() -> str:
|
|
try:
|
|
return subprocess.check_output(
|
|
["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL
|
|
).decode().strip()
|
|
except Exception:
|
|
return "unknown"
|
|
|
|
|
|
def _load_toml(path: Path) -> dict:
|
|
with open(path, "rb") as f:
|
|
return tomllib.load(f)
|
|
|
|
|
|
def _save_config(cfg: dict, out_dir: Path) -> None:
|
|
lines = [f"# git: {_git_hash()}", ""]
|
|
for section, values in cfg.items():
|
|
lines.append(f"[{section}]")
|
|
for k, v in values.items():
|
|
lines.append(f"{k:<12} = {repr(v) if isinstance(v, str) else v}")
|
|
lines.append("")
|
|
(out_dir / "config.toml").write_text("\n".join(lines))
|
|
|
|
|
|
@app.command()
|
|
def train(
|
|
data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")],
|
|
config: Annotated[Optional[Path], typer.Option(help="TOML config file (overridden by explicit flags)")] = None,
|
|
mode: Annotated[Optional[Mode], typer.Option(help="Generative model: flow matching or DDPM")] = None,
|
|
epochs: Annotated[Optional[int], typer.Option()] = None,
|
|
batch_size: Annotated[Optional[int], typer.Option()] = None,
|
|
lr: Annotated[Optional[float], typer.Option()] = None,
|
|
hidden_dim: Annotated[Optional[int], typer.Option()] = None,
|
|
n_blocks: Annotated[Optional[int], typer.Option()] = None,
|
|
emb_dim: Annotated[Optional[int], typer.Option()] = None,
|
|
val_fraction: Annotated[Optional[float], typer.Option()] = None,
|
|
shuffle_buffer: Annotated[int, typer.Option(help="Rows held in RAM per worker for shuffling")] = 65536,
|
|
out: Annotated[Optional[Path], typer.Option(help="Checkpoint dir (default: auto from hyperparams)")] = None,
|
|
device: Annotated[Optional[str], typer.Option(help="cpu | cuda | mps (default: auto)")] = None,
|
|
num_workers: Annotated[Optional[int], typer.Option()] = None,
|
|
) -> None:
|
|
"""Train the GIANT surrogate model."""
|
|
# Defaults → config file → explicit CLI flags.
|
|
cfg: dict = {
|
|
"train": {
|
|
"mode": "flow", "epochs": 100, "batch_size": 4096, "lr": 3e-4,
|
|
"val_fraction": 0.1, "num_workers": 4,
|
|
},
|
|
"model": {
|
|
"hidden_dim": 256, "n_blocks": 6, "emb_dim": 16,
|
|
},
|
|
}
|
|
|
|
if config is not None:
|
|
file_cfg = _load_toml(config)
|
|
for section in ("train", "model"):
|
|
cfg[section].update(file_cfg.get(section, {}))
|
|
|
|
cli_train = {k: v for k, v in {
|
|
"mode": mode.value if mode is not None else None,
|
|
"epochs": epochs, "batch_size": batch_size, "lr": lr,
|
|
"val_fraction": val_fraction, "num_workers": num_workers,
|
|
}.items() if v is not None}
|
|
cli_model = {k: v for k, v in {
|
|
"hidden_dim": hidden_dim, "n_blocks": n_blocks, "emb_dim": emb_dim,
|
|
}.items() if v is not None}
|
|
cfg["train"].update(cli_train)
|
|
cfg["model"].update(cli_model)
|
|
|
|
t, m = cfg["train"], cfg["model"]
|
|
|
|
_device = torch.device(device) if device else _auto_device()
|
|
out_dir = out or Path(
|
|
f"checkpoints/{t['mode']}"
|
|
f"_h{m['hidden_dim']}"
|
|
f"_b{m['n_blocks']}"
|
|
f"_e{m['emb_dim']}"
|
|
f"_lr{t['lr']}"
|
|
f"_bs{t['batch_size']}"
|
|
)
|
|
|
|
typer.echo(f"device: {_device}")
|
|
typer.echo(f"out_dir: {out_dir}")
|
|
|
|
files = find_parquet_files(data)
|
|
typer.echo(f"found {len(files)} parquet file(s)")
|
|
|
|
typer.echo("scanning event IDs …")
|
|
all_event_ids = np.concatenate([load_event_ids(f) for f in files])
|
|
train_events, val_events = make_event_split(all_event_ids, val_fraction=t["val_fraction"])
|
|
events_arr = np.array(sorted(train_events))
|
|
n_train_steps = int(np.isin(all_event_ids, events_arr).sum())
|
|
typer.echo(
|
|
f" {len(all_event_ids):,} steps | "
|
|
f"{len(train_events)} train events (~{n_train_steps:,} steps) | "
|
|
f"{len(val_events)} val events"
|
|
)
|
|
|
|
typer.echo("building vocabulary maps …")
|
|
pdg_map, mat_map = build_index_maps_from_files(files)
|
|
typer.echo(f" {len(pdg_map)} PDG codes | {len(mat_map)} materials")
|
|
|
|
typer.echo("fitting normalizer (streaming) …")
|
|
cond_acc = _WelfordAccumulator(9)
|
|
tgt_acc = _WelfordAccumulator(6)
|
|
for path in files:
|
|
for chunk in iter_file_chunks(path):
|
|
mask = np.isin(chunk["event_id"], events_arr)
|
|
if not mask.any():
|
|
continue
|
|
chunk_tr = {k: v[mask] for k, v in chunk.items()}
|
|
cond_cont, _, target, _, _ = build_features(chunk_tr, pdg_map, mat_map)
|
|
cond_acc.update(cond_cont)
|
|
tgt_acc.update(target)
|
|
cond_norm = cond_acc.to_normalizer()
|
|
tgt_norm = tgt_acc.to_normalizer()
|
|
|
|
train_ds = StreamingStepsDataset(
|
|
files=files, split_events=train_events,
|
|
pdg_map=pdg_map, mat_map=mat_map,
|
|
cond_normalizer=cond_norm, target_normalizer=tgt_norm,
|
|
batch_size=t["batch_size"],
|
|
shuffle_buffer=shuffle_buffer, shuffle=True,
|
|
)
|
|
val_ds = StreamingStepsDataset(
|
|
files=files, split_events=val_events,
|
|
pdg_map=pdg_map, mat_map=mat_map,
|
|
cond_normalizer=cond_norm, target_normalizer=tgt_norm,
|
|
batch_size=t["batch_size"],
|
|
shuffle=False,
|
|
)
|
|
|
|
# Dataset yields whole batches already, so batch_size=None tells DataLoader
|
|
# to pass them through instead of re-collating row-by-row in Python.
|
|
pin = _device.type == "cuda"
|
|
train_loader = DataLoader(
|
|
train_ds, batch_size=None,
|
|
num_workers=t["num_workers"], pin_memory=pin,
|
|
)
|
|
val_loader = DataLoader(
|
|
val_ds, batch_size=None,
|
|
num_workers=t["num_workers"], pin_memory=pin,
|
|
)
|
|
|
|
model = DenoisingMLP(
|
|
pdg_vocab=len(pdg_map), mat_vocab=len(mat_map),
|
|
hidden_dim=m["hidden_dim"], n_blocks=m["n_blocks"], emb_dim=m["emb_dim"],
|
|
)
|
|
typer.echo(f"model: {sum(p.numel() for p in model.parameters()):,} parameters")
|
|
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
_save_config(cfg, out_dir)
|
|
|
|
model_config = {
|
|
"pdg_vocab": len(pdg_map), "mat_vocab": len(mat_map),
|
|
"hidden_dim": m["hidden_dim"], "n_blocks": m["n_blocks"], "emb_dim": m["emb_dim"],
|
|
}
|
|
|
|
run_training(
|
|
model=model,
|
|
train_loader=train_loader, val_loader=val_loader,
|
|
mode=t["mode"], epochs=t["epochs"], lr=t["lr"],
|
|
device=_device, out_dir=out_dir,
|
|
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()},
|
|
model_config=model_config,
|
|
)
|
|
|
|
|
|
@app.command()
|
|
def predict(
|
|
data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")],
|
|
checkpoint: Annotated[Path, typer.Option(help="Path to checkpoint .pt file (best.pt or last.pt)")],
|
|
out: Annotated[Optional[Path], typer.Option(help="Output parquet path (default: <data>_predicted.parquet)")] = None,
|
|
batch_size: Annotated[int, typer.Option(help="Inference batch size")] = 4096,
|
|
steps: Annotated[int, typer.Option(help="Flow matching ODE steps")] = 10,
|
|
device: Annotated[Optional[str], typer.Option(help="cpu | cuda | mps (default: auto)")] = None,
|
|
) -> None:
|
|
"""Run trained model on a parquet file and save predictions."""
|
|
_device = torch.device(device) if device else _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)
|
|
|
|
model_cfg = ckpt["model_config"]
|
|
pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
|
|
mat_map = {int(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"])
|
|
|
|
model = DenoisingMLP(**model_cfg)
|
|
model.load_state_dict(ckpt["model"])
|
|
model.to(_device).eval()
|
|
typer.echo(f"loaded checkpoint: {checkpoint}")
|
|
|
|
# --- Output path ---
|
|
if out is None:
|
|
stem = data.stem if data.is_file() else data.name
|
|
out = data.parent / f"{stem}_predicted.parquet"
|
|
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
|
|
|
|
for path in files:
|
|
for chunk in iter_cond_chunks(path):
|
|
N = len(chunk["event_id"])
|
|
cond_cont, cond_cat = build_cond_features(chunk, pdg_map, mat_map, cond_norm)
|
|
|
|
# Inference in batch_size slices
|
|
pred_parts = []
|
|
for start in range(0, N, batch_size):
|
|
end = min(start + batch_size, N)
|
|
cc = torch.from_numpy(cond_cont[start:end]).float().to(_device)
|
|
ck = torch.from_numpy(cond_cat[start:end]).long().to(_device)
|
|
pred_parts.append(sample_flow(model, cc, ck, steps=steps).cpu().numpy())
|
|
pred = np.concatenate(pred_parts, axis=0) # (N, 6) normalised
|
|
|
|
# Inverse-normalise → local frame, log-scaled scalars
|
|
raw = tgt_norm.inverse_transform(pred)
|
|
|
|
step_length = inv_log_transform(raw[:, 0])
|
|
delta_e = inv_log_transform(raw[:, 1])
|
|
edep = inv_log_transform(raw[:, 2])
|
|
|
|
# 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(chunk["pre_dir"], post_dir_local)
|
|
|
|
table = pa.table({
|
|
"event_id": chunk["event_id"],
|
|
"pdg": chunk["pdg"],
|
|
"pre_x": chunk["pre_pos"][:, 0],
|
|
"pre_y": chunk["pre_pos"][:, 1],
|
|
"pre_z": chunk["pre_pos"][:, 2],
|
|
"pre_E": chunk["pre_E"],
|
|
"pre_dx": chunk["pre_dir"][:, 0],
|
|
"pre_dy": chunk["pre_dir"][:, 1],
|
|
"pre_dz": chunk["pre_dir"][:, 2],
|
|
"material": chunk["material"],
|
|
"layer_id": chunk["layer_id"],
|
|
"n_sec": chunk["n_sec"],
|
|
"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],
|
|
})
|
|
|
|
if writer is None:
|
|
writer = pq.ParquetWriter(out, table.schema)
|
|
writer.write_table(table)
|
|
total += N
|
|
|
|
if writer is not None:
|
|
writer.close()
|
|
|
|
typer.echo(f"wrote {total:,} rows → {out}")
|
|
|
|
if __name__ == "__main__":
|
|
app()
|