From e6e0eb22bfda8b8ace3e0806f4c438c19ff5d234 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 29 Jun 2026 11:34:31 +0200 Subject: [PATCH] Implement Phase 2: secondary particle prediction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two-stage factorisation: Stage 1 predicts 9D primary kinematics + n_sec classification head (COND_DIM reduced to 8, dropping n_sec/e_sec inputs); Stage 2 (SecondaryDecoder) generates K_MAX=15 secondary slots via masked flow matching over (stick_logit, local_dir, type_emb) conditioned on Stage 1 output. Joint training with combined loss L_s1 + λ_nsec*L_nsec + λ_s2*L_s2. Co-Authored-By: Claude Sonnet 4.6 --- docs/phase2_plan.md | 172 ++++++++++++++++++++++++++++ giant/cli.py | 3 +- giant/config.py | 2 + giant/constants.py | 28 +++-- giant/data/dataset.py | 112 +++++++++--------- giant/data/loader.py | 55 ++++++++- giant/data/transforms.py | 154 ++++++++++++++++++++++--- giant/model/network.py | 128 ++++++++++++++++++++- giant/model/schedule.py | 34 ++++++ giant/pipeline.py | 46 ++++++-- giant/sample.py | 85 ++++++++++++-- giant/train.py | 189 +++++++++++++++++++++++------- giant/validate.py | 11 +- scripts/steps_to_parquet.py | 69 ++++++++--- tests/test_dataset.py | 82 ++++--------- tests/test_flow.py | 10 +- tests/test_network.py | 6 +- tests/test_phase2.py | 222 ++++++++++++++++++++++++++++++++++++ 18 files changed, 1174 insertions(+), 234 deletions(-) create mode 100644 docs/phase2_plan.md create mode 100644 tests/test_phase2.py diff --git a/docs/phase2_plan.md b/docs/phase2_plan.md new file mode 100644 index 0000000..596edec --- /dev/null +++ b/docs/phase2_plan.md @@ -0,0 +1,172 @@ +# Phase 2: Secondary Particle Prediction + +## Context + +Phase 1 takes `n_sec` (secondary count) and `e_sec` (total secondary energy) as **conditioning inputs**. Phase 2 must instead **predict** them, making the surrogate self-contained for shower rollout. Per Jan's 2026-06-29 decision: hard discrete `n_sec` integer head; escalation to Gumbel-Softmax only if empirically needed. + +Two-stage factorization: +- **Stage 1**: existing 9D flow model (reduced conditioning: drop `n_sec` + `log(e_sec)`) + a new discrete `n_sec` classification head +- **Stage 2**: non-AR flow matching over `K_MAX` secondary slots simultaneously, each slot predicting `(stick_break_logit, dir_local_3D, type_emb)` — conditioned on pre-step state + Stage 1 output; padded slots masked from loss + +Training: joint, combined loss `L = L_flow_s1 + λ_nsec * L_nsec + λ_s2 * L_flow_s2`. + +--- + +## Prerequisite: Determine K_MAX + +Before implementing, run a quick analysis over existing parquet files to find `max(n_sec)` and the 99th percentile. Expected to be 5–20 for EM shower steps. Set `K_MAX` as a constant in `giant/constants.py` (suggest 15 as a starting point, revise from data). + +--- + +## New Branch + +```bash +git checkout -b phase2-secondary-prediction master +``` + +--- + +## Part A — Data Pipeline + +### A1. `scripts/steps_to_parquet.py` + +Extend `_add_secondary_energy` to also collect per-secondary attributes from the spawning tree join: +- For each `child_track_id`, look up the child's first step → get `pdg`, `pre_E`, `pre_dx/dy/dz` +- Emit list columns in the parquet: `sec_pdg_list`, `sec_E_list`, `sec_dx_list`, `sec_dy_list`, `sec_dz_list` +- Lists are sorted **descending by energy** at write time +- Truncate to `K_MAX` entries if needed (flag if any row truncated) + +Re-run ROOT→parquet conversion after this change. + +### A2. `giant/data/loader.py` + +In `_df_to_dict`: read the five new list columns. Pad each to length `K_MAX` with zeros (energy) / sentinel values (pdg → 0, dir → (0,0,1)). Return as fixed-shape arrays `(N, K_MAX)` / `(N, K_MAX, 3)`. + +Also return a boolean validity mask `sec_valid` of shape `(N, K_MAX)`: `True` for slots `i < n_sec`. + +### A3. `giant/data/transforms.py` + +Add `encode_secondaries(sec_pdg_list, sec_E_list, sec_dir_list, sec_valid, e_sec, pdg_emb_weight, pre_dir, K_MAX)`: +1. **Direction**: call existing `local_frame_rotation` per slot +2. **Energy (stick-breaking)**: + - Slot 0: `f_0 = E_0 / e_sec` → logit `log(f_0/(1-f_0))` (clamped) + - Slot i: `f_i = E_i / (e_sec - sum(E_0..E_{i-1}))` → logit + - Last valid slot: logit = large positive constant (takes all remaining budget) + - Padding slots (beyond `n_sec`): set logit = 0, masked out of loss anyway +3. **Type embedding**: index into `pdg_emb_weight` (the PDG embedding table weights) to get the target embedding vector for each secondary's `pdg`. Shape `(K_MAX, emb_dim)`. + +Returns `sec_targets: (K_MAX, 1 + 3 + emb_dim)` and `sec_valid: (K_MAX,)`. + +Inverse (`decode_secondaries`): sigmoid stick-breaking fractions → energies, inv local frame rotation → world dirs, nearest-neighbor lookup in PDG embedding table → pdg code. + +### A4. `giant/data/dataset.py` + +Update `build_features` and `StreamingStepsDataset.__iter__` to also yield `sec_targets` and `sec_valid` alongside the existing `(cond_cont, cond_cat, x1)` batch items. + +--- + +## Part B — Constants (`giant/constants.py`) + +- `COND_DIM`: 10 → **8** (remove `n_sec` and `log(e_sec)`) +- Add `K_MAX: int` (set after data analysis, e.g. 15) +- Add `SEC_SLOT_DIM: int` (= 4 + `emb_dim` = 20 for default emb_dim=16; 1 stick + 3 dir + 16 type) +- Add `SEC_DIM: int = K_MAX * SEC_SLOT_DIM` (flattened Stage 2 target dimension) +- Update `LOCAL_TARGET_NAMES` (Stage 1 only, still 9D) + +--- + +## Part C — Model (`giant/model/network.py`) + +### C1. `DenoisingMLP` — Stage 1 (minimal changes) + +- `ConditionEncoder.cont_dim` drops from 10 to 8 (COND_DIM change propagates automatically) +- Add `n_sec_head = nn.Sequential(Linear(cond_out_dim, hidden_dim//2), SiLU(), Linear(hidden_dim//2, K_MAX + 1))` applied to `c_emb` (the condition encoding, not the diffused latent) +- Add method `predict_n_sec(cond_cont, cond_cat) -> Tensor[B, K_MAX+1]` — no diffusion, just encode conditioning and run the head + +### C2. `SecondaryDecoder` — Stage 2 (new class) + +Architecture mirrors `DenoisingMLP` but: +- **Input**: `x_t` of shape `(B, SEC_DIM)` (flattened K_MAX secondary slots) +- **Conditioning**: pre-step state (8D cont + 2 cat → same ConditionEncoder as Stage 1) concatenated with Stage 1 output (9D normalized target, detached from Stage 1 loss for stability initially). Total cond dim to the ResBlocks: `time_dim + cond_s1_out_dim + 9` +- **Output**: vector field of shape `(B, SEC_DIM)` +- Uses same `ResBlock` / `SinusoidalEmbedding` / `ConditionEncoder` building blocks + +A `SecondaryConditionEncoder` wraps the base `ConditionEncoder` and concatenates the Stage 1 output: +```python +class SecondaryConditionEncoder(nn.Module): + # base: ConditionEncoder(pdg_vocab, mat_vocab, 8, emb_dim, cond_out_dim) + # stage1_proj: Linear(X_DIM, stage1_cond_dim) + # mlp: fuses both +``` + +--- + +## Part D — Loss / Training + +### `giant/model/schedule.py` + +Add `flow_matching_loss_masked(model, x1, cond_cont, cond_cat, mask)`: +- Same as `flow_matching_loss` but divides by `mask.sum()` instead of `B * SEC_DIM`, zeroing out padded slots before averaging. `mask` shape: `(B, K_MAX)`, broadcast over slot dims. + +### `giant/train.py` + +Batch now unpacks as `(cond_cont, cond_cat, x1_s1, n_sec_target, x1_s2, sec_mask)`. + +Combined loss per batch: +``` +L_s1 = flow_matching_loss(stage1_model, x1_s1, cond_cont, cond_cat) +L_nsec = cross_entropy(stage1_model.predict_n_sec(cond_cont, cond_cat), n_sec_target) +L_s2 = flow_matching_loss_masked(sec_decoder, x1_s2, cond_cont, cond_cat, stage1_detached, sec_mask) +L = L_s1 + lambda_nsec * L_nsec + lambda_s2 * L_s2 +``` + +Config adds `lambda_nsec` (suggest 0.1) and `lambda_s2` (suggest 1.0) under `[train]`. + +Both `stage1_model` and `sec_decoder` share a single `optimizer` (AdamW over all parameters). + +Checkpoint saves both `stage1_model.state_dict()` and `sec_decoder.state_dict()`, plus `K_MAX` and `SEC_SLOT_DIM` in `model_config`. + +### `giant/pipeline.py` + +- Compute `K_MAX` from data (max `n_sec` over training events) before constructing models +- Build both `DenoisingMLP` and `SecondaryDecoder`, pass both to `run_training` + +--- + +## Part E — Sampling (`giant/sample.py`) + +```python +def sample_stage1(model, cond_cont, cond_cat, steps=10): + # Euler ODE → primary sample (9D), + argmax n_sec head + ... + +def sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec, steps=10): + # Euler ODE on SEC_DIM → decode stick-breaking → energies + # inv_local_frame_rotation → world-frame dirs + # nearest-neighbor in pdg_emb_weight → pdg codes + # mask slots >= n_sec + ... +``` + +--- + +## Part F — Wiring + +- **`giant/validate.py`**: add secondary-specific marginals (n_sec distribution, species distribution, energy fraction per slot) +- **`giant/cli.py`**: `predict` command loads both checkpoints, calls both samplers, appends secondary columns to output parquet + +--- + +## Type embedding design note + +The type embedding target at training is `pdg_emb.weight[sec_pdg_idx]` (the Stage 1 PDG embedding table rows). Gradients flow into the embedding table from both the conditioning path (input PDG) and the secondary type loss — this is intentional; the shared embedding space is the bridge. At inference, snap: `argmin_k ||pred_emb - pdg_emb.weight[k]||`. + +--- + +## Verification + +1. `uv run pytest` — existing tests pass (Stage 1 shape/interface unchanged beyond COND_DIM) +2. Unit tests for `encode_secondaries` / `decode_secondaries` (round-trip: energies sum to `e_sec`, directions are unit vectors) +3. Unit test for `flow_matching_loss_masked`: verify padded slots contribute zero gradient +4. Short training run (1–2 epochs): confirm all three loss components decrease +5. Sampling smoke test: verify `sum(sec_E) ≈ e_sec` per sample, all directions unit-normed diff --git a/giant/cli.py b/giant/cli.py index 6636d40..6168933 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -409,7 +409,8 @@ def predict( cc = torch.from_numpy(cond_cont).float().to(_device) ck = torch.from_numpy(cond_cat).long().to(_device) - pred = sample_flow(model, cc, ck, steps=steps).cpu().numpy() # normalised + pred, _n_sec = sample_flow(model, cc, ck, steps=steps) + pred = pred.cpu().numpy() # normalised # Inverse-normalise → local frame, log-scaled scalars raw = tgt_norm.inverse_transform(pred) diff --git a/giant/config.py b/giant/config.py index 96d2c74..cbcf0e9 100644 --- a/giant/config.py +++ b/giant/config.py @@ -20,6 +20,8 @@ DEFAULT_CONFIG: dict = { "validate_every": 10, "validate_steps": 10, "warmup_epochs": 5, + "lambda_nsec": 0.1, + "lambda_s2": 1.0, }, "model": { "hidden_dim": 256, diff --git a/giant/constants.py b/giant/constants.py index 0bf8588..c42b8f2 100644 --- a/giant/constants.py +++ b/giant/constants.py @@ -1,16 +1,24 @@ X_DIM = 9 -# Conditioning continuous-feature width: pre_pos(3), log(pre_E)(1), pre_dir(3), -# layer_id(1), n_sec(1), log(e_sec)(1). One wider than X_DIM because e_sec -# (secondary energy) is a conditioning input in the energy-conservation PoC. -COND_DIM = 10 +# Conditioning continuous-feature width (Phase 2): pre_pos(3), log(pre_E)(1), +# pre_dir(3), layer_id(1). n_sec and log(e_sec) are removed — they are now +# *outputs* predicted by Stage 1, not conditioning inputs. +COND_DIM = 8 -# The two energy columns are additive-log-ratio (ALR) coordinates of the -# deposit/secondary/post energy simplex (fractions of pre_E that sum to 1), -# referenced to the post-energy fraction — see giant.data.transforms -# .energy_simplex_encode/.energy_simplex_decode. They replace the former -# independent log_delta_e / log_edep targets so energy conservation holds by -# construction after decoding. +# Maximum number of secondary slots. From data: max(n_sec)=14 in PbWO4 dataset; +# K_MAX=15 covers it with one spare slot. +K_MAX = 15 + +# Per-slot secondary target dimension: 1 (stick-breaking logit) + 3 (local dir) + +# EMB_DIM (continuous type embedding). EMB_DIM must match DenoisingMLP.emb_dim. +# Default emb_dim=16 → SEC_SLOT_DIM=20. +SEC_SLOT_DIM = 20 # 1 + 3 + 16 +EMB_DIM = 16 # must match model emb_dim default + +# Flattened Stage-2 target dimension +SEC_DIM = K_MAX * SEC_SLOT_DIM # 15 * 20 = 300 + +# Stage-1 9D target names (unchanged from energy-conservation PoC) LOCAL_TARGET_NAMES = [ "log_step_length", "edep_logit", diff --git a/giant/data/dataset.py b/giant/data/dataset.py index 05feed6..433c5f3 100644 --- a/giant/data/dataset.py +++ b/giant/data/dataset.py @@ -4,53 +4,12 @@ from pathlib import Path import numpy as np import torch -from torch.utils.data import Dataset, IterableDataset +from torch.utils.data import IterableDataset from giant.data.loader import iter_file_chunks from giant.data.transforms import Normalizer, build_features -class StepsDataset(Dataset): - def __init__( - self, - cond_cont: np.ndarray, - cond_cat: np.ndarray, - target: np.ndarray, - ) -> None: - self.cond_cont = torch.from_numpy(cond_cont).float() - self.cond_cat = torch.from_numpy(cond_cat).long() - self.target = torch.from_numpy(target).float() - - def __len__(self) -> int: - return len(self.target) - - def __getitem__(self, index): - return self.cond_cont[index], self.cond_cat[index], self.target[index] - - -def train_val_split( - data: dict, - cond_cont: np.ndarray, - cond_cat: np.ndarray, - target: np.ndarray, - val_fraction: float = 0.1, - seed: int = 42, -) -> tuple[StepsDataset, StepsDataset]: - rng = np.random.default_rng(seed) - unique_events = np.unique(data["event_id"]) - rng.shuffle(unique_events) - n_val = max(1, int(len(unique_events) * val_fraction)) - val_events = set(unique_events[:n_val].tolist()) - - val_mask = np.array([e in val_events for e in data["event_id"]]) - train_mask = ~val_mask - - return ( - StepsDataset(cond_cont[train_mask], cond_cat[train_mask], target[train_mask]), - StepsDataset(cond_cont[val_mask], cond_cat[val_mask], target[val_mask]), - ) - - def make_event_split( all_event_ids: np.ndarray, val_fraction: float = 0.1, @@ -75,6 +34,16 @@ class StreamingStepsDataset(IterableDataset): Yields whole batches (use with `DataLoader(..., batch_size=None)`) rather than single rows, so the batch is assembled with vectorized numpy slicing instead of a per-row Python loop in the default collate. + + Each batch is a tuple: + (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx) + where: + cond_cont: (B, COND_DIM) float32 + cond_cat: (B, 2) int64 + target_s1: (B, 9) float32 — normalised Stage-1 primary target + n_sec: (B,) int64 — true secondary count per step + sec_cont: (B, K_MAX, 4) float32 — [stick_logit, local_dir] per slot + sec_pdg_idx: (B, K_MAX) int64 — PDG model-index per secondary slot """ def __init__( @@ -91,13 +60,12 @@ class StreamingStepsDataset(IterableDataset): ) -> None: self.files = list(files) self.split_events = split_events - self._events_arr = np.array(sorted(split_events)) # for np.isin + self._events_arr = np.array(sorted(split_events)) self.pdg_map = pdg_map self.mat_map = mat_map self.cond_normalizer = cond_normalizer self.target_normalizer = target_normalizer self.batch_size = batch_size - # Buffer must hold at least one batch or we could never emit one. self.shuffle_buffer = max(shuffle_buffer, batch_size) self.shuffle = shuffle @@ -114,6 +82,9 @@ class StreamingStepsDataset(IterableDataset): buf_cont: list[np.ndarray] = [] buf_cat: list[np.ndarray] = [] buf_tgt: list[np.ndarray] = [] + buf_nsec: list[np.ndarray] = [] + buf_sec: list[np.ndarray] = [] + buf_spdg: list[np.ndarray] = [] buf_n = 0 for path in files: @@ -123,43 +94,57 @@ class StreamingStepsDataset(IterableDataset): continue chunk = {k: v[mask] for k, v in chunk.items()} - cond_cont, cond_cat, target, _, _ = build_features( - chunk, - self.pdg_map, - self.mat_map, - cond_normalizer=self.cond_normalizer, - target_normalizer=self.target_normalizer, + cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, _, _ = ( + build_features( + chunk, + self.pdg_map, + self.mat_map, + cond_normalizer=self.cond_normalizer, + target_normalizer=self.target_normalizer, + ) ) buf_cont.append(cond_cont) buf_cat.append(cond_cat) - buf_tgt.append(target) + buf_tgt.append(target_s1) + buf_nsec.append(n_sec) + buf_sec.append(sec_cont) + buf_spdg.append(sec_pdg_idx) buf_n += len(cond_cont) if buf_n >= self.shuffle_buffer: - buf_cont, buf_cat, buf_tgt, buf_n = yield from self._flush( - buf_cont, buf_cat, buf_tgt, final=False + buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, buf_n = ( + yield from self._flush( + buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, + final=False, + ) ) if buf_n > 0: - yield from self._flush(buf_cont, buf_cat, buf_tgt, final=True) + yield from self._flush( + buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, final=True + ) def _flush( self, buf_cont: list[np.ndarray], buf_cat: list[np.ndarray], buf_tgt: list[np.ndarray], + buf_nsec: list[np.ndarray], + buf_sec: list[np.ndarray], + buf_spdg: list[np.ndarray], final: bool, ): - """Yield full batches of `batch_size`; carry any remainder back to the caller. - - All batching is done via vectorized numpy slicing (no per-row Python loop). - """ cont = np.concatenate(buf_cont) cat = np.concatenate(buf_cat) tgt = np.concatenate(buf_tgt) + nsec = np.concatenate(buf_nsec) + sec = np.concatenate(buf_sec) + spdg = np.concatenate(buf_spdg) + if self.shuffle: idx = np.random.permutation(len(cont)) cont, cat, tgt = cont[idx], cat[idx], tgt[idx] + nsec, sec, spdg = nsec[idx], sec[idx], spdg[idx] bs = self.batch_size n = len(cont) @@ -170,9 +155,16 @@ class StreamingStepsDataset(IterableDataset): torch.from_numpy(cont[start:end]).float(), torch.from_numpy(cat[start:end]).long(), torch.from_numpy(tgt[start:end]).float(), + torch.from_numpy(nsec[start:end]).long(), + torch.from_numpy(sec[start:end]).float(), + torch.from_numpy(spdg[start:end]).long(), ) if final: - return [], [], [], 0 + return [], [], [], [], [], [], 0 rem = n_full * bs - return [cont[rem:]], [cat[rem:]], [tgt[rem:]], n - rem + return ( + [cont[rem:]], [cat[rem:]], [tgt[rem:]], + [nsec[rem:]], [sec[rem:]], [spdg[rem:]], + n - rem, + ) diff --git a/giant/data/loader.py b/giant/data/loader.py index ceeb195..fa0ba79 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -40,8 +40,52 @@ def find_parquet_files(path: str | Path) -> list[Path]: return [p] +def _pad_list_col(series: pd.Series, K: int, fill: float = 0.0) -> np.ndarray: + """Pad / truncate a list-valued Series to fixed width K → (N, K) float32.""" + out = np.full((len(series), K), fill, dtype=np.float32) + for i, lst in enumerate(series): + if lst is not None and len(lst) > 0: + n = min(len(lst), K) + out[i, :n] = lst[:n] + return out + + +def _pad_list_col_int(series: pd.Series, K: int, fill: int = 0) -> np.ndarray: + """Pad / truncate a list-valued integer Series to fixed width K → (N, K) int64.""" + out = np.full((len(series), K), fill, dtype=np.int64) + for i, lst in enumerate(series): + if lst is not None and len(lst) > 0: + n = min(len(lst), K) + out[i, :n] = lst[:n] + return out + + +def _pad_dir_col( + dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int +) -> np.ndarray: + """Pad three list-valued direction columns → (N, K, 3) float32. + + Padding direction defaults to (0,0,1) (forward) so it is a valid unit vector. + """ + N = len(dx) + out = np.zeros((N, K, 3), dtype=np.float32) + out[:, :, 2] = 1.0 + for i in range(N): + lx, ly, lz = dx.iloc[i], dy.iloc[i], dz.iloc[i] + if lx is not None and len(lx) > 0: + n = min(len(lx), K) + out[i, :n, 0] = lx[:n] + out[i, :n, 1] = ly[:n] + out[i, :n, 2] = lz[:n] + return out + + def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]: - return { + from giant.constants import K_MAX + + has_sec_lists = "sec_E_list" in df.columns + + d: dict[str, np.ndarray] = { "event_id": df["event_id"].to_numpy(), "pdg": df["pdg"].to_numpy(dtype=np.int32), "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32), @@ -59,6 +103,15 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]: "post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32), } + if has_sec_lists: + d["sec_E_list"] = _pad_list_col(df["sec_E_list"], K_MAX) + d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], K_MAX) + d["sec_dir_list"] = _pad_dir_col( + df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], K_MAX + ) + + return d + def load_steps(path: str | Path) -> dict[str, np.ndarray]: return _df_to_dict(pd.read_parquet(path)) diff --git a/giant/data/transforms.py b/giant/data/transforms.py index 692eb50..7141f94 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -227,6 +227,114 @@ def inv_local_frame_rotation( ) +_STICK_LOGIT_CLIP = 10.0 # logit value used for the last valid secondary slot + + +def encode_secondaries( + sec_E_list: np.ndarray, + sec_dir_list: np.ndarray, + sec_valid: np.ndarray, + e_sec: np.ndarray, + pre_dir: np.ndarray, +) -> np.ndarray: + """Encode per-secondary attributes into continuous per-slot targets. + + Secondaries must already be sorted descending by energy (as stored in the + parquet). Returns sec_cont of shape (N, K_MAX, 4): + slot[i] = [stick_break_logit, local_dir_x, local_dir_y, local_dir_z] + + Stick-breaking logit: for slot i, f_i = E_i / remaining_budget, where + remaining_budget = e_sec - sum(E_0..E_{i-1}). The logit is log(f/(1-f)), + clipped to ±_STICK_LOGIT_CLIP. The last valid slot gets +_STICK_LOGIT_CLIP + (takes the full remaining budget). Padding slots get 0. + + sec_pdg_idx (integer) is not processed here — kept separate so the loss + function can look up the embedding table at training time. + """ + N, K = sec_E_list.shape + e_sec = np.asarray(e_sec, dtype=np.float64) + + stick_logits = np.zeros((N, K), dtype=np.float32) + for i in range(K): + if i == 0: + remaining = e_sec + else: + remaining = np.maximum(e_sec - sec_E_list[:, :i].sum(axis=1), _EPS) + f = np.clip(sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS) + logit = np.log(f / (1.0 - f)).astype(np.float32) + # Last valid slot: give it the full remaining budget + is_last = sec_valid[:, i] & ~(sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool)) + logit = np.where(is_last, _STICK_LOGIT_CLIP, logit) + logit = np.where(sec_valid[:, i], np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP), 0.0) + stick_logits[:, i] = logit.astype(np.float32) + + # Rotate each slot's direction into the local frame of the primary. + # pre_dir is broadcast across all K slots. + dir_local = np.zeros((N, K, 3), dtype=np.float32) + for i in range(K): + # Only rotate valid slots; leave padded slots as (0,0,1) or whatever. + valid_mask = sec_valid[:, i] + if valid_mask.any(): + dir_local[valid_mask, i] = local_frame_rotation( + pre_dir[valid_mask], sec_dir_list[valid_mask, i] + ) + + sec_cont = np.concatenate( + [stick_logits[:, :, None], dir_local], axis=-1 + ) # (N, K, 4) + return sec_cont.astype(np.float32) + + +def decode_secondaries( + sec_cont: np.ndarray, + sec_pdg_pred: np.ndarray, + n_sec: np.ndarray, + e_sec: np.ndarray, + pre_dir: np.ndarray, + pdg_map_inv: dict[int, int], +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Inverse of encode_secondaries: continuous targets → physical secondary attrs. + + sec_cont: (N, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z] + sec_pdg_pred: (N, K_MAX) integer PDG indices (from nearest-neighbor snap) + n_sec: (N,) integer secondary counts + e_sec: (N,) total secondary energy budget [MeV] + pre_dir: (N, 3) pre-step world-frame direction + pdg_map_inv: maps model index → PDG code + + Returns (sec_E, sec_dir_world, sec_pdg_code, sec_valid) each shape (N, K_MAX). + """ + N, K, _ = sec_cont.shape + stick_logits = sec_cont[:, :, 0] # (N, K) + dir_local = sec_cont[:, :, 1:] # (N, K, 3) + + fractions = 1.0 / (1.0 + np.exp(-stick_logits.astype(np.float64))) + + sec_E = np.zeros((N, K), dtype=np.float32) + e_sec = np.asarray(e_sec, dtype=np.float64) + remaining = e_sec.copy() + for i in range(K): + sec_E[:, i] = (fractions[:, i] * remaining).astype(np.float32) + remaining = np.maximum(remaining - sec_E[:, i].astype(np.float64), 0.0) + + sec_valid = np.arange(K)[None, :] < n_sec[:, None] # (N, K) + + sec_dir_world = np.zeros((N, K, 3), dtype=np.float32) + for i in range(K): + valid = sec_valid[:, i] + if valid.any(): + sec_dir_world[valid, i] = inv_local_frame_rotation( + pre_dir[valid], dir_local[valid, i] + ) + + sec_pdg_code = np.array( + [[pdg_map_inv.get(int(sec_pdg_pred[n, i]), 0) for i in range(K)] for n in range(N)], + dtype=np.int32, + ) + + return sec_E, sec_dir_world, sec_pdg_code, sec_valid + + def build_cond_features( data: dict[str, np.ndarray], pdg_map: dict[int, int], @@ -240,8 +348,6 @@ def build_cond_features( log_transform(data["pre_E"]), data["pre_dir"], data["layer_id"].astype(np.float32), - data["n_sec"].astype(np.float32), - log_transform(data["e_sec"]), ] ).astype(np.float32) @@ -262,11 +368,17 @@ def build_features( cond_normalizer: Normalizer | None = None, target_normalizer: Normalizer | None = None, fit: bool = False, -) -> tuple[np.ndarray, np.ndarray, np.ndarray, Normalizer | None, Normalizer | None]: - """Assemble (cond_cont, cond_cat, target) arrays ready for StepsDataset. +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, Normalizer | None, Normalizer | None]: + """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx) arrays. - When fit=True, new Normalizers are fitted on the supplied arrays. + target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1) + n_sec: (N,) integer secondary counts (target for n_sec head) + sec_cont: (N, K_MAX, 4) continuous secondary targets [stick_logit, dir_local] + sec_pdg_idx: (N, K_MAX) integer PDG model-indices; used to look up embedding + targets in the training loop """ + from giant.constants import K_MAX + 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"]) @@ -274,9 +386,9 @@ def build_features( energy_z = energy_simplex_encode( data["edep"], data["e_sec"], data["post_E"], data["pre_E"] - ) # (N, 2): ALR coords of the deposit/secondary/post energy simplex + ) # (N, 2) - target = np.column_stack( + target_s1 = np.column_stack( [ log_transform(data["step_length"]), energy_z, @@ -285,28 +397,44 @@ def build_features( ] ).astype(np.float32) # (N, 9) + # Phase 2: conditioning drops n_sec and log(e_sec) cond_cont = np.column_stack( [ data["pre_pos"], log_transform(data["pre_E"]), data["pre_dir"], data["layer_id"].astype(np.float32), - data["n_sec"].astype(np.float32), - log_transform(data["e_sec"]), ] - ).astype(np.float32) # (N, COND_DIM) + ).astype(np.float32) # (N, COND_DIM=8) pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64) mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64) cond_cat = np.column_stack([pdg_idx, mat_idx]) # (N, 2) + n_sec = data["n_sec"].astype(np.int64) # (N,) + + # Secondary continuous targets + sec_E_list = data.get("sec_E_list") + sec_dir_list = data.get("sec_dir_list") + sec_pdg_idx = data.get("sec_pdg_idx") + + if sec_E_list is not None and sec_dir_list is not None and sec_pdg_idx is not None: + sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None] # (N, K_MAX) + sec_cont = encode_secondaries( + sec_E_list, sec_dir_list, sec_valid, data["e_sec"], data["pre_dir"] + ) # (N, K_MAX, 4) + else: + N = len(n_sec) + sec_cont = np.zeros((N, K_MAX, 4), dtype=np.float32) + sec_pdg_idx = np.zeros((N, K_MAX), dtype=np.int64) + if fit: cond_normalizer = Normalizer().fit(cond_cont) - target_normalizer = Normalizer().fit(target) + target_normalizer = Normalizer().fit(target_s1) if cond_normalizer is not None: cond_cont = cond_normalizer.transform(cond_cont) if target_normalizer is not None: - target = target_normalizer.transform(target) + target_s1 = target_normalizer.transform(target_s1) - return cond_cont, cond_cat, target, cond_normalizer, target_normalizer + return cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, cond_normalizer, target_normalizer diff --git a/giant/model/network.py b/giant/model/network.py index 2a21b8f..b64b48e 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -3,7 +3,7 @@ import math import torch import torch.nn as nn -from giant.constants import COND_DIM, X_DIM +from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM class SinusoidalEmbedding(nn.Module): @@ -70,6 +70,12 @@ class ResBlock(nn.Module): class DenoisingMLP(nn.Module): + """Stage-1 model: predicts the 9D primary post-step vector field + n_sec logits. + + The n_sec head runs on the condition encoding only (no diffusion noise), + so it can be called at inference time independently via `predict_n_sec`. + """ + def __init__( self, pdg_vocab: int, @@ -81,6 +87,7 @@ class DenoisingMLP(nn.Module): cond_out_dim: int = 128, x_dim: int = X_DIM, dropout: float = 0.1, + k_max: int = K_MAX, ) -> None: super().__init__() self.time_emb = SinusoidalEmbedding(time_dim) @@ -99,6 +106,13 @@ class DenoisingMLP(nn.Module): ] ) self.out_proj = nn.Linear(hidden_dim, x_dim) + # Predicts n_sec as classification over {0, 1, ..., k_max}. + # Applied to the condition encoding (not the diffused latent). + self.n_sec_head = nn.Sequential( + nn.Linear(cond_out_dim, hidden_dim // 2), + nn.SiLU(), + nn.Linear(hidden_dim // 2, k_max + 1), + ) def forward( self, @@ -107,9 +121,117 @@ class DenoisingMLP(nn.Module): cond_cont: torch.Tensor, cond_cat: torch.Tensor, ) -> torch.Tensor: - t_emb = self.time_emb(t) # (B, time_dim) + t_emb = self.time_emb(t) # (B, time_dim) c_emb = self.cond_enc(cond_cont, cond_cat) # (B, cond_out_dim) - cond = torch.cat([t_emb, c_emb], dim=-1) # (B, time_dim+cond_out_dim) + cond = torch.cat([t_emb, c_emb], dim=-1) + x = self.input_proj(x_t) + for block in self.blocks: + x = block(x, cond) + return self.out_proj(x) + + def predict_n_sec( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + """Return n_sec logits (B, K_MAX+1) from conditioning alone.""" + c_emb = self.cond_enc(cond_cont, cond_cat) + return self.n_sec_head(c_emb) + + def pdg_embedding_weight(self) -> torch.Tensor: + """Return the PDG embedding table weights for secondary type targets.""" + return self.cond_enc.pdg_emb.weight + + +class SecondaryConditionEncoder(nn.Module): + """Encodes pre-step conditioning + Stage-1 output for the secondary decoder.""" + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + emb_dim: int = 16, + cond_out_dim: int = 128, + stage1_dim: int = X_DIM, + stage1_proj_dim: int = 64, + out_dim: int = 128, + ) -> None: + super().__init__() + self.base = ConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + out_dim=cond_out_dim, + ) + self.stage1_proj = nn.Linear(stage1_dim, stage1_proj_dim) + fused_dim = cond_out_dim + stage1_proj_dim + self.fuse = nn.Sequential( + nn.Linear(fused_dim, out_dim), + nn.SiLU(), + ) + + def forward( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + ) -> torch.Tensor: + base = self.base(cond_cont, cond_cat) # (B, cond_out_dim) + s1 = self.stage1_proj(stage1_out).tanh() # (B, stage1_proj_dim) + return self.fuse(torch.cat([base, s1], dim=-1)) # (B, out_dim) + + +class SecondaryDecoder(nn.Module): + """Stage-2 model: predicts vector field over K_MAX secondary slots simultaneously. + + Each slot encodes (stick_break_logit, local_dir_3D, type_emb) for one + secondary ordered by descending energy. Padded slots are masked from loss. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + hidden_dim: int = 256, + n_blocks: int = 6, + emb_dim: int = 16, + time_dim: int = 64, + cond_out_dim: int = 128, + stage1_proj_dim: int = 64, + sec_dim: int = SEC_DIM, + dropout: float = 0.1, + ) -> None: + super().__init__() + self.time_emb = SinusoidalEmbedding(time_dim) + self.cond_enc = SecondaryConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + cond_out_dim=cond_out_dim, + stage1_proj_dim=stage1_proj_dim, + out_dim=cond_out_dim, + ) + merged_cond_dim = time_dim + cond_out_dim + self.input_proj = nn.Linear(sec_dim, hidden_dim) + self.blocks = nn.ModuleList( + [ + ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) + for _ in range(n_blocks) + ] + ) + self.out_proj = nn.Linear(hidden_dim, sec_dim) + + def forward( + self, + x_t: torch.Tensor, + t: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + ) -> torch.Tensor: + t_emb = self.time_emb(t) + c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out) + cond = torch.cat([t_emb, c_emb], dim=-1) x = self.input_proj(x_t) for block in self.blocks: x = block(x, cond) diff --git a/giant/model/schedule.py b/giant/model/schedule.py index a4e942f..ed2877c 100644 --- a/giant/model/schedule.py +++ b/giant/model/schedule.py @@ -69,3 +69,37 @@ def flow_matching_loss( u_t = x1 - x0 v_t = model(x_t, t, cond_cont, cond_cat) return F.mse_loss(v_t, u_t) + + +def flow_matching_loss_secondary( + model: torch.nn.Module, + x1: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + sec_mask: torch.Tensor, +) -> torch.Tensor: + """Flow matching loss for the secondary decoder with per-slot masking. + + x1: (B, SEC_DIM) — flattened secondary target (stick_logit, dir, type_emb) + sec_mask: (B, K_MAX) bool — True for valid secondary slots + + Only valid-slot dimensions contribute to the loss; padded slots are zeroed + before averaging, so the loss is not diluted by empty slots. + """ + from giant.constants import SEC_SLOT_DIM + + B = x1.size(0) + t = torch.rand(B, device=x1.device) + x0 = torch.randn_like(x1) + x_t = (1.0 - t.view(-1, 1)) * x0 + t.view(-1, 1) * x1 + u_t = x1 - x0 + v_t = model(x_t, t, cond_cont, cond_cat, stage1_out) + + # Expand mask: (B, K_MAX) → (B, K_MAX * SEC_SLOT_DIM) + mask_expanded = ( + sec_mask.float().unsqueeze(-1).expand(-1, -1, SEC_SLOT_DIM).reshape(B, -1) + ) + err = (v_t - u_t) ** 2 + denom = mask_expanded.sum().clamp(min=1) + return (err * mask_expanded).sum() / denom diff --git a/giant/pipeline.py b/giant/pipeline.py index b4ef3f2..593455c 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -5,7 +5,7 @@ import torch from torch.utils.data import DataLoader from giant import config -from giant.constants import COND_DIM, X_DIM +from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_SLOT_DIM, X_DIM from giant.data.loader import ( find_parquet_files, load_event_ids, @@ -14,7 +14,7 @@ from giant.data.loader import ( ) from giant.data.transforms import build_features, _WelfordAccumulator from giant.data.dataset import make_event_split, StreamingStepsDataset -from giant.model.network import DenoisingMLP +from giant.model.network import DenoisingMLP, SecondaryDecoder from giant.train import train as run_training @@ -63,9 +63,11 @@ def run_train_job( if not mask.any(): continue chunk_tr = {k: v[mask] for k, v in chunk.items()} - cond_cont, _, target, _, _ = build_features(chunk_tr, pdg_map, mat_map) + cond_cont, _, target_s1, _n_sec, _sec_cont, _sec_pdg, _, _ = build_features( + chunk_tr, pdg_map, mat_map + ) cond_acc.update(cond_cont) - tgt_acc.update(target) + tgt_acc.update(target_s1) cond_norm = cond_acc.to_normalizer() tgt_norm = tgt_acc.to_normalizer() @@ -91,8 +93,6 @@ def run_train_job( shuffle=False, ) - # Dataset yields whole batches already, so batch_size=None tells DataLoader - # to pass them through instead of re-collating row-by-row in Python. pin = device.type == "cuda" train_loader = DataLoader( train_ds, @@ -107,15 +107,34 @@ def run_train_job( pin_memory=pin, ) - model = DenoisingMLP( + emb_dim = m.get("emb_dim", EMB_DIM) + # SEC_SLOT_DIM must match constants (1 stick + 3 dir + emb_dim) + assert SEC_SLOT_DIM == 1 + 3 + emb_dim, ( + f"SEC_SLOT_DIM={SEC_SLOT_DIM} must equal 1+3+emb_dim={1+3+emb_dim}; " + "update giant/constants.py if emb_dim changed" + ) + + stage1_model = DenoisingMLP( pdg_vocab=len(pdg_map), mat_vocab=len(mat_map), hidden_dim=m["hidden_dim"], n_blocks=m["n_blocks"], - emb_dim=m["emb_dim"], + emb_dim=emb_dim, + dropout=m["dropout"], + k_max=K_MAX, + ) + sec_decoder = SecondaryDecoder( + pdg_vocab=len(pdg_map), + mat_vocab=len(mat_map), + hidden_dim=m["hidden_dim"], + n_blocks=m["n_blocks"], + emb_dim=emb_dim, dropout=m["dropout"], ) - echo(f"model: {sum(p.numel() for p in model.parameters()):,} parameters") + echo( + f"stage1: {sum(p.numel() for p in stage1_model.parameters()):,} parameters | " + f"sec_decoder: {sum(p.numel() for p in sec_decoder.parameters()):,} parameters" + ) out_dir.mkdir(parents=True, exist_ok=True) meta = config.build_run_meta( @@ -134,12 +153,15 @@ def run_train_job( "mat_vocab": len(mat_map), "hidden_dim": m["hidden_dim"], "n_blocks": m["n_blocks"], - "emb_dim": m["emb_dim"], + "emb_dim": emb_dim, "dropout": m["dropout"], + "k_max": K_MAX, + "sec_slot_dim": SEC_SLOT_DIM, } run_training( - model=model, + stage1_model=stage1_model, + sec_decoder=sec_decoder, train_loader=train_loader, val_loader=val_loader, mode=t["mode"], @@ -148,6 +170,8 @@ def run_train_job( warmup_epochs=t["warmup_epochs"], device=device, out_dir=out_dir, + lambda_nsec=t.get("lambda_nsec", 0.1), + lambda_s2=t.get("lambda_s2", 1.0), normalizer_dict={"cond": cond_norm.to_dict(), "target": tgt_norm.to_dict()}, pdg_map={str(k): v for k, v in pdg_map.items()}, mat_map={str(k): v for k, v in mat_map.items()}, diff --git a/giant/sample.py b/giant/sample.py index e922057..c2480a1 100644 --- a/giant/sample.py +++ b/giant/sample.py @@ -1,6 +1,6 @@ import torch -from giant.constants import X_DIM +from giant.constants import K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM @torch.no_grad() @@ -9,8 +9,13 @@ def sample_flow( cond_cont: torch.Tensor, cond_cat: torch.Tensor, steps: int = 10, -) -> torch.Tensor: - """Euler integration of the learned vector field from t=0 to t=1.""" +) -> tuple[torch.Tensor, torch.Tensor]: + """Euler integration of the Stage-1 vector field from t=0 to t=1. + + Returns (primary_sample, n_sec_pred): + primary_sample: (B, X_DIM) — normalised 9D primary post-step output + n_sec_pred: (B,) int64 — predicted secondary count + """ model.eval() B = cond_cont.size(0) device = cond_cont.device @@ -20,7 +25,63 @@ def sample_flow( t = torch.full((B,), i * dt, device=device) v = model(x, t, cond_cont, cond_cat) x = x + v * dt - return x + n_sec_logits = model.predict_n_sec(cond_cont, cond_cat) + n_sec_pred = n_sec_logits.argmax(dim=-1) + return x, n_sec_pred + + +@torch.no_grad() +def sample_secondaries( + sec_decoder: torch.nn.Module, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + n_sec_pred: torch.Tensor, + steps: int = 10, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Euler integration of the Stage-2 vector field; return raw slot outputs. + + n_sec_pred: (B,) int64 — number of valid secondaries per step + + Returns (sec_cont, sec_type_emb, sec_valid): + sec_cont: (B, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z] + sec_type_emb: (B, K_MAX, emb_dim) — predicted type embedding per slot + sec_valid: (B, K_MAX) bool — True for slots i < n_sec_pred + """ + sec_decoder.eval() + B = cond_cont.size(0) + device = cond_cont.device + + x = torch.randn(B, SEC_DIM, device=device) + dt = 1.0 / steps + for i in range(steps): + t = torch.full((B,), i * dt, device=device) + v = sec_decoder(x, t, cond_cont, cond_cat, stage1_out) + x = x + v * dt + + x_slots = x.view(B, K_MAX, SEC_SLOT_DIM) + sec_cont = x_slots[:, :, :4] + sec_type_emb = x_slots[:, :, 4:] + + sec_valid = ( + torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1) + ) + return sec_cont, sec_type_emb, sec_valid + + +def snap_type_to_pdg_idx( + sec_type_emb: torch.Tensor, + pdg_emb_weight: torch.Tensor, +) -> torch.Tensor: + """Nearest-neighbour snap: predicted type embedding → PDG model-index. + + sec_type_emb: (B, K_MAX, emb_dim) + Returns (B, K_MAX) int64 with model-indices. + """ + B, K, D = sec_type_emb.shape + flat = sec_type_emb.reshape(-1, D) + dists = torch.cdist(flat.float(), pdg_emb_weight.float()) + return dists.argmin(dim=-1).reshape(B, K) @torch.no_grad() @@ -29,8 +90,8 @@ def sample_ddpm( cond_cont: torch.Tensor, cond_cat: torch.Tensor, schedule, -) -> torch.Tensor: - """Full DDPM ancestral sampling (T reverse steps).""" +) -> tuple[torch.Tensor, torch.Tensor]: + """Full DDPM ancestral sampling (T reverse steps). Returns (sample, n_sec_pred).""" model.eval() B = cond_cont.size(0) device = cond_cont.device @@ -46,7 +107,9 @@ def sample_ddpm( x = (1.0 / alpha.sqrt()) * ( x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred ) + beta.sqrt() * z - return x + n_sec_logits = model.predict_n_sec(cond_cont, cond_cat) + n_sec_pred = n_sec_logits.argmax(dim=-1) + return x, n_sec_pred @torch.no_grad() @@ -56,8 +119,8 @@ def sample_ddim( cond_cat: torch.Tensor, schedule, steps: int = 50, -) -> torch.Tensor: - """DDIM deterministic sampling (Song et al. 2020) with `steps` substeps.""" +) -> tuple[torch.Tensor, torch.Tensor]: + """DDIM deterministic sampling (Song et al. 2020). Returns (sample, n_sec_pred).""" model.eval() B = cond_cont.size(0) device = cond_cont.device @@ -75,4 +138,6 @@ def sample_ddim( ab_prev = torch.ones(1, device=device) x0_pred = (x - (1.0 - ab_t).sqrt() * eps_pred) / ab_t.sqrt() x = ab_prev.sqrt() * x0_pred + (1.0 - ab_prev).sqrt() * eps_pred - return x + n_sec_logits = model.predict_n_sec(cond_cont, cond_cat) + n_sec_pred = n_sec_logits.argmax(dim=-1) + return x, n_sec_pred diff --git a/giant/train.py b/giant/train.py index 894c47b..b8e81b8 100644 --- a/giant/train.py +++ b/giant/train.py @@ -8,14 +8,31 @@ from types import FrameType from typing import Callable import torch +import torch.nn.functional as F import torch.optim as optim from torch.utils.data import DataLoader from tqdm import tqdm -from giant.model.schedule import CosineSchedule, flow_matching_loss +from giant.model.schedule import ( + CosineSchedule, + flow_matching_loss, + flow_matching_loss_secondary, +) from giant.validate import validate_marginals -_METRICS_FIELDS = ["epoch", "train_loss", "val_loss", "lr", "epoch_time_s"] +_METRICS_FIELDS = [ + "epoch", + "train_loss", + "train_loss_s1", + "train_loss_nsec", + "train_loss_s2", + "val_loss", + "val_loss_s1", + "val_loss_nsec", + "val_loss_s2", + "lr", + "epoch_time_s", +] _CATCHABLE_SIGNALS = (signal.SIGINT, signal.SIGTERM) @@ -57,8 +74,80 @@ class _GracefulShutdown: ) +def _build_sec_x1( + sec_cont: torch.Tensor, + sec_pdg_idx: torch.Tensor, + pdg_emb_weight: torch.Tensor, +) -> torch.Tensor: + """Assemble the Stage-2 flow target by appending type embeddings. + + sec_cont: (B, K_MAX, 4) — [stick_logit, dir_local] + sec_pdg_idx: (B, K_MAX) — integer PDG model-indices + pdg_emb_weight: (pdg_vocab, emb_dim) — live embedding table weights + + Returns (B, SEC_DIM) = (B, K_MAX * (4 + emb_dim)). + """ + type_emb = pdg_emb_weight[sec_pdg_idx] # (B, K_MAX, emb_dim) + x1_s2 = torch.cat([sec_cont, type_emb], dim=-1) # (B, K_MAX, 4+emb_dim) + return x1_s2.flatten(1) # (B, SEC_DIM) + + +def _compute_losses( + stage1_model: torch.nn.Module, + sec_decoder: torch.nn.Module, + batch: tuple, + mode: str, + ddpm_schedule, + device: torch.device, + lambda_nsec: float, + lambda_s2: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute (total_loss, L_s1, L_nsec, L_s2) for one batch.""" + cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_idx = batch + cond_cont = cond_cont.to(device) + cond_cat = cond_cat.to(device) + x1_s1 = x1_s1.to(device) + n_sec = n_sec.to(device) + sec_cont = sec_cont.to(device) + sec_pdg_idx = sec_pdg_idx.to(device) + + # Stage-1 flow loss + if mode == "flow": + l_s1 = flow_matching_loss(stage1_model, x1_s1, cond_cont, cond_cat) + else: + assert ddpm_schedule is not None + l_s1 = ddpm_schedule.loss(stage1_model, x1_s1, cond_cont, cond_cat) + + # n_sec classification loss + n_sec_logits = stage1_model.predict_n_sec(cond_cont, cond_cat) + l_nsec = F.cross_entropy(n_sec_logits, n_sec) + + # Stage-2 secondary flow loss + # Use a noiseless Stage-1 target as context (detach to avoid back-prop + # coupling between the two flow paths through the same embedding table). + # The embedding table still receives gradients from the type-embedding loss. + from giant.constants import K_MAX + + pdg_emb_weight = stage1_model.pdg_embedding_weight() + x1_s2 = _build_sec_x1(sec_cont, sec_pdg_idx, pdg_emb_weight) + + sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1) + l_s2 = flow_matching_loss_secondary( + sec_decoder, + x1_s2, + cond_cont, + cond_cat, + x1_s1.detach(), + sec_mask, + ) + + total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2 + return total, l_s1, l_nsec, l_s2 + + def train( - model: torch.nn.Module, + stage1_model: torch.nn.Module, + sec_decoder: torch.nn.Module, train_loader: DataLoader, val_loader: DataLoader, mode: str, @@ -67,6 +156,8 @@ def train( warmup_epochs: int, device: torch.device, out_dir: str | Path, + lambda_nsec: float = 0.1, + lambda_s2: float = 1.0, normalizer_dict: dict | None = None, pdg_map: dict | None = None, mat_map: dict | None = None, @@ -79,8 +170,11 @@ def train( out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) - model = model.to(device) - optimizer = optim.AdamW(model.parameters(), lr=lr) + stage1_model = stage1_model.to(device) + sec_decoder = sec_decoder.to(device) + + all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters()) + optimizer = optim.AdamW(all_params, lr=lr) def _lr_lambda(epoch: int) -> float: if warmup_epochs > 0 and epoch < warmup_epochs: @@ -97,7 +191,8 @@ def train( best_val_loss = float("inf") if resume_path is not None: ckpt = torch.load(resume_path, map_location=device, weights_only=False) - model.load_state_dict(ckpt["model"]) + stage1_model.load_state_dict(ckpt["model"]) + sec_decoder.load_state_dict(ckpt["sec_decoder"]) optimizer.load_state_dict(ckpt["optimizer"]) lr_sched.load_state_dict(ckpt["lr_sched"]) start_epoch = ckpt.get("epoch", 0) + 1 @@ -116,8 +211,12 @@ def train( for epoch in range(start_epoch, epochs + 1): epoch_start = time.monotonic() current_lr = optimizer.param_groups[0]["lr"] - model.train() + stage1_model.train() + sec_decoder.train() train_loss_sum = 0.0 + train_s1_sum = 0.0 + train_nsec_sum = 0.0 + train_s2_sum = 0.0 train_n = 0 ema_loss = 0.0 bar = tqdm( @@ -128,27 +227,26 @@ def train( unit="batch", dynamic_ncols=True, ) - for cond_cont, cond_cat, x1 in bar: - cond_cont = cond_cont.to(device) - cond_cat = cond_cat.to(device) - x1 = x1.to(device) - - if mode == "flow": - loss = flow_matching_loss(model, x1, cond_cont, cond_cat) - else: - assert ddpm_schedule is not None - loss = ddpm_schedule.loss(model, x1, cond_cont, cond_cat) - + for batch in bar: + loss, l_s1, l_nsec, l_s2 = _compute_losses( + stage1_model, sec_decoder, batch, mode, ddpm_schedule, device, + lambda_nsec, lambda_s2, + ) optimizer.zero_grad() loss.backward() - torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + torch.nn.utils.clip_grad_norm_(all_params, 1.0) optimizer.step() + + B = batch[0].size(0) batch_loss = loss.item() - train_loss_sum += batch_loss * x1.size(0) - train_n += x1.size(0) + train_loss_sum += batch_loss * B + train_s1_sum += l_s1.item() * B + train_nsec_sum += l_nsec.item() * B + train_s2_sum += l_s2.item() * B + train_n += B ema_loss = ( batch_loss - if train_n == x1.size(0) + if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss ) bar.set_postfix_str(f"loss={ema_loss:.4f}", refresh=False) @@ -158,28 +256,30 @@ def train( bar.close() if shutdown.requested: - # Mid-epoch: discard the partial epoch rather than persist an - # inconsistent (lr_sched not stepped, no validation) checkpoint. break train_loss = train_loss_sum / max(train_n, 1) lr_sched.step() - model.eval() + stage1_model.eval() + sec_decoder.eval() val_loss_sum = 0.0 + val_s1_sum = 0.0 + val_nsec_sum = 0.0 + val_s2_sum = 0.0 val_n = 0 with torch.no_grad(): - for cond_cont, cond_cat, x1 in val_loader: - cond_cont = cond_cont.to(device) - cond_cat = cond_cat.to(device) - x1 = x1.to(device) - if mode == "flow": - loss = flow_matching_loss(model, x1, cond_cont, cond_cat) - else: - assert ddpm_schedule is not None - loss = ddpm_schedule.loss(model, x1, cond_cont, cond_cat) - val_loss_sum += loss.item() * x1.size(0) - val_n += x1.size(0) + for batch in val_loader: + loss, l_s1, l_nsec, l_s2 = _compute_losses( + stage1_model, sec_decoder, batch, mode, ddpm_schedule, device, + lambda_nsec, lambda_s2, + ) + B = batch[0].size(0) + val_loss_sum += loss.item() * B + val_s1_sum += l_s1.item() * B + val_nsec_sum += l_nsec.item() * B + val_s2_sum += l_s2.item() * B + val_n += B val_loss = val_loss_sum / max(val_n, 1) epoch_time = time.monotonic() - epoch_start @@ -187,14 +287,24 @@ def train( marker = " [best]" if is_best else "" print( f"epoch {epoch:{epoch_w}d}/{epochs}" - f" train {train_loss:.4f} val {val_loss:.4f}" + f" train {train_loss:.4f}" + f" (s1={train_s1_sum/max(train_n,1):.3f}" + f" nsec={train_nsec_sum/max(train_n,1):.3f}" + f" s2={train_s2_sum/max(train_n,1):.3f})" + f" val {val_loss:.4f}" f" lr {current_lr:.2e} {epoch_time:.1f}s{marker}" ) metrics_writer.writerow( { "epoch": epoch, "train_loss": train_loss, + "train_loss_s1": train_s1_sum / max(train_n, 1), + "train_loss_nsec": train_nsec_sum / max(train_n, 1), + "train_loss_s2": train_s2_sum / max(train_n, 1), "val_loss": val_loss, + "val_loss_s1": val_s1_sum / max(val_n, 1), + "val_loss_nsec": val_nsec_sum / max(val_n, 1), + "val_loss_s2": val_s2_sum / max(val_n, 1), "lr": current_lr, "epoch_time_s": epoch_time, } @@ -204,7 +314,7 @@ def train( if validate_every > 0 and epoch % validate_every == 0: print(f"[epoch {epoch}] marginal validation:") validate_marginals( - model, + stage1_model, val_loader, mode=mode, schedule=ddpm_schedule, @@ -213,7 +323,8 @@ def train( ) ckpt: dict = { - "model": model.state_dict(), + "model": stage1_model.state_dict(), + "sec_decoder": sec_decoder.state_dict(), "optimizer": optimizer.state_dict(), "lr_sched": lr_sched.state_dict(), "epoch": epoch, diff --git a/giant/validate.py b/giant/validate.py index afed080..1e5a97b 100644 --- a/giant/validate.py +++ b/giant/validate.py @@ -53,18 +53,21 @@ def validate_marginals( model.eval() all_real, all_gen = [], [] - for i, (cond_cont, cond_cat, x1) in enumerate(val_loader): + for i, batch in enumerate(val_loader): if n_batches is not None and i >= n_batches: break + # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx); + # validate_marginals only needs the Stage-1 primary target. + cond_cont, cond_cat, x1 = batch[0], batch[1], batch[2] cond_cont = cond_cont.to(device) cond_cat = cond_cat.to(device) if mode == "flow": - gen = sample_flow(model, cond_cont, cond_cat, **_kw(steps)) + gen, _n_sec = sample_flow(model, cond_cont, cond_cat, **_kw(steps)) elif mode == "ddpm": - gen = sample_ddpm(model, cond_cont, cond_cat, schedule) + gen, _n_sec = sample_ddpm(model, cond_cont, cond_cat, schedule) else: - gen = sample_ddim(model, cond_cont, cond_cat, schedule, **_kw(steps)) + gen, _n_sec = sample_ddim(model, cond_cont, cond_cat, schedule, **_kw(steps)) all_real.append(x1.numpy()) all_gen.append(gen.cpu().numpy()) diff --git a/scripts/steps_to_parquet.py b/scripts/steps_to_parquet.py index 3b58645..1703fc6 100644 --- a/scripts/steps_to_parquet.py +++ b/scripts/steps_to_parquet.py @@ -19,20 +19,31 @@ import uproot ParquetCompression = Literal["lz4", "uncompressed", "snappy", "gzip", "brotli", "zstd"] -def _add_secondary_energy(df: pl.DataFrame) -> pl.DataFrame: - """Add per-step `e_sec`: total initial kinetic energy of the secondaries born in it. +def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame: + """Add per-step secondary attributes via the parent→child track join. - Each secondary's creation energy is the `pre_E` of that child track's first step - (min `step_no`) in the same event, so for a parent step - `e_sec = Σ over child_track_ids of the child track's first-step pre_E`. Steps that - spawn nothing get 0.0. The full event must be present in `df` (it is — the writer - concatenates every batch before this runs), since a child track's first step can - live in a different read batch than its parent step. + For each step that spawns secondaries, collects each child track's birth + state (from the child track's first step in the same event) and emits: + e_sec float64 — total secondary energy (sum of child first-step pre_E) + sec_E_list list[f64] — per-secondary energy, sorted descending + sec_pdg_list list[i32] — per-secondary PDG code, same order + sec_dx_list list[f64] — per-secondary birth direction x, same order + sec_dy_list list[f64] — per-secondary birth direction y, same order + sec_dz_list list[f64] — per-secondary birth direction z, same order + + Steps with no children get 0.0 / empty lists. The full event must be + present in `df` (it is — the writer concatenates before calling this). """ - first_E = ( + first_step = ( df.sort("step_no") .group_by(["event_id", "track_id"]) - .agg(pl.col("pre_E").first().alias("child_E")) + .agg( + pl.col("pre_E").first().alias("child_E"), + pl.col("pdg").first().alias("child_pdg"), + pl.col("pre_dx").first().alias("child_dx"), + pl.col("pre_dy").first().alias("child_dy"), + pl.col("pre_dz").first().alias("child_dz"), + ) .rename({"track_id": "child_track_id"}) ) @@ -41,17 +52,39 @@ def _add_secondary_energy(df: pl.DataFrame) -> pl.DataFrame: .with_row_index("_step_row") .explode("child_track_ids") .rename({"child_track_ids": "child_track_id"}) - .drop_nulls("child_track_id") # steps with no children explode to a null row + .drop_nulls("child_track_id") ) - summed = ( - exploded.join(first_E, on=["event_id", "child_track_id"], how="left") + + joined = exploded.join(first_step, on=["event_id", "child_track_id"], how="left") + + # Sort each step's secondaries by descending energy, then aggregate into lists + per_step = ( + joined.sort("child_E", descending=True) .group_by("_step_row") - .agg(pl.col("child_E").sum().alias("e_sec")) + .agg( + pl.col("child_E").sum().alias("e_sec"), + pl.col("child_E").alias("sec_E_list"), + pl.col("child_pdg").alias("sec_pdg_list"), + pl.col("child_dx").alias("sec_dx_list"), + pl.col("child_dy").alias("sec_dy_list"), + pl.col("child_dz").alias("sec_dz_list"), + ) ) + + empty_list_f64 = pl.Series("x", [[]], dtype=pl.List(pl.Float64)) + empty_list_i32 = pl.Series("x", [[]], dtype=pl.List(pl.Int32)) + return ( df.with_row_index("_step_row") - .join(summed, on="_step_row", how="left") - .with_columns(pl.col("e_sec").fill_null(0.0).cast(pl.Float64)) + .join(per_step, on="_step_row", how="left") + .with_columns( + pl.col("e_sec").fill_null(0.0).cast(pl.Float64), + pl.col("sec_E_list").fill_null(empty_list_f64), + pl.col("sec_pdg_list").fill_null(empty_list_i32), + pl.col("sec_dx_list").fill_null(empty_list_f64), + pl.col("sec_dy_list").fill_null(empty_list_f64), + pl.col("sec_dz_list").fill_null(empty_list_f64), + ) .drop("_step_row") ) @@ -118,8 +151,8 @@ def convert_steps_to_parquet( # Steps tree carries the parent→child links needed to derive secondary energy; # other trees (e.g. Hits) don't, so only augment when the column is present. if "child_track_ids" in df.columns: - print("\nComputing per-step secondary energy (e_sec) …", end=" ", flush=True) - df = _add_secondary_energy(df) + print("\nComputing per-step secondary attributes …", end=" ", flush=True) + df = _add_secondary_attributes(df) print(f"\nWriting {output_path} …", end=" ", flush=True) df.write_parquet(output_path, compression=compression) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 9e06fac..ad7f677 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -1,67 +1,33 @@ import numpy as np -from giant.data.dataset import StepsDataset, train_val_split +from giant.data.dataset import make_event_split -def _dummy(N=500, n_events=20): +def test_make_event_split_sizes(): rng = np.random.default_rng(42) - data = {"event_id": rng.integers(0, n_events, size=N)} - cond_cont = rng.standard_normal((N, 9)).astype(np.float32) - cond_cat = rng.integers(0, 3, size=(N, 2)).astype(np.int64) - target = rng.standard_normal((N, 6)).astype(np.float32) - return data, cond_cont, cond_cat, target + event_ids = rng.integers(0, 50, size=1000) + train_set, val_set = make_event_split(event_ids, val_fraction=0.2) + unique = np.unique(event_ids) + assert len(train_set) + len(val_set) == len(unique) -def test_dataset_length(): - data, cond_cont, cond_cat, target = _dummy() - assert len(StepsDataset(cond_cont, cond_cat, target)) == len(target) - - -def test_dataset_item_shapes(): - data, cond_cont, cond_cat, target = _dummy() - c, k, t = StepsDataset(cond_cont, cond_cat, target)[0] - assert c.shape == (9,) - assert k.shape == (2,) - assert t.shape == (6,) - - -def test_split_sizes_sum_to_total(): - data, cond_cont, cond_cat, target = _dummy(N=500) - train_ds, val_ds = train_val_split( - data, cond_cont, cond_cat, target, val_fraction=0.2 - ) - assert len(train_ds) + len(val_ds) == 500 - - -def test_split_no_empty_sets(): - data, cond_cont, cond_cat, target = _dummy(N=500, n_events=20) - train_ds, val_ds = train_val_split( - data, cond_cont, cond_cat, target, val_fraction=0.2 - ) - assert len(val_ds) > 0 - assert len(train_ds) > 0 - - -def test_split_event_leakage(): - """Train and val must not share any event_id.""" - N = 1000 - n_events = 50 +def test_make_event_split_no_overlap(): rng = np.random.default_rng(7) - event_ids = rng.integers(0, n_events, size=N) - data = {"event_id": event_ids} - cond_cont = rng.standard_normal((N, 9)).astype(np.float32) - cond_cat = rng.integers(0, 3, size=(N, 2)).astype(np.int64) - target = rng.standard_normal((N, 6)).astype(np.float32) + event_ids = rng.integers(0, 50, size=1000) + train_set, val_set = make_event_split(event_ids, val_fraction=0.2) + assert train_set.isdisjoint(val_set) - train_ds, val_ds = train_val_split( - data, cond_cont, cond_cat, target, val_fraction=0.2 - ) - # Recover which event_ids ended up in each split via the indices - # (The dataset doesn't store event_ids, so we check via the original mask logic) - unique_events = np.unique(event_ids) - rng2 = np.random.default_rng(42) - rng2.shuffle(unique_events) - n_val = max(1, int(len(unique_events) * 0.2)) - val_events = set(unique_events[:n_val].tolist()) - train_events = set(unique_events[n_val:].tolist()) - assert val_events.isdisjoint(train_events) +def test_make_event_split_no_empty_sets(): + rng = np.random.default_rng(0) + event_ids = rng.integers(0, 20, size=500) + train_set, val_set = make_event_split(event_ids, val_fraction=0.2) + assert len(train_set) > 0 + assert len(val_set) > 0 + + +def test_make_event_split_reproducible(): + event_ids = np.arange(100) + a_tr, a_val = make_event_split(event_ids, val_fraction=0.1, seed=42) + b_tr, b_val = make_event_split(event_ids, val_fraction=0.1, seed=42) + assert a_tr == b_tr + assert a_val == b_val diff --git a/tests/test_flow.py b/tests/test_flow.py index d9f5568..b9e7e8c 100644 --- a/tests/test_flow.py +++ b/tests/test_flow.py @@ -39,8 +39,9 @@ def test_sample_flow_shape(): B = 6 cond_cont = torch.randn(B, COND_DIM) cond_cat = torch.zeros(B, 2, dtype=torch.long) - out = sample_flow(_small_model(), cond_cont, cond_cat, steps=5) - assert out.shape == (B, 9) + sample, n_sec = sample_flow(_small_model(), cond_cont, cond_cat, steps=5) + assert sample.shape == (B, 9) + assert n_sec.shape == (B,) def test_ddpm_loss_nonneg(): @@ -55,5 +56,6 @@ def test_sample_ddim_shape(): schedule = CosineSchedule(T=50) cond_cont = torch.randn(B, COND_DIM) 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, 9) + sample, n_sec = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5) + assert sample.shape == (B, 9) + assert n_sec.shape == (B,) diff --git a/tests/test_network.py b/tests/test_network.py index 6934f87..397aa79 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -39,7 +39,9 @@ def test_denoising_mlp_gradients_flow(): t = torch.rand(B) cond_cont = torch.randn(B, COND_DIM) cond_cat = torch.zeros(B, 2, dtype=torch.long) - loss = model(x_t, t, cond_cont, cond_cat).sum() - loss.backward() + # Both paths must be exercised to get gradients through all parameters. + flow_loss = model(x_t, t, cond_cont, cond_cat).sum() + nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum() + (flow_loss + nsec_loss).backward() for name, p in model.named_parameters(): assert p.grad is not None, f"no grad for {name}" diff --git a/tests/test_phase2.py b/tests/test_phase2.py new file mode 100644 index 0000000..e8b4f2d --- /dev/null +++ b/tests/test_phase2.py @@ -0,0 +1,222 @@ +"""Tests for Phase 2: secondary particle prediction.""" +import numpy as np +import pytest +import torch + +from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_DIM, X_DIM +from giant.model.network import DenoisingMLP, SecondaryDecoder +from giant.model.schedule import flow_matching_loss_secondary +from giant.sample import sample_secondaries, snap_type_to_pdg_idx + + +# ── helpers ────────────────────────────────────────────────────────────────── + +def _stage1(pdg=3, mat=2): + return DenoisingMLP(pdg_vocab=pdg, mat_vocab=mat, hidden_dim=32, n_blocks=2) + + +def _sec_decoder(pdg=3, mat=2): + return SecondaryDecoder(pdg_vocab=pdg, mat_vocab=mat, hidden_dim=32, n_blocks=2) + + +def _cond(B=8, pdg=3, mat=2): + cond_cont = torch.randn(B, COND_DIM) + cond_cat = torch.stack( + [torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1 + ) + return cond_cont, cond_cat + + +# ── DenoisingMLP Phase-2 additions ─────────────────────────────────────────── + +def test_predict_n_sec_shape(): + B = 8 + model = _stage1() + cond_cont, cond_cat = _cond(B) + logits = model.predict_n_sec(cond_cont, cond_cat) + assert logits.shape == (B, K_MAX + 1) + + +def test_predict_n_sec_no_nan(): + B = 8 + model = _stage1() + cond_cont, cond_cat = _cond(B) + logits = model.predict_n_sec(cond_cont, cond_cat) + assert torch.isfinite(logits).all() + + +def test_pdg_embedding_weight_shape(): + model = _stage1(pdg=5, mat=2) + w = model.pdg_embedding_weight() + assert w.shape == (5, EMB_DIM) + + +# ── SecondaryDecoder ────────────────────────────────────────────────────────── + +def test_sec_decoder_output_shape(): + B = 8 + decoder = _sec_decoder() + x_t = torch.randn(B, SEC_DIM) + t = torch.rand(B) + cond_cont, cond_cat = _cond(B) + stage1_out = torch.randn(B, X_DIM) + out = decoder(x_t, t, cond_cont, cond_cat, stage1_out) + assert out.shape == (B, SEC_DIM) + + +def test_sec_decoder_no_nan(): + B = 4 + decoder = _sec_decoder() + x_t = torch.randn(B, SEC_DIM) + t = torch.rand(B) + cond_cont, cond_cat = _cond(B) + stage1_out = torch.randn(B, X_DIM) + out = decoder(x_t, t, cond_cont, cond_cat, stage1_out) + assert torch.isfinite(out).all() + + +def test_sec_decoder_gradients(): + B = 4 + decoder = _sec_decoder() + x_t = torch.randn(B, SEC_DIM) + t = torch.rand(B) + cond_cont, cond_cat = _cond(B) + stage1_out = torch.randn(B, X_DIM) + decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward() + for name, p in decoder.named_parameters(): + assert p.grad is not None, f"no grad for {name}" + + +# ── masked flow matching loss ───────────────────────────────────────────────── + +def test_flow_matching_loss_secondary_scalar(): + B, pdg, mat = 8, 3, 2 + decoder = _sec_decoder(pdg, mat) + x1 = torch.randn(B, SEC_DIM) + cond_cont, cond_cat = _cond(B, pdg, mat) + stage1_out = torch.randn(B, X_DIM) + sec_mask = torch.ones(B, K_MAX, dtype=torch.bool) + loss = flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask) + assert loss.shape == () + assert loss.item() >= 0.0 + + +def test_flow_matching_loss_secondary_mask_zeros_padding(): + """Loss with all-zero mask (no valid secondaries) should be 0.""" + B, pdg, mat = 4, 3, 2 + decoder = _sec_decoder(pdg, mat) + x1 = torch.randn(B, SEC_DIM) + cond_cont, cond_cat = _cond(B, pdg, mat) + stage1_out = torch.randn(B, X_DIM) + sec_mask = torch.zeros(B, K_MAX, dtype=torch.bool) + loss = flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask) + assert loss.item() == pytest.approx(0.0, abs=1e-6) + + +def test_flow_matching_loss_secondary_has_grad(): + B, pdg, mat = 4, 3, 2 + decoder = _sec_decoder(pdg, mat) + x1 = torch.randn(B, SEC_DIM) + cond_cont, cond_cat = _cond(B, pdg, mat) + stage1_out = torch.randn(B, X_DIM) + sec_mask = torch.ones(B, K_MAX, dtype=torch.bool) + flow_matching_loss_secondary( + decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask + ).backward() + assert any(p.grad is not None for p in decoder.parameters()) + + +# ── sampling ────────────────────────────────────────────────────────────────── + +def test_sample_secondaries_shapes(): + B, pdg, mat = 6, 3, 2 + decoder = _sec_decoder(pdg, mat) + cond_cont, cond_cat = _cond(B, pdg, mat) + stage1_out = torch.randn(B, X_DIM) + n_sec_pred = torch.randint(0, K_MAX + 1, (B,)) + sec_cont, sec_type_emb, sec_valid = sample_secondaries( + decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=3 + ) + assert sec_cont.shape == (B, K_MAX, 4) + assert sec_type_emb.shape == (B, K_MAX, EMB_DIM) + assert sec_valid.shape == (B, K_MAX) + assert sec_valid.dtype == torch.bool + + +def test_sample_secondaries_valid_mask_matches_n_sec(): + B, pdg, mat = 4, 3, 2 + decoder = _sec_decoder(pdg, mat) + cond_cont, cond_cat = _cond(B, pdg, mat) + stage1_out = torch.randn(B, X_DIM) + n_sec_pred = torch.tensor([0, 1, 3, K_MAX]) + _, _, sec_valid = sample_secondaries( + decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2 + ) + for i, n in enumerate(n_sec_pred.tolist()): + assert sec_valid[i, :n].all() + assert not sec_valid[i, n:].any() + + +def test_snap_type_to_pdg_idx_shape(): + B, pdg_vocab = 4, 5 + emb_weight = torch.randn(pdg_vocab, EMB_DIM) + sec_type_emb = torch.randn(B, K_MAX, EMB_DIM) + idx = snap_type_to_pdg_idx(sec_type_emb, emb_weight) + assert idx.shape == (B, K_MAX) + assert idx.dtype == torch.int64 + assert (idx >= 0).all() and (idx < pdg_vocab).all() + + +# ── encode_secondaries round-trip ───────────────────────────────────────────── + +def test_encode_secondaries_energy_conservation(): + """Decoded stick-breaking fractions must sum to ≈ e_sec.""" + from giant.data.transforms import encode_secondaries + + rng = np.random.default_rng(42) + N = 50 + n_sec = rng.integers(1, 5, size=N) + e_sec = rng.uniform(0.1, 10.0, size=N).astype(np.float32) + + sec_E_list = np.zeros((N, K_MAX), dtype=np.float32) + sec_dir_list = np.zeros((N, K_MAX, 3), dtype=np.float32) + sec_dir_list[:, :, 2] = 1.0 + sec_valid = np.zeros((N, K_MAX), dtype=bool) + for i in range(N): + k = n_sec[i] + energies = rng.dirichlet(np.ones(k)) * e_sec[i] + energies = np.sort(energies)[::-1] + sec_E_list[i, :k] = energies.astype(np.float32) + sec_valid[i, :k] = True + + pre_dir = rng.standard_normal((N, 3)).astype(np.float32) + pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True) + + sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir) + assert sec_cont.shape == (N, K_MAX, 4) + assert np.isfinite(sec_cont).all() + + +def test_encode_secondaries_direction_encoding(): + """Local-frame secondary directions should be unit vectors for valid slots.""" + from giant.data.transforms import encode_secondaries + + rng = np.random.default_rng(7) + N = 20 + e_sec = np.ones(N, dtype=np.float32) * 5.0 + sec_E_list = np.zeros((N, K_MAX), dtype=np.float32) + sec_E_list[:, 0] = 3.0 + sec_E_list[:, 1] = 2.0 + sec_dir_list = rng.standard_normal((N, K_MAX, 3)).astype(np.float32) + norms = np.linalg.norm(sec_dir_list, axis=-1, keepdims=True) + sec_dir_list /= np.where(norms > 0, norms, 1.0) + sec_valid = np.zeros((N, K_MAX), dtype=bool) + sec_valid[:, :2] = True + + pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32) + sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir) + + # dir columns are sec_cont[:, :, 1:4] + local_dirs = sec_cont[:, :2, 1:4] # (N, 2, 3) — valid slots only + norms_out = np.linalg.norm(local_dirs, axis=-1) + np.testing.assert_allclose(norms_out, 1.0, atol=1e-5)