# 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