Add post_pos as a model target via travel_dir decomposition

step_length already encodes |post_pos - pre_pos| by definition, so a raw
post_pos target would duplicate that magnitude and could drift inconsistent
with step_length during sampling. Instead add travel_dir, a unit vector
(local frame) giving only the direction of pre_pos->post_pos; post_pos is
reconstructed at inference as pre_pos + step_length * travel_dir, keeping
the two self-consistent. Target grows from 6D to 9D.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

giant/cli.py

@@ -25,6 +25,7 @@ from giant.data.transforms import (
     build_cond_features,
     inv_local_frame_rotation,
     inv_log_transform,
+    reconstruct_post_pos,
     _WelfordAccumulator,
     Normalizer,
 )
@@ -158,7 +159,7 @@ def train(
 
     typer.echo("fitting normalizer (streaming) …")
     cond_acc = _WelfordAccumulator(9)
-    tgt_acc = _WelfordAccumulator(6)
+    tgt_acc = _WelfordAccumulator(9)
     for path in files:
         for chunk in iter_file_chunks(path):
             mask = np.isin(chunk["event_id"], events_arr)
@@ -279,7 +280,7 @@ def predict(
                 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
+            pred = np.concatenate(pred_parts, axis=0)  # (N, 9) normalised
 
             # Inverse-normalise → local frame, log-scaled scalars
             raw = tgt_norm.inverse_transform(pred)
@@ -294,6 +295,15 @@ def predict(
             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
+            )
+
             table = pa.table({
                 "event_id":     chunk["event_id"],
                 "pdg":          chunk["pdg"],
@@ -313,6 +323,9 @@ def predict(
                 "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: