Add --coord local mode to predict for raw-space prediction debugging
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 <noreply@anthropic.com>
This commit is contained in:
+92
-44
@@ -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: <data>_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: <data>_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)
|
||||
|
||||
Reference in New Issue
Block a user