From 646a9d7a72d05769b55daabb07ceccf9fe871216 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 17 Jun 2026 11:10:27 +0200 Subject: [PATCH] Add giant predict command - iter_cond_chunks: column-projected row-group streaming; post-step variables are never read from disk during inference - build_cond_features: assembles conditioning arrays without any target or post-step fields - inv_local_frame_rotation: Rodrigues R^T (negative angle) to rotate predicted post_dir back from local frame to world frame - giant predict: loads checkpoint, streams input, runs flow matching sampler, inverse-normalises and inverse-rotates outputs, writes predictions incrementally as parquet via PyArrow ParquetWriter - train now saves model_config in checkpoint so predict can reconstruct the architecture without extra CLI flags Co-Authored-By: Claude Sonnet 4.6 --- giant/cli.py | 122 ++++++++++++++++++++++++++++++++++++++- giant/data/loader.py | 28 +++++++++ giant/data/transforms.py | 47 +++++++++++++++ giant/train.py | 3 + 4 files changed, 199 insertions(+), 1 deletion(-) diff --git a/giant/cli.py b/giant/cli.py index cdf7258..65c6cac 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -10,15 +10,27 @@ 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, _WelfordAccumulator +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) @@ -191,6 +203,11 @@ def train( 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, @@ -199,4 +216,107 @@ def train( 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: _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_energy": chunk["pre_energy"], + "pre_dir_x": chunk["pre_dir"][:, 0], + "pre_dir_y": chunk["pre_dir"][:, 1], + "pre_dir_z": chunk["pre_dir"][:, 2], + "material_id": chunk["material"], + "layer_id": chunk["layer_id"], + "n_secondaries": chunk["n_sec"], + "step_length": step_length, + "delta_e": delta_e, + "edep": edep, + "post_dir_x": post_dir_world[:, 0], + "post_dir_y": post_dir_world[:, 1], + "post_dir_z": 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}") diff --git a/giant/data/loader.py b/giant/data/loader.py index a2bdd08..d9aab78 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -49,6 +49,34 @@ def iter_file_chunks(path: str | Path) -> Iterator[dict[str, np.ndarray]]: yield _df_to_dict(pf.read_row_group(i).to_pandas()) +_COND_COLS = [ + "event_id", "pdg", + "pre_x", "pre_y", "pre_z", "pre_energy", + "pre_dir_x", "pre_dir_y", "pre_dir_z", + "material_id", "layer_id", "n_secondaries", +] + + +def _cond_df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]: + return { + "event_id": df["event_id"].to_numpy(), + "pdg": df["pdg"].to_numpy(dtype=np.int32), + "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32), + "pre_energy": df["pre_energy"].to_numpy(dtype=np.float32), + "pre_dir": df[["pre_dir_x", "pre_dir_y", "pre_dir_z"]].to_numpy(dtype=np.float32), + "material": df["material_id"].to_numpy(dtype=np.int32), + "layer_id": df["layer_id"].to_numpy(dtype=np.int32), + "n_sec": df["n_secondaries"].to_numpy(dtype=np.int32), + } + + +def iter_cond_chunks(path: str | Path) -> Iterator[dict[str, np.ndarray]]: + """Yield conditioning-only row-groups (no post-step columns read from disk).""" + pf = pq.ParquetFile(path) + for i in range(pf.num_row_groups): + yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas()) + + def build_index_maps( data: dict[str, np.ndarray], ) -> tuple[dict[int, int], dict[int, int]]: diff --git a/giant/data/transforms.py b/giant/data/transforms.py index 0f01fcf..bd74970 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -98,6 +98,53 @@ class _WelfordAccumulator: return norm +def inv_local_frame_rotation(pre_dir: np.ndarray, post_dir_local: np.ndarray) -> np.ndarray: + """Inverse of local_frame_rotation: rotate from local frame back to world frame. + + Applies R^T (same axis, negative angle) to post_dir_local. + """ + z = np.array([[0.0, 0.0, 1.0]], dtype=np.float32) + + cos_t = np.clip((pre_dir * z).sum(axis=1, keepdims=True), -1.0, 1.0) + sin_t = np.sqrt(np.maximum(0.0, 1.0 - cos_t ** 2)) + + axis = np.cross(pre_dir, z) + axis_norm = np.linalg.norm(axis, axis=1, keepdims=True) + safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm) + axis = np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm) + + kxv = np.cross(axis, post_dir_local) + kdv = (axis * post_dir_local).sum(axis=1, keepdims=True) + + # Negative angle: sin_t → -sin_t + return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(np.float32) + + +def build_cond_features( + data: dict[str, np.ndarray], + pdg_map: dict[int, int], + mat_map: dict[int, int], + cond_normalizer: "Normalizer | None" = None, +) -> tuple[np.ndarray, np.ndarray]: + """Build conditioning arrays only — no target, no post-step variables.""" + cond_cont = np.column_stack([ + data["pre_pos"], + log_transform(data["pre_energy"]), + data["pre_dir"], + data["layer_id"].astype(np.float32), + data["n_sec"].astype(np.float32), + ]).astype(np.float32) + + pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64) + mat_idx = np.array([mat_map[int(m)] for m in data["material"]], dtype=np.int64) + cond_cat = np.column_stack([pdg_idx, mat_idx]) + + if cond_normalizer is not None: + cond_cont = cond_normalizer.transform(cond_cont) + + return cond_cont, cond_cat + + def build_features( data: dict[str, np.ndarray], pdg_map: dict[int, int], diff --git a/giant/train.py b/giant/train.py index e23b3a8..e803b86 100644 --- a/giant/train.py +++ b/giant/train.py @@ -19,6 +19,7 @@ def train( normalizer_dict: dict | None = None, pdg_map: dict | None = None, mat_map: dict | None = None, + model_config: dict | None = None, ) -> None: out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) @@ -81,6 +82,8 @@ def train( ckpt["pdg_map"] = pdg_map if mat_map is not None: ckpt["mat_map"] = mat_map + if model_config is not None: + ckpt["model_config"] = model_config torch.save(ckpt, out_dir / "best.pt") torch.save({"model": model.state_dict()}, out_dir / "last.pt")