From f1a82b5853439e4851c32b18e83123a4049c2cb2 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 18 Jun 2026 10:55:58 +0200 Subject: [PATCH] Add --coord local mode to predict for raw-space prediction debugging MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Outputs the model's 9D prediction (denormalised only — still local frame, log-scaled scalars) alongside the matching ground-truth target for the same input rows, so they're directly comparable in the space the loss is actually computed in. Also fixes mat_map keys being cast with int() instead of str() when loading a checkpoint in predict. Co-Authored-By: Claude Sonnet 4.6 --- giant/cli.py | 136 ++++++++++++++++++++++++++++++++++----------------- 1 file changed, 92 insertions(+), 44 deletions(-) diff --git a/giant/cli.py b/giant/cli.py index 39c0508..9523e8f 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -36,6 +36,18 @@ from giant.train import train as run_training app = typer.Typer(no_args_is_help=True) +_LOCAL_TARGET_NAMES = [ + "log_step_length", + "log_delta_e", + "log_edep", + "post_dx", + "post_dy", + "post_dz", + "travel_dx", + "travel_dy", + "travel_dz", +] + @app.callback() def _main() -> None: @@ -47,6 +59,11 @@ class Mode(str, Enum): ddpm = "ddpm" +class Coord(str, Enum): + global_ = "global" + local = "local" + + def _auto_device() -> torch.device: if torch.cuda.is_available(): return torch.device("cuda") @@ -229,7 +246,13 @@ def train( 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, + coord: Annotated[Coord, typer.Option( + 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(help="Output parquet path (default: _predicted[_local].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, @@ -246,7 +269,7 @@ def predict( 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()} + 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"]) @@ -258,7 +281,8 @@ def predict( # --- Output path --- if out is None: stem = data.stem if data.is_file() else data.name - out = data.parent / f"{stem}_predicted.parquet" + suffix = "_predicted_local.parquet" if coord == Coord.local else "_predicted.parquet" + out = data.parent / f"{stem}{suffix}" typer.echo(f"output: {out}") # --- Stream input, generate predictions, write output --- @@ -267,11 +291,17 @@ def predict( writer: pq.ParquetWriter | None = None total = 0 + chunk_iter = iter_file_chunks if coord == Coord.local else iter_cond_chunks for path in files: - for chunk in iter_cond_chunks(path): + for chunk in chunk_iter(path): N = len(chunk["event_id"]) - cond_cont, cond_cat = build_cond_features(chunk, pdg_map, mat_map, cond_norm) + + if coord == Coord.local: + cond_cont, cond_cat, target_raw, _, _ = build_features(chunk, pdg_map, mat_map) + cond_cont = cond_norm.transform(cond_cont) + else: + cond_cont, cond_cat = build_cond_features(chunk, pdg_map, mat_map, cond_norm) # Inference in batch_size slices pred_parts = [] @@ -285,48 +315,66 @@ def predict( # 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]) + if coord == Coord.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"], + **{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]) + 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) + # 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) - # 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( - chunk["pre_pos"], chunk["pre_dir"], step_length, travel_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( + chunk["pre_pos"], chunk["pre_dir"], step_length, travel_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], - "post_x": post_pos_world[:, 0], - "post_y": post_pos_world[:, 1], - "post_z": post_pos_world[:, 2], - }) + 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], + "post_x": post_pos_world[:, 0], + "post_y": post_pos_world[:, 1], + "post_z": post_pos_world[:, 2], + }) if writer is None: writer = pq.ParquetWriter(out, table.schema)