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>

tests/test_materials.py

@@ -0,0 +1,105 @@
+import numpy as np
+import pytest
+
+from giant.materials import (
+    MaterialProperties,
+    MaterialPropertiesNotFilledError,
+    UnknownMaterialError,
+    get_material_properties,
+    material_properties_array,
+)
+
+
+def test_get_material_properties_unknown_name_raises():
+    with pytest.raises(UnknownMaterialError):
+        get_material_properties("G4_Unobtainium")
+
+
+def test_get_material_properties_unfilled_entry_raises():
+    """G4_LYSO is not a stock Geant4 NIST material (confirmed against the
+    vendored Geant4 11.4.1 build) and is the one entry still shipped unfilled."""
+    with pytest.raises(MaterialPropertiesNotFilledError):
+        get_material_properties("G4_LYSO")
+
+
+def test_get_material_properties_returns_filled_entry_from_injected_table():
+    table = {
+        "G4_Pb": MaterialProperties(
+            z_eff=82.0, a_eff=207.2, density=11.35, x0=0.5612, lambda_int=17.59
+        )
+    }
+    props = get_material_properties("G4_Pb", table)
+    assert props.z_eff == 82.0
+    assert props.a_eff == 207.2
+    assert props.density == 11.35
+    assert props.x0 == 0.5612
+    assert props.lambda_int == 17.59
+
+
+def test_material_properties_array_shape_and_values():
+    table = {
+        "G4_Pb": MaterialProperties(82.0, 207.2, 11.35, 0.5612, 17.59),
+        "G4_W": MaterialProperties(74.0, 183.84, 19.3, 0.3504, 9.95),
+    }
+    names = np.array(["G4_Pb", "G4_W", "G4_Pb"], dtype=object)
+    arr = material_properties_array(names, table)
+    assert arr.shape == (3, 5)
+    assert arr.dtype == np.float32
+    np.testing.assert_allclose(arr[0], [82.0, 207.2, 11.35, 0.5612, 17.59], rtol=1e-5)
+    np.testing.assert_allclose(arr[1], [74.0, 183.84, 19.3, 0.3504, 9.95], rtol=1e-5)
+
+
+def test_all_known_materials_present_in_stub_table():
+    """Every material referenced elsewhere in the repo must at least have a
+    stub entry (even if unfilled) -- an unknown name should never be the
+    failure mode a physicist hits when populating the table."""
+    from giant.materials import MATERIAL_PROPERTIES
+
+    expected = {
+        "G4_PbWO4",
+        "G4_CESIUM_IODIDE",
+        "G4_Pb",
+        "G4_W",
+        "G4_Cu",
+        "G4_Fe",
+        "G4_BRASS",
+        "G4_POLYSTYRENE",
+        "G4_PLASTIC_SC_VINYLTOLUENE",
+        "G4_BGO",
+        "G4_LYSO",
+        "G4_AIR",
+        "G4_lAr",
+    }
+    assert expected <= set(MATERIAL_PROPERTIES.keys())
+
+
+def test_all_materials_filled_except_lyso():
+    """G4_LYSO is the sole intentionally-unfilled entry (not a stock Geant4
+    NIST material); every other known material has real Geant4-derived
+    values -- see the module docstring for provenance."""
+    from giant.materials import MATERIAL_PROPERTIES
+
+    for name, props in MATERIAL_PROPERTIES.items():
+        if name == "G4_LYSO":
+            assert all(v is None for v in props)
+        else:
+            assert all(v is not None for v in props), f"{name} unexpectedly unfilled"
+
+
+def test_elemental_material_z_eff_matches_atomic_number():
+    """Single-element materials' z_eff must equal the element's real Z."""
+    pb = get_material_properties("G4_Pb")
+    assert pb.z_eff == pytest.approx(82.0)
+    w = get_material_properties("G4_W")
+    assert w.z_eff == pytest.approx(74.0)
+    fe = get_material_properties("G4_Fe")
+    assert fe.z_eff == pytest.approx(26.0)
+
+
+def test_pbwo4_values_match_known_cms_ecal_reference():
+    """PbWO4 (CMS ECAL crystal) has well-known reference values: X0~0.89cm,
+    density 8.28 g/cm^3 -- sanity check the Geant4-derived numbers land there."""
+    pbwo4 = get_material_properties("G4_PbWO4")
+    assert pbwo4.density == pytest.approx(8.28)
+    assert pbwo4.x0 == pytest.approx(0.89, abs=0.01)
+    assert pbwo4.z_eff == pytest.approx(31.33, abs=0.01)

tests/test_particles.py

@@ -0,0 +1,118 @@
+import numpy as np
+import pytest
+
+from giant.particles import (
+    nearest_known_pdg,
+    particle_mass_charge,
+    particle_phys_array,
+)
+
+
+def test_photon_massless_neutral():
+    mass, charge = particle_mass_charge(22)
+    assert mass == pytest.approx(0.0)
+    assert charge == pytest.approx(0.0)
+
+
+def test_electron_mass_charge():
+    mass, charge = particle_mass_charge(11)
+    assert mass == pytest.approx(0.51099895069, rel=1e-6)
+    assert charge == pytest.approx(-1.0)
+
+
+def test_positron_is_charge_conjugate_of_electron():
+    mass_e, charge_e = particle_mass_charge(11)
+    mass_p, charge_p = particle_mass_charge(-11)
+    assert mass_p == pytest.approx(mass_e)
+    assert charge_p == pytest.approx(-charge_e)
+
+
+def test_proton_mass_charge():
+    mass, charge = particle_mass_charge(2212)
+    assert mass == pytest.approx(938.27208943, rel=1e-6)
+    assert charge == pytest.approx(1.0)
+
+
+def test_neutrino_unmeasured_mass_treated_as_zero():
+    """PDG tables store an unmeasured neutrino mass as None -- must not
+    propagate a None/NaN into a physical conditioning feature."""
+    mass, charge = particle_mass_charge(12)
+    assert mass == pytest.approx(0.0)
+    assert charge == pytest.approx(0.0)
+
+
+def test_ground_state_nucleus_resolved_via_particle_package():
+    """He-4 (Z=2, A=4) is a common nuclide in `particle`'s ground-state table."""
+    mass, charge = particle_mass_charge(1000020040)
+    assert charge == pytest.approx(2.0)
+    assert mass == pytest.approx(
+        4 * 931.494, rel=0.05
+    )  # near A*amu, binding-energy-corrected
+
+
+def test_nuclear_isomer_falls_back_to_z_a_decode():
+    """An excited/isomer nuclear code (nonzero trailing digit) is absent from
+    `particle`'s ground-state-only nuclide table -- confirmed necessary for
+    ~32% of the nuclear codes in the multi-material dataset. Fe-56 isomer:
+    Z=26, A=56, isomer level 1 -> pdgid 1000260561."""
+    pdg = 1000260561
+    mass, charge = particle_mass_charge(pdg)
+    assert charge == pytest.approx(26.0)
+    assert mass == pytest.approx(56 * 931.494, rel=1e-6)
+
+
+def test_invalid_pdg_code_raises():
+    with pytest.raises(ValueError):
+        particle_mass_charge(999999999)
+
+
+def test_particle_mass_charge_is_cached():
+    particle_mass_charge.cache_clear()
+    particle_mass_charge(22)
+    particle_mass_charge(22)
+    info = particle_mass_charge.cache_info()
+    assert info.hits >= 1
+
+
+def test_particle_phys_array_shape_and_dtype():
+    arr = particle_phys_array(np.array([22, 11, 2212]))
+    assert arr.shape == (3, 2)
+    assert arr.dtype == np.float32
+    np.testing.assert_allclose(arr[0], [0.0, 0.0])
+    np.testing.assert_allclose(arr[2], [938.27208943, 1.0], rtol=1e-5)
+
+
+# ── nearest_known_pdg (reporting-only nearest-neighbour label) ──────────────
+
+
+def test_nearest_known_pdg_exact_match():
+    candidates = [22, 11, -11, 2212, 2112]
+    mass_e, charge_e = particle_mass_charge(11)
+    result = nearest_known_pdg(np.array([mass_e]), np.array([charge_e]), candidates)
+    assert result[0] == 11
+
+
+def test_nearest_known_pdg_prioritises_charge_match():
+    """Charge is a small conserved quantum number and should usually match
+    exactly even when the queried mass is noisy/imperfect."""
+    candidates = [22, 11, -11, 2212]
+    # Close to electron mass but not exact, positive charge like the positron.
+    result = nearest_known_pdg(np.array([0.6]), np.array([1.0]), candidates)
+    assert result[0] == -11
+
+
+def test_nearest_known_pdg_empty_candidates_raises():
+    with pytest.raises(ValueError):
+        nearest_known_pdg(np.array([1.0]), np.array([0.0]), [])
+
+
+def test_nearest_known_pdg_shape():
+    candidates = [22, 11, -11, 2212, 2112]
+    n = 10
+    result = nearest_known_pdg(
+        np.random.default_rng(0).uniform(0, 1000, n),
+        np.random.default_rng(1).uniform(-1, 1, n),
+        candidates,
+    )
+    assert result.shape == (n,)
+    assert set(result.tolist()) <= set(candidates)

tests/test_phase2.py

@@ -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)

tests/test_rollout.py

@@ -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),

tests/test_router.py

@@ -576,11 +576,9 @@ def test_routed_denoising_mlp_predict_n_sec_shape():
     assert logits.shape == (B, K_MAX + 1)
 
 
-def test_routed_denoising_mlp_pdg_embedding_weight_shape():
+def test_routed_denoising_mlp_has_no_pdg_embedding_weight_method():
     model = _routed_stage1(pdg=5, mat=2)
-    from giant.constants import EMB_DIM
-
-    assert model.pdg_embedding_weight().shape == (5, EMB_DIM)
+    assert not hasattr(model, "pdg_embedding_weight")
 
 
 # ── RoutedSecondaryDecoder ───────────────────────────────────────────────────

tests/test_transforms.py

@@ -1,7 +1,8 @@
 import numpy as np
 import pytest
-from giant.constants import K_MAX
+from giant.constants import COND_DIM, COND_DIM_BASE, K_MAX
 from giant.data.transforms import (
+    build_cond_features,
     build_features,
     energy_simplex_decode,
     energy_simplex_encode,
@@ -235,7 +236,7 @@ def test_build_features_clamps_n_sec_label_to_k_max():
     pdg_map = {11: 0}
     mat_map = {"PbWO4": 0}
 
-    _, _, _, n_sec, _, _, _, _, _ = build_features(data, pdg_map, mat_map)
+    _, _, _, n_sec, _, _, _, _ = build_features(data, pdg_map, mat_map)
 
     assert n_sec.max() <= K_MAX
     np.testing.assert_array_equal(n_sec, [0, 5, K_MAX])
@@ -312,9 +313,87 @@ def test_build_features_require_secondaries_ok_when_no_secondaries():
     data = _step_data_no_sec_lists(np.zeros(3, dtype=np.int32))
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
-    _, _, _, _, sec_cont, sec_pdg_idx, *_ = build_features(
+    _, _, _, _, sec_cont, *_ = build_features(
         data, pdg_map, mat_map, require_secondaries=True
     )
 
     assert not sec_cont.any()
-    assert not sec_pdg_idx.any()
+
+
+# ── physical-property conditioning ────────────────────────────────────────────
+
+
+@pytest.fixture
+def fake_material_props(monkeypatch):
+    """Inject a fully-populated fake materials table for "physical" mode
+    tests, independent of when the real giant/materials.py table is filled
+    in by the user (see giant.materials.MaterialPropertiesNotFilledError)."""
+    import giant.materials as gm
+
+    fake = {
+        "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_build_features_embedding_mode_zero_fills_physical_columns():
+    data = _minimal_step_data(3)
+    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
+
+    cond_cont, *_ = build_features(data, pdg_map, mat_map, conditioning="embedding")
+
+    assert cond_cont.shape[1] == COND_DIM
+    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0)
+
+
+def test_build_features_physical_mode_shape_and_values(fake_material_props):
+    from giant.particles import particle_mass_charge
+
+    data = _minimal_step_data(3)
+    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
+
+    cond_cont, *_ = build_features(data, pdg_map, mat_map, conditioning="physical")
+
+    assert cond_cont.shape[1] == COND_DIM
+    mass, charge = particle_mass_charge(11)
+    expected_log_mass = log_transform(np.array([mass]))[0]
+    np.testing.assert_allclose(
+        cond_cont[:, COND_DIM_BASE], expected_log_mass, atol=1e-5
+    )
+    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], charge)
+    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 2], 75.6)  # z_eff
+    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 3], 205.3)  # a_eff
+
+
+def test_build_features_physical_mode_unfilled_material_raises():
+    """G4_LYSO is the one material giant/materials.py still ships unfilled
+    (not a stock Geant4 NIST material) — must fail loudly, not silently."""
+    from giant.materials import MaterialPropertiesNotFilledError
+
+    data = _minimal_step_data(2)
+    data["material"] = np.full(2, "G4_LYSO", dtype=object)
+    pdg_map, mat_map = {11: 0}, {"G4_LYSO": 0}
+
+    with pytest.raises(MaterialPropertiesNotFilledError):
+        build_features(data, pdg_map, mat_map, conditioning="physical")
+
+
+def test_build_cond_features_mass_charge_override(fake_material_props):
+    """rollout.py's secondaries carry their own predicted mass/charge — when
+    present in `data`, these bypass the pdg-based lookup entirely (the "no
+    snapping" design: a track's own future conditioning must use its actual
+    predicted physical identity, not a value re-derived from a PDG code)."""
+    data = _minimal_step_data(2)
+    data["mass"] = np.array([123.0, 456.0], dtype=np.float32)
+    data["charge"] = np.array([2.0, -2.0], dtype=np.float32)
+    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
+
+    cond_cont, _ = build_cond_features(data, pdg_map, mat_map, conditioning="physical")
+
+    np.testing.assert_allclose(
+        cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0]))
+    )
+    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], [2.0, -2.0])