From 72bd65ff9f2b03af14f84e56953abaa32e5a3ea1 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 18 Jun 2026 10:36:55 +0200 Subject: [PATCH] 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 --- CLAUDE.md | 4 +-- giant/cli.py | 17 +++++++++++-- giant/data/loader.py | 1 + giant/data/transforms.py | 35 +++++++++++++++++++++++++- giant/model/network.py | 2 +- giant/sample.py | 6 ++--- giant/validate.py | 5 +++- scripts/train.py | 2 +- tests/test_flow.py | 6 ++--- tests/test_network.py | 6 ++--- tests/test_transforms.py | 54 ++++++++++++++++++++++++++++++++++++++++ 11 files changed, 121 insertions(+), 17 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 0302cf8..c4e6d2f 100644 --- a/CLAUDE.md +++ b/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. diff --git a/giant/cli.py b/giant/cli.py index 3a8d6c8..39c0508 100644 --- a/giant/cli.py +++ b/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: diff --git a/giant/data/loader.py b/giant/data/loader.py index 5572829..2c908e9 100644 --- a/giant/data/loader.py +++ b/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), } diff --git a/giant/data/transforms.py b/giant/data/transforms.py index a05c983..0c545f6 100644 --- a/giant/data/transforms.py +++ b/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"], diff --git a/giant/model/network.py b/giant/model/network.py index 31f5427..5ae91a2 100644 --- a/giant/model/network.py +++ b/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) diff --git a/giant/sample.py b/giant/sample.py index e5358a2..5ae4be5 100644 --- a/giant/sample.py +++ b/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) diff --git a/giant/validate.py b/giant/validate.py index aab2e0b..83630ca 100644 --- a/giant/validate.py +++ b/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 diff --git a/scripts/train.py b/scripts/train.py index 80ee63d..1c321da 100644 --- a/scripts/train.py +++ b/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) diff --git a/tests/test_flow.py b/tests/test_flow.py index 5959271..00901ef 100644 --- a/tests/test_flow.py +++ b/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) diff --git a/tests/test_network.py b/tests/test_network.py index a499913..ba5eee9 100644 --- a/tests/test_network.py +++ b/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) diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 0a9a196..51e0da5 100644 --- a/tests/test_transforms.py +++ b/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)