diff --git a/docs/phase2_plan.md b/docs/phase2_plan.md deleted file mode 100644 index 596edec..0000000 --- a/docs/phase2_plan.md +++ /dev/null @@ -1,172 +0,0 @@ -# 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/test-cuda.py b/test-cuda.py deleted file mode 100644 index 2b95e64..0000000 --- a/test-cuda.py +++ /dev/null @@ -1,19 +0,0 @@ -import torch -import torch.version - -print(f"PyTorch version: {torch.__version__}") -print(f"CUDA available: {torch.cuda.is_available()}") - -if torch.cuda.is_available(): - print(f"CUDA version: {torch.version.cuda}") - print(f"Device count: {torch.cuda.device_count()}") - print(f"Device name: {torch.cuda.get_device_name(0)}") - - # Run a small tensor op on the GPU - a = torch.randn(1000, 1000, device="cuda") - b = torch.randn(1000, 1000, device="cuda") - c = a @ b - torch.cuda.synchronize() - print(f"Matrix multiply: OK (result shape {c.shape}, device {c.device})") -else: - print("No CUDA device found — check driver/CUDA installation.")