import argparse import subprocess import tomllib from pathlib import Path import numpy as np import torch from torch.utils.data import DataLoader from giant.data.loader import ( find_parquet_files, load_event_ids, iter_file_chunks, build_index_maps_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 from giant.train import train as run_training 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_config(path: str) -> 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)) def main() -> None: parser = argparse.ArgumentParser(description="Train GIANT surrogate model") parser.add_argument("--config", default=None, help="Path to TOML config file") parser.add_argument("--data", required=True, help="Path to parquet file or directory") parser.add_argument("--mode", choices=["flow", "ddpm"]) parser.add_argument("--epochs", type=int) parser.add_argument("--batch-size", type=int) parser.add_argument("--lr", type=float) parser.add_argument("--hidden-dim", type=int) parser.add_argument("--n-blocks", type=int) parser.add_argument("--emb-dim", type=int) parser.add_argument("--val-fraction", type=float) parser.add_argument("--shuffle-buffer", type=int, default=65536, help="Rows held in RAM for shuffling per worker (default: 65536)") parser.add_argument("--out", default=None, help="Checkpoint output directory (default: auto from hyperparams)") parser.add_argument("--device", default=None, help="cpu | cuda | mps (default: auto)") parser.add_argument("--num-workers", type=int) args = parser.parse_args() # 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 args.config: file_cfg = _load_config(args.config) for section in ("train", "model"): cfg[section].update(file_cfg.get(section, {})) cli_train = {k: v for k, v in { "mode": args.mode, "epochs": args.epochs, "batch_size": args.batch_size, "lr": args.lr, "val_fraction": args.val_fraction, "num_workers": args.num_workers, }.items() if v is not None} cli_model = {k: v for k, v in { "hidden_dim": args.hidden_dim, "n_blocks": args.n_blocks, "emb_dim": args.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(args.device) if args.device else _auto_device() out_dir = Path(args.out or ( 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']}" )) print(f"device: {device}") print(f"out_dir: {out_dir}") # --- Discover files --- files = find_parquet_files(args.data) print(f"found {len(files)} parquet file(s)") # --- Scan event IDs (single column, cheap) --- print("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"]) n_train_steps = (np.isin(all_event_ids, np.array(sorted(train_events)))).sum() print(f" {len(all_event_ids):,} steps | " f"{len(train_events)} train events (~{n_train_steps:,} steps) | " f"{len(val_events)} val events") # --- Scan PDG / material vocabularies (2 columns, cheap) --- print("building vocabulary maps …") pdg_map, mat_map = build_index_maps_from_files(files) print(f" {len(pdg_map)} PDG codes | {len(mat_map)} materials") # --- Streaming normalizer fit over training data --- print("fitting normalizer (streaming) …") events_arr = np.array(sorted(train_events)) cond_acc = _WelfordAccumulator(9) tgt_acc = _WelfordAccumulator(9) 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() # --- Streaming datasets --- 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=args.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"], ) print(f"model: {sum(p.numel() for p in model.parameters()):,} parameters") out_dir.mkdir(parents=True, exist_ok=True) _save_config(cfg, out_dir) 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()}, ) if __name__ == "__main__": main()