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 <noreply@anthropic.com>

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

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

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

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