Condition on material/particle physical properties instead of learned embeddings
Adds model.conditioning = "physical" | "embedding": physical mode routes particle mass/charge and material Z_eff/A_eff/density/X0/lambda_int through small MLPs to replace the learned PDG/material embedding tables, so the surrogate generalizes to PDG codes/materials outside the training vocab instead of memorizing it. "embedding" stays available as the comparison baseline (old checkpoints without the key default to it). Stage 2 now regresses a secondary's mass/charge directly against a fixed physics-derived target instead of a learned/snapped embedding, and uses no snapping at inference — the model's raw predicted (mass, charge) is the secondary's physical identity, including for its own further rollout steps. A separate reporting-only nearest-known-PDG lookup (never fed back into the model) populates output pdg columns / the embedding-mode rollout fallback. giant/materials.py's table is populated with Geant4's own built-in NIST constants (Z_eff, A_eff, density, X0, lambda_int), extracted directly from the Geant4 11.4.1 build vendored in minicalosim via G4NistManager rather than hand-typed literature values. G4_LYSO is left unfilled: confirmed (both by runtime lookup and by searching minicalosim's history) that it's never actually a constructed Geant4 material there, only documentation/UI color-map text. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
+11
-21
@@ -43,10 +43,15 @@ def sample_secondaries(
|
||||
|
||||
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
|
||||
Returns (sec_cont, sec_phys, sec_valid):
|
||||
sec_cont: (B, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z]
|
||||
sec_phys: (B, K_MAX, PARTICLE_PHYS_DIM) — predicted [log_mass, charge]
|
||||
per slot (normalised iff the checkpoint's sec_phys
|
||||
normalizer was applied at training time — denormalize
|
||||
before treating as physical units; see
|
||||
giant.data.transforms.decode_secondaries). Used as-is —
|
||||
no snapping to a discrete PDG code.
|
||||
sec_valid: (B, K_MAX) bool — True for slots i < n_sec_pred
|
||||
"""
|
||||
sec_decoder.eval()
|
||||
B = cond_cont.size(0)
|
||||
@@ -61,27 +66,12 @@ def sample_secondaries(
|
||||
|
||||
x_slots = x.view(B, K_MAX, SEC_SLOT_DIM)
|
||||
sec_cont = x_slots[:, :, :4]
|
||||
sec_type_emb = x_slots[:, :, 4:]
|
||||
sec_phys = 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)
|
||||
return sec_cont, sec_phys, sec_valid
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
|
||||
Reference in New Issue
Block a user