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>

CLAUDE.md

@@ -17,7 +17,7 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
 
 **Data pipeline** (`giant/data/`): parquet files from miniCaloSim are loaded into numpy arrays (`loader.py`), then log-transformed and rotated into a local coordinate frame where `pre_dir = ẑ` (`transforms.py`), before being wrapped in a PyTorch `Dataset` (`dataset.py`). Train/val split is by `event_id` to avoid leaking correlated steps from the same shower.
 
-**Output space (6D):** `step_length` (log), `ΔE` (log), `edep` (log), and `post_dir` as a unit vector in the local frame.
+**Output space (9D):** `step_length` (log), `ΔE` (log), `edep` (log), `post_dir` (post-scattering momentum direction, unit vector in the local frame), and `travel_dir` (direction of `post_pos - pre_pos`, unit vector in the local frame). `post_pos` itself is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, since `step_length` already encodes that displacement's magnitude and duplicating it would let the two become inconsistent.
 
 **Conditioning vector:** PDG code (embedding), pre-step position, log(pre-energy), pre-step direction, material (embedding), layer ID, number of secondaries (Phase 1 only — see Roadmap below).
 
@@ -29,6 +29,6 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
 
 ## Roadmap
 
-Phase 1 (current): number of secondaries is a conditioning input — model predicts only 6D post-step kinematics.
+Phase 1 (current): number of secondaries is a conditioning input — model predicts only 9D post-step kinematics (including derived post_pos).
 
 Phase 2 (target): model must jointly predict the number of secondaries and all their properties (energy, direction, species), requiring an extended output space and likely a set-based or autoregressive generation scheme for the variable-length secondary list.

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:

giant/data/loader.py

@@ -30,6 +30,7 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
         "delta_e": (df["pre_E"] - df["post_E"]).to_numpy(dtype=np.float32),
         "edep": df["edep"].to_numpy(dtype=np.float32),
         "post_dir": df[["post_dx", "post_dy", "post_dz"]].to_numpy(dtype=np.float32),
+        "post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32),
     }
 
 

giant/data/transforms.py

@@ -98,6 +98,35 @@ class _WelfordAccumulator:
         return norm
 
 
+def travel_direction(pre_pos: np.ndarray, post_pos: np.ndarray) -> np.ndarray:
+    """World-frame unit vector pointing from pre_pos to post_pos.
+
+    Kept independent of `step_length`: that scalar already encodes the
+    magnitude of this displacement, so this function only ever returns
+    direction (norm-guarded the same way as `local_frame_rotation`'s axis).
+    """
+    disp = post_pos - pre_pos
+    norm = np.linalg.norm(disp, axis=1, keepdims=True)
+    safe_norm = np.where(norm < 1e-7, 1.0, norm)
+    return np.where(norm < 1e-7, np.array([[0.0, 0.0, 1.0]]), disp / safe_norm).astype(np.float32)
+
+
+def reconstruct_post_pos(
+    pre_pos: np.ndarray,
+    pre_dir: np.ndarray,
+    step_length: np.ndarray,
+    travel_dir_local: np.ndarray,
+) -> np.ndarray:
+    """Inverse of the travel_direction/local_frame_rotation encoding.
+
+    Single source of truth for combining the magnitude (`step_length`) and
+    direction (`travel_dir_local`) back into a world-frame post_pos, so
+    `step_length` and post_pos stay consistent by construction.
+    """
+    travel_dir_world = inv_local_frame_rotation(pre_dir, travel_dir_local)
+    return (pre_pos + step_length.reshape(-1, 1) * travel_dir_world).astype(np.float32)
+
+
 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.
 
@@ -158,13 +187,17 @@ def build_features(
     When fit=True, new Normalizers are fitted on the supplied arrays.
     """
     post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])
+    travel_dir_local = local_frame_rotation(
+        data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"])
+    )
 
     target = np.column_stack([
         log_transform(data["step_length"]),
         log_transform(data["delta_e"]),
         log_transform(data["edep"]),
         post_dir_local,
-    ]).astype(np.float32)  # (N, 6)
+        travel_dir_local,
+    ]).astype(np.float32)  # (N, 9)
 
     cond_cont = np.column_stack([
         data["pre_pos"],

giant/model/network.py

@@ -73,7 +73,7 @@ class DenoisingMLP(nn.Module):
         emb_dim: int = 16,
         time_dim: int = 64,
         cond_out_dim: int = 128,
-        x_dim: int = 6,
+        x_dim: int = 9,
     ) -> None:
         super().__init__()
         self.time_emb = SinusoidalEmbedding(time_dim)

giant/sample.py

@@ -12,7 +12,7 @@ def sample_flow(
     model.eval()
     B = cond_cont.size(0)
     device = cond_cont.device
-    x = torch.randn(B, 6, device=device)
+    x = torch.randn(B, 9, device=device)
     dt = 1.0 / steps
     for i in range(steps):
         t = torch.full((B,), i * dt, device=device)
@@ -32,7 +32,7 @@ def sample_ddpm(
     model.eval()
     B = cond_cont.size(0)
     device = cond_cont.device
-    x = torch.randn(B, 6, device=device)
+    x = torch.randn(B, 9, device=device)
     T = schedule.T
     for i in reversed(range(T)):
         t_norm = torch.full((B,), i / T, device=device)
@@ -63,7 +63,7 @@ def sample_ddim(
     device = cond_cont.device
     T = schedule.T
     timesteps = torch.linspace(T - 1, 0, steps, dtype=torch.long, device=device)
-    x = torch.randn(B, 6, device=device)
+    x = torch.randn(B, 9, device=device)
     for step_idx, ts in enumerate(timesteps):
         t_idx = int(ts.item())
         t_norm = torch.full((B,), t_idx / T, device=device)

giant/validate.py

@@ -11,6 +11,9 @@ _TARGET_NAMES = [
     "post_dx",
     "post_dy",
     "post_dz",
+    "travel_dx",
+    "travel_dy",
+    "travel_dz",
 ]
 
 
@@ -24,7 +27,7 @@ def validate_marginals(
 ) -> dict[str, np.ndarray]:
     """Compare per-dimension marginals of generated vs. real steps.
 
-    Returns {"real": (N,6), "generated": (N,6)} in normalised space.
+    Returns {"real": (N,9), "generated": (N,9)} in normalised space.
     """
     if device is None:
         device = next(model.parameters()).device

scripts/train.py

@@ -133,7 +133,7 @@ def main() -> None:
     print("fitting normalizer (streaming) …")
     events_arr = np.array(sorted(train_events))
     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)

tests/test_flow.py

@@ -10,7 +10,7 @@ def _small_model():
 
 
 def _batch(B=8):
-    x1 = torch.randn(B, 6)
+    x1 = torch.randn(B, 9)
     cond_cont = torch.randn(B, 9)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     return x1, cond_cont, cond_cat
@@ -40,7 +40,7 @@ def test_sample_flow_shape():
     cond_cont = torch.randn(B, 9)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     out = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
-    assert out.shape == (B, 6)
+    assert out.shape == (B, 9)
 
 
 def test_ddpm_loss_nonneg():
@@ -56,4 +56,4 @@ def test_sample_ddim_shape():
     cond_cont = torch.randn(B, 9)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     out = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
-    assert out.shape == (B, 6)
+    assert out.shape == (B, 9)

tests/test_network.py

@@ -18,7 +18,7 @@ def test_sinusoidal_embedding_batch_1():
 def test_denoising_mlp_output_shape():
     B = 8
     model = DenoisingMLP(pdg_vocab=5, mat_vocab=3)
-    x_t = torch.randn(B, 6)
+    x_t = torch.randn(B, 9)
     t = torch.rand(B)
     cond_cont = torch.randn(B, 9)
     cond_cat = torch.stack([
@@ -26,13 +26,13 @@ def test_denoising_mlp_output_shape():
         torch.randint(0, 3, (B,)),
     ], dim=1)
     out = model(x_t, t, cond_cont, cond_cat)
-    assert out.shape == (B, 6)
+    assert out.shape == (B, 9)
 
 
 def test_denoising_mlp_gradients_flow():
     B = 4
     model = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
-    x_t = torch.randn(B, 6)
+    x_t = torch.randn(B, 9)
     t = torch.rand(B)
     cond_cont = torch.randn(B, 9)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)

tests/test_transforms.py

@@ -5,6 +5,8 @@ from giant.data.transforms import (
     local_frame_rotation,
     log_transform,
     Normalizer,
+    reconstruct_post_pos,
+    travel_direction,
 )
 
 
@@ -50,6 +52,58 @@ def test_local_frame_rotation_preserves_norm():
     np.testing.assert_allclose(np.linalg.norm(result, axis=1), 1.0, atol=1e-5)
 
 
+def test_travel_direction_is_unit_norm():
+    rng = np.random.default_rng(5)
+    N = 50
+    pre_pos = rng.standard_normal((N, 3)).astype(np.float32)
+    post_pos = pre_pos + rng.standard_normal((N, 3)).astype(np.float32)
+    result = travel_direction(pre_pos, post_pos)
+    np.testing.assert_allclose(np.linalg.norm(result, axis=1), 1.0, atol=1e-5)
+
+
+def test_travel_direction_matches_normalized_displacement():
+    rng = np.random.default_rng(6)
+    N = 50
+    pre_pos = rng.standard_normal((N, 3)).astype(np.float32)
+    disp = rng.standard_normal((N, 3)).astype(np.float32)
+    post_pos = pre_pos + disp
+    expected = disp / np.linalg.norm(disp, axis=1, keepdims=True)
+    np.testing.assert_allclose(travel_direction(pre_pos, post_pos), expected, atol=1e-5)
+
+
+def test_reconstruct_post_pos_straight_line():
+    """When post_pos = pre_pos + L * pre_dir, travel_dir equals pre_dir, so its
+    local-frame encoding is ẑ — reconstruction must recover post_pos exactly."""
+    rng = np.random.default_rng(7)
+    N = 20
+    pre_pos = rng.standard_normal((N, 3)).astype(np.float32)
+    pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
+    pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
+    step_length = rng.uniform(0.1, 5.0, size=N).astype(np.float32)
+    post_pos = pre_pos + step_length[:, None] * pre_dir
+
+    travel_dir_local = local_frame_rotation(pre_dir, travel_direction(pre_pos, post_pos))
+    np.testing.assert_allclose(travel_dir_local, np.tile([0, 0, 1], (N, 1)), atol=1e-4)
+
+    reconstructed = reconstruct_post_pos(pre_pos, pre_dir, step_length, travel_dir_local)
+    np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)
+
+
+def test_reconstruct_post_pos_general_roundtrip():
+    """Full encode (build_features-style) -> decode (cli.py predict-style) path."""
+    rng = np.random.default_rng(8)
+    N = 100
+    pre_pos = rng.standard_normal((N, 3)).astype(np.float32)
+    pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
+    pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
+    post_pos = pre_pos + rng.standard_normal((N, 3)).astype(np.float32)
+    step_length = np.linalg.norm(post_pos - pre_pos, axis=1).astype(np.float32)
+
+    travel_dir_local = local_frame_rotation(pre_dir, travel_direction(pre_pos, post_pos))
+    reconstructed = reconstruct_post_pos(pre_pos, pre_dir, step_length, travel_dir_local)
+    np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)
+
+
 def test_normalizer_roundtrip():
     rng = np.random.default_rng(3)
     X = rng.standard_normal((200, 9)).astype(np.float32)