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:
+136
-52
@@ -4,21 +4,33 @@ import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_DIM, X_DIM
|
||||
from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, 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
|
||||
from giant.sample import sample_secondaries
|
||||
|
||||
|
||||
# ── helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _stage1(pdg=3, mat=2):
|
||||
return DenoisingMLP(pdg_vocab=pdg, mat_vocab=mat, hidden_dim=32, n_blocks=2)
|
||||
def _stage1(pdg=3, mat=2, conditioning="embedding"):
|
||||
return DenoisingMLP(
|
||||
pdg_vocab=pdg,
|
||||
mat_vocab=mat,
|
||||
hidden_dim=32,
|
||||
n_blocks=2,
|
||||
conditioning=conditioning,
|
||||
)
|
||||
|
||||
|
||||
def _sec_decoder(pdg=3, mat=2):
|
||||
return SecondaryDecoder(pdg_vocab=pdg, mat_vocab=mat, hidden_dim=32, n_blocks=2)
|
||||
def _sec_decoder(pdg=3, mat=2, conditioning="embedding"):
|
||||
return SecondaryDecoder(
|
||||
pdg_vocab=pdg,
|
||||
mat_vocab=mat,
|
||||
hidden_dim=32,
|
||||
n_blocks=2,
|
||||
conditioning=conditioning,
|
||||
)
|
||||
|
||||
|
||||
def _cond(B=8, pdg=3, mat=2):
|
||||
@@ -48,18 +60,35 @@ def test_predict_n_sec_no_nan():
|
||||
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)
|
||||
@pytest.mark.parametrize("conditioning", ["embedding", "physical"])
|
||||
def test_no_pdg_embedding_weight_method(conditioning):
|
||||
"""The Stage-2 species output no longer needs a shared embedding table."""
|
||||
model = _stage1(pdg=5, mat=2, conditioning=conditioning)
|
||||
assert not hasattr(model, "pdg_embedding_weight")
|
||||
|
||||
|
||||
def test_condition_encoder_physical_mode_has_no_embedding_tables():
|
||||
model = _stage1(pdg=5, mat=2, conditioning="physical")
|
||||
assert not hasattr(model.cond_enc, "pdg_emb")
|
||||
assert not hasattr(model.cond_enc, "mat_emb")
|
||||
assert hasattr(model.cond_enc, "particle_mlp")
|
||||
assert hasattr(model.cond_enc, "material_mlp")
|
||||
|
||||
|
||||
def test_condition_encoder_embedding_mode_has_embedding_tables():
|
||||
model = _stage1(pdg=5, mat=2, conditioning="embedding")
|
||||
assert hasattr(model.cond_enc, "pdg_emb")
|
||||
assert hasattr(model.cond_enc, "mat_emb")
|
||||
assert not hasattr(model.cond_enc, "particle_mlp")
|
||||
|
||||
|
||||
# ── SecondaryDecoder ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_sec_decoder_output_shape():
|
||||
@pytest.mark.parametrize("conditioning", ["embedding", "physical"])
|
||||
def test_sec_decoder_output_shape(conditioning):
|
||||
B = 8
|
||||
decoder = _sec_decoder()
|
||||
decoder = _sec_decoder(conditioning=conditioning)
|
||||
x_t = torch.randn(B, SEC_DIM)
|
||||
t = torch.rand(B)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
@@ -144,11 +173,11 @@ def test_sample_secondaries_shapes():
|
||||
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(
|
||||
sec_cont, sec_phys, 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_phys.shape == (B, K_MAX, PARTICLE_PHYS_DIM)
|
||||
assert sec_valid.shape == (B, K_MAX)
|
||||
assert sec_valid.dtype == torch.bool
|
||||
|
||||
@@ -167,16 +196,6 @@ def test_sample_secondaries_valid_mask_matches_n_sec():
|
||||
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 ─────────────────────────────────────────────
|
||||
|
||||
|
||||
@@ -192,6 +211,7 @@ def test_encode_secondaries_energy_conservation():
|
||||
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_pdg_list = np.zeros((N, K_MAX), dtype=np.int64)
|
||||
sec_valid = np.zeros((N, K_MAX), dtype=bool)
|
||||
for i in range(N):
|
||||
k = n_sec[i]
|
||||
@@ -199,12 +219,15 @@ def test_encode_secondaries_energy_conservation():
|
||||
energies = np.sort(energies)[::-1]
|
||||
sec_E_list[i, :k] = energies.astype(np.float32)
|
||||
sec_valid[i, :k] = True
|
||||
sec_pdg_list[i, :k] = 22 # photon — resolvable by giant.particles
|
||||
|
||||
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)
|
||||
sec_cont = encode_secondaries(
|
||||
sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
|
||||
)
|
||||
assert sec_cont.shape == (N, K_MAX, 6)
|
||||
assert np.isfinite(sec_cont).all()
|
||||
|
||||
|
||||
@@ -233,13 +256,54 @@ def test_encode_secondaries_direction_encoding():
|
||||
np.testing.assert_allclose(norms_out, 1.0, atol=1e-5)
|
||||
|
||||
|
||||
def test_encode_secondaries_physical_columns_without_pdg_list():
|
||||
"""Omitting sec_pdg_list zero-fills the physical columns (no crash)."""
|
||||
from giant.data.transforms import encode_secondaries
|
||||
|
||||
N = 3
|
||||
e_sec = np.ones(N, dtype=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)
|
||||
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)
|
||||
np.testing.assert_allclose(sec_cont[:, :, 4:6], 0.0)
|
||||
|
||||
|
||||
def test_encode_secondaries_physical_columns_match_ground_truth_pdg():
|
||||
"""log_mass/charge for a valid slot match giant.particles for that PDG."""
|
||||
from giant.data.transforms import encode_secondaries, log_transform
|
||||
from giant.particles import particle_mass_charge
|
||||
|
||||
N = 1
|
||||
e_sec = np.array([5.0], dtype=np.float32)
|
||||
sec_E_list = np.zeros((N, K_MAX), dtype=np.float32)
|
||||
sec_E_list[0, 0] = 5.0
|
||||
sec_dir_list = np.zeros((N, K_MAX, 3), dtype=np.float32)
|
||||
sec_dir_list[0, 0] = [0, 0, 1]
|
||||
sec_pdg_list = np.zeros((N, K_MAX), dtype=np.int64)
|
||||
sec_pdg_list[0, 0] = 11 # electron
|
||||
sec_valid = np.zeros((N, K_MAX), dtype=bool)
|
||||
sec_valid[0, 0] = True
|
||||
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
sec_cont = encode_secondaries(
|
||||
sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
|
||||
)
|
||||
mass, charge = particle_mass_charge(11)
|
||||
assert sec_cont[0, 0, 4] == pytest.approx(log_transform(np.array([mass]))[0])
|
||||
assert sec_cont[0, 0, 5] == pytest.approx(charge)
|
||||
|
||||
|
||||
# ── decode_secondaries: exact energy conservation ────────────────────────────
|
||||
|
||||
|
||||
def _random_sec_cont(rng, N, stick_logit_scale=1.0):
|
||||
sec_cont = rng.standard_normal((N, K_MAX, 4)).astype(np.float32)
|
||||
sec_cont = rng.standard_normal((N, K_MAX, 6)).astype(np.float32)
|
||||
sec_cont[:, :, 0] *= stick_logit_scale
|
||||
dirs = sec_cont[:, :, 1:]
|
||||
dirs = sec_cont[:, :, 1:4]
|
||||
dirs /= np.linalg.norm(dirs, axis=-1, keepdims=True)
|
||||
return sec_cont
|
||||
|
||||
@@ -257,13 +321,12 @@ def test_decode_secondaries_valid_slots_sum_to_e_sec():
|
||||
rng = np.random.default_rng(0)
|
||||
N = 200
|
||||
sec_cont = _random_sec_cont(rng, N)
|
||||
sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
|
||||
n_sec = rng.integers(0, K_MAX + 1, size=N)
|
||||
e_sec = rng.uniform(0.0, 50.0, size=N).astype(np.float32)
|
||||
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
|
||||
|
||||
sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries(
|
||||
sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22}
|
||||
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
|
||||
sec_cont, n_sec, e_sec, pre_dir
|
||||
)
|
||||
|
||||
valid_sum = (sec_E * sec_valid).sum(axis=1)
|
||||
@@ -283,13 +346,12 @@ def test_decode_secondaries_zero_n_sec_has_zero_energy():
|
||||
rng = np.random.default_rng(1)
|
||||
N = 10
|
||||
sec_cont = _random_sec_cont(rng, N)
|
||||
sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
|
||||
n_sec = np.zeros(N, dtype=np.int64)
|
||||
e_sec = rng.uniform(1.0, 10.0, size=N).astype(np.float32)
|
||||
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
|
||||
|
||||
sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries(
|
||||
sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22}
|
||||
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
|
||||
sec_cont, n_sec, e_sec, pre_dir
|
||||
)
|
||||
|
||||
assert not sec_valid.any()
|
||||
@@ -307,12 +369,11 @@ def test_decode_secondaries_degenerate_row_falls_back_to_even_split():
|
||||
n_sec = np.array([0, 1, 3, K_MAX])
|
||||
for i, k in enumerate(n_sec):
|
||||
sec_cont[i, :k, 0] = -80.0
|
||||
sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
|
||||
e_sec = np.array([0.0, 4.0, 9.0, 30.0], dtype=np.float32)
|
||||
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
|
||||
|
||||
sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries(
|
||||
sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22}
|
||||
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
|
||||
sec_cont, n_sec, e_sec, pre_dir
|
||||
)
|
||||
|
||||
for i, k in enumerate(n_sec):
|
||||
@@ -336,25 +397,48 @@ def test_decode_secondaries_rescale_preserves_relative_shares():
|
||||
sec_cont = _random_sec_cont(rng, N)
|
||||
n_sec = np.array([4])
|
||||
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
|
||||
|
||||
sec_E_small, _, _, sec_valid = decode_secondaries(
|
||||
sec_cont,
|
||||
sec_pdg_pred,
|
||||
n_sec,
|
||||
np.array([5.0], dtype=np.float32),
|
||||
pre_dir,
|
||||
{0: 22},
|
||||
sec_E_small, _, _, _, sec_valid = decode_secondaries(
|
||||
sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir
|
||||
)
|
||||
sec_E_large, _, _, _ = decode_secondaries(
|
||||
sec_cont,
|
||||
sec_pdg_pred,
|
||||
n_sec,
|
||||
np.array([50.0], dtype=np.float32),
|
||||
pre_dir,
|
||||
{0: 22},
|
||||
sec_E_large, _, _, _, _ = decode_secondaries(
|
||||
sec_cont, n_sec, np.array([50.0], dtype=np.float32), pre_dir
|
||||
)
|
||||
|
||||
ratio_small = sec_E_small[0, :4] / sec_E_small[0, 0]
|
||||
ratio_large = sec_E_large[0, :4] / sec_E_large[0, 0]
|
||||
np.testing.assert_allclose(ratio_small, ratio_large, rtol=1e-4)
|
||||
|
||||
|
||||
def test_decode_secondaries_mass_charge_round_trip_with_normalizer():
|
||||
from giant.data.transforms import Normalizer, decode_secondaries, encode_secondaries
|
||||
|
||||
N = 1
|
||||
e_sec = np.array([5.0], dtype=np.float32)
|
||||
sec_E_list = np.zeros((N, K_MAX), dtype=np.float32)
|
||||
sec_E_list[0, 0] = 5.0
|
||||
sec_dir_list = np.zeros((N, K_MAX, 3), dtype=np.float32)
|
||||
sec_dir_list[0, 0] = [0, 0, 1]
|
||||
sec_pdg_list = np.zeros((N, K_MAX), dtype=np.int64)
|
||||
sec_pdg_list[0, 0] = 2212 # proton
|
||||
sec_valid = np.zeros((N, K_MAX), dtype=bool)
|
||||
sec_valid[0, 0] = True
|
||||
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
sec_cont = encode_secondaries(
|
||||
sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
|
||||
)
|
||||
norm = Normalizer()
|
||||
norm.mean = np.array([-2.0, 0.5], dtype=np.float32)
|
||||
norm.std = np.array([3.0, 1.5], dtype=np.float32)
|
||||
sec_cont_normed = sec_cont.copy()
|
||||
sec_cont_normed[:, :, 4:6] = norm.transform(
|
||||
sec_cont[:, :, 4:6].reshape(-1, 2)
|
||||
).reshape(N, K_MAX, 2)
|
||||
|
||||
n_sec = np.array([1])
|
||||
_, _, sec_mass, sec_charge, _ = decode_secondaries(
|
||||
sec_cont_normed, n_sec, e_sec, pre_dir, sec_phys_normalizer=norm
|
||||
)
|
||||
assert sec_mass[0, 0] == pytest.approx(938.27208943, abs=1e-2)
|
||||
assert sec_charge[0, 0] == pytest.approx(1.0, abs=1e-4)
|
||||
|
||||
Reference in New Issue
Block a user