giant/docs/phase2_plan.md

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

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:

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)

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