lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
faac19edf373f7c43adb6a6da52ecbc8eda35a49
giant/tests/test_phase2.py
T
lars 68fb99bed8 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>
2026-07-17 15:12:54 +02:00

445 lines
16 KiB
Python
Raw Blame History

"""Tests for Phase 2: secondary particle prediction."""
import numpy as np
import pytest
import torch
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
# ── helpers ──────────────────────────────────────────────────────────────────
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, 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):
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.stack(
[torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
)
return cond_cont, cond_cat
# ── DenoisingMLP Phase-2 additions ───────────────────────────────────────────
def test_predict_n_sec_shape():
B = 8
model = _stage1()
cond_cont, cond_cat = _cond(B)
logits = model.predict_n_sec(cond_cont, cond_cat)
assert logits.shape == (B, K_MAX + 1)
def test_predict_n_sec_no_nan():
B = 8
model = _stage1()
cond_cont, cond_cat = _cond(B)
logits = model.predict_n_sec(cond_cont, cond_cat)
assert torch.isfinite(logits).all()
@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 ──────────────────────────────────────────────────────────
@pytest.mark.parametrize("conditioning", ["embedding", "physical"])
def test_sec_decoder_output_shape(conditioning):
B = 8
decoder = _sec_decoder(conditioning=conditioning)
x_t = torch.randn(B, SEC_DIM)
t = torch.rand(B)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
out = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
assert out.shape == (B, SEC_DIM)
def test_sec_decoder_no_nan():
B = 4
decoder = _sec_decoder()
x_t = torch.randn(B, SEC_DIM)
t = torch.rand(B)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
out = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
assert torch.isfinite(out).all()
def test_sec_decoder_gradients():
B = 4
decoder = _sec_decoder()
x_t = torch.randn(B, SEC_DIM)
t = torch.rand(B)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward()
for name, p in decoder.named_parameters():
assert p.grad is not None, f"no grad for {name}"
# ── masked flow matching loss ─────────────────────────────────────────────────
def test_flow_matching_loss_secondary_scalar():
B, pdg, mat = 8, 3, 2
decoder = _sec_decoder(pdg, mat)
x1 = torch.randn(B, SEC_DIM)
cond_cont, cond_cat = _cond(B, pdg, mat)
stage1_out = torch.randn(B, X_DIM)
sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
loss = flow_matching_loss_secondary(
decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
)
assert loss.shape == ()
assert loss.item() >= 0.0
def test_flow_matching_loss_secondary_mask_zeros_padding():
"""Loss with all-zero mask (no valid secondaries) should be 0."""
B, pdg, mat = 4, 3, 2
decoder = _sec_decoder(pdg, mat)
x1 = torch.randn(B, SEC_DIM)
cond_cont, cond_cat = _cond(B, pdg, mat)
stage1_out = torch.randn(B, X_DIM)
sec_mask = torch.zeros(B, K_MAX, dtype=torch.bool)
loss = flow_matching_loss_secondary(
decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
)
assert loss.item() == pytest.approx(0.0, abs=1e-6)
def test_flow_matching_loss_secondary_has_grad():
B, pdg, mat = 4, 3, 2
decoder = _sec_decoder(pdg, mat)
x1 = torch.randn(B, SEC_DIM)
cond_cont, cond_cat = _cond(B, pdg, mat)
stage1_out = torch.randn(B, X_DIM)
sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
flow_matching_loss_secondary(
decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
).backward()
assert any(p.grad is not None for p in decoder.parameters())
# ── sampling ──────────────────────────────────────────────────────────────────
def test_sample_secondaries_shapes():
B, pdg, mat = 6, 3, 2
decoder = _sec_decoder(pdg, mat)
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_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_phys.shape == (B, K_MAX, PARTICLE_PHYS_DIM)
assert sec_valid.shape == (B, K_MAX)
assert sec_valid.dtype == torch.bool
def test_sample_secondaries_valid_mask_matches_n_sec():
B, pdg, mat = 4, 3, 2
decoder = _sec_decoder(pdg, mat)
cond_cont, cond_cat = _cond(B, pdg, mat)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.tensor([0, 1, 3, K_MAX])
_, _, sec_valid = sample_secondaries(
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
)
for i, n in enumerate(n_sec_pred.tolist()):
assert sec_valid[i, :n].all()
assert not sec_valid[i, n:].any()
# ── encode_secondaries round-trip ─────────────────────────────────────────────
def test_encode_secondaries_energy_conservation():
"""Decoded stick-breaking fractions must sum to ≈ e_sec."""
from giant.data.transforms import encode_secondaries
rng = np.random.default_rng(42)
N = 50
n_sec = rng.integers(1, 5, size=N)
e_sec = rng.uniform(0.1, 10.0, size=N).astype(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_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]
energies = rng.dirichlet(np.ones(k)) * e_sec[i]
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, sec_pdg_list=sec_pdg_list
)
assert sec_cont.shape == (N, K_MAX, 6)
assert np.isfinite(sec_cont).all()
def test_encode_secondaries_direction_encoding():
"""Local-frame secondary directions should be unit vectors for valid slots."""
from giant.data.transforms import encode_secondaries
rng = np.random.default_rng(7)
N = 20
e_sec = np.ones(N, dtype=np.float32) * 5.0
sec_E_list = np.zeros((N, K_MAX), dtype=np.float32)
sec_E_list[:, 0] = 3.0
sec_E_list[:, 1] = 2.0
sec_dir_list = rng.standard_normal((N, K_MAX, 3)).astype(np.float32)
norms = np.linalg.norm(sec_dir_list, axis=-1, keepdims=True)
sec_dir_list /= np.where(norms > 0, norms, 1.0)
sec_valid = np.zeros((N, K_MAX), dtype=bool)
sec_valid[:, :2] = True
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)
# dir columns are sec_cont[:, :, 1:4]
local_dirs = sec_cont[:, :2, 1:4] # (N, 2, 3) — valid slots only
norms_out = np.linalg.norm(local_dirs, axis=-1)
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, 6)).astype(np.float32)
sec_cont[:, :, 0] *= stick_logit_scale
dirs = sec_cont[:, :, 1:4]
dirs /= np.linalg.norm(dirs, axis=-1, keepdims=True)
return sec_cont
def test_decode_secondaries_valid_slots_sum_to_e_sec():
"""The valid slots' energies must sum to exactly e_sec, not just <= e_sec.
Rows with n_sec=0 are excluded: there's no slot to put the budget in, so
valid_sum is correctly 0 regardless of e_sec there (see
test_decode_secondaries_zero_n_sec_has_zero_energy) — the shortfall in
that case is handled downstream (e.g. rollout.py dumps it into edep).
"""
from giant.data.transforms import decode_secondaries
rng = np.random.default_rng(0)
N = 200
sec_cont = _random_sec_cont(rng, N)
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, _mass, _charge, sec_valid = decode_secondaries(
sec_cont, n_sec, e_sec, pre_dir
)
valid_sum = (sec_E * sec_valid).sum(axis=1)
has_secondaries = n_sec > 0
np.testing.assert_allclose(
valid_sum[has_secondaries],
e_sec[has_secondaries],
atol=1e-3,
rtol=1e-5,
)
def test_decode_secondaries_zero_n_sec_has_zero_energy():
"""n_sec=0 rows get no secondaries and no forced energy assignment."""
from giant.data.transforms import decode_secondaries
rng = np.random.default_rng(1)
N = 10
sec_cont = _random_sec_cont(rng, N)
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, _mass, _charge, sec_valid = decode_secondaries(
sec_cont, n_sec, e_sec, pre_dir
)
assert not sec_valid.any()
np.testing.assert_allclose(sec_E, 0.0)
def test_decode_secondaries_degenerate_row_falls_back_to_even_split():
"""All-zero stick fractions for the valid slots fall back to an even split."""
from giant.data.transforms import decode_secondaries
rng = np.random.default_rng(2)
N = 4
sec_cont = _random_sec_cont(rng, N)
# Drive every valid slot's stick-breaking fraction to ~0 (huge negative logit).
n_sec = np.array([0, 1, 3, K_MAX])
for i, k in enumerate(n_sec):
sec_cont[i, :k, 0] = -80.0
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, _mass, _charge, sec_valid = decode_secondaries(
sec_cont, n_sec, e_sec, pre_dir
)
for i, k in enumerate(n_sec):
if k == 0:
continue
np.testing.assert_allclose(sec_E[i, :k], e_sec[i] / k, atol=1e-4)
np.testing.assert_allclose(sec_E[i, :k].sum(), e_sec[i], atol=1e-3)
def test_decode_secondaries_rescale_preserves_relative_shares():
"""Rescaling should keep each valid slot's *share* of the budget unchanged.
A shortfall shouldn't get dumped into whichever slot is last by energy
rank — it should be spread proportionally, i.e. sec_E[i] / sec_E[j] for
two valid slots must match before and after the e_sec rescale.
"""
from giant.data.transforms import decode_secondaries
rng = np.random.default_rng(3)
N = 1
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_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, 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 View Git Blame Copy Permalink