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:
+42
-8
@@ -20,17 +20,22 @@ PDG_MAP = {22: 0, 11: 1, -11: 2}
|
||||
MAT_MAP = {"G4_AIR": 0, "G4_PbWO4": 1}
|
||||
|
||||
|
||||
def _models():
|
||||
s1 = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
|
||||
s2 = SecondaryDecoder(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
|
||||
def _models(conditioning="embedding"):
|
||||
s1 = DenoisingMLP(
|
||||
pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, conditioning=conditioning
|
||||
)
|
||||
s2 = SecondaryDecoder(
|
||||
pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, conditioning=conditioning
|
||||
)
|
||||
return s1.eval(), s2.eval()
|
||||
|
||||
|
||||
def _norms():
|
||||
rng = np.random.default_rng(0)
|
||||
cond = Normalizer().fit(rng.standard_normal((1000, 8)).astype(np.float32))
|
||||
cond = Normalizer().fit(rng.standard_normal((1000, 15)).astype(np.float32))
|
||||
tgt = Normalizer().fit(rng.standard_normal((1000, 9)).astype(np.float32))
|
||||
return cond, tgt
|
||||
sec_phys = Normalizer().fit(rng.standard_normal((1000, 2)).astype(np.float32))
|
||||
return cond, tgt, sec_phys
|
||||
|
||||
|
||||
def _oracle():
|
||||
@@ -63,11 +68,12 @@ def _run(
|
||||
max_steps=30,
|
||||
max_tracks_per_event=300,
|
||||
seeds=None,
|
||||
conditioning="embedding",
|
||||
):
|
||||
torch.manual_seed(0)
|
||||
np.random.seed(0)
|
||||
s1, s2 = _models()
|
||||
cond, tgt = _norms()
|
||||
s1, s2 = _models(conditioning)
|
||||
cond, tgt, sec_phys = _norms()
|
||||
return rollout(
|
||||
s1,
|
||||
s2,
|
||||
@@ -75,6 +81,7 @@ def _run(
|
||||
seeds or _seeds(),
|
||||
cond,
|
||||
tgt,
|
||||
sec_phys,
|
||||
PDG_MAP,
|
||||
MAT_MAP,
|
||||
energy_cutoff=energy_cutoff,
|
||||
@@ -83,9 +90,35 @@ def _run(
|
||||
batch_size=128,
|
||||
max_tracks_per_event=max_tracks_per_event,
|
||||
escape_threshold=escape_threshold,
|
||||
conditioning=conditioning,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_material_props(monkeypatch):
|
||||
import giant.materials as gm
|
||||
|
||||
fake = {
|
||||
"G4_AIR": gm.MaterialProperties(
|
||||
z_eff=7.3, a_eff=14.4, density=1.2e-3, x0=3.0e4, lambda_int=7.0e5
|
||||
),
|
||||
"G4_PbWO4": gm.MaterialProperties(
|
||||
z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7
|
||||
),
|
||||
}
|
||||
monkeypatch.setattr(gm, "MATERIAL_PROPERTIES", fake)
|
||||
return fake
|
||||
|
||||
|
||||
def test_rollout_physical_conditioning_end_to_end(fake_material_props):
|
||||
"""Physical-mode rollout runs to completion; spawned secondaries carry
|
||||
mass/charge forward (no snapping) and the output pdg column is populated
|
||||
via the reporting-only nearest-known-PDG label."""
|
||||
rec = _run(conditioning="physical")
|
||||
assert len(rec["event_id"]) > 0
|
||||
assert set(np.unique(rec["pdg"]).tolist()) <= set(PDG_MAP.keys())
|
||||
|
||||
|
||||
def test_seed_frontier_track_ids():
|
||||
seeds = _seeds(3)
|
||||
fr, counts = make_seed_frontier(**seeds)
|
||||
@@ -171,7 +204,7 @@ def _run_streaming(on_chunk, **kwargs):
|
||||
torch.manual_seed(0)
|
||||
np.random.seed(0)
|
||||
s1, s2 = _models()
|
||||
cond, tgt = _norms()
|
||||
cond, tgt, sec_phys = _norms()
|
||||
seeds = kwargs.pop("seeds", None) or _seeds()
|
||||
return rollout(
|
||||
s1,
|
||||
@@ -180,6 +213,7 @@ def _run_streaming(on_chunk, **kwargs):
|
||||
seeds,
|
||||
cond,
|
||||
tgt,
|
||||
sec_phys,
|
||||
PDG_MAP,
|
||||
MAT_MAP,
|
||||
energy_cutoff=kwargs.pop("energy_cutoff", 1.0),
|
||||
|
||||
Reference in New Issue
Block a user