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>

giant/validate.py

@@ -8,9 +8,10 @@ from giant.sample import (
     sample_ddpm,
     sample_ddim,
     sample_secondaries,
-    snap_type_to_pdg_idx,
 )
 
+_SEC_PHYS_NAMES = ["log_mass", "charge"]
+
 
 def _kw(steps: int | None) -> dict[str, int]:
     return {} if steps is None else {"steps": steps}
@@ -63,12 +64,15 @@ def validate_marginals(
     mode, which always runs the full schedule.
 
     When `sec_decoder` is given, also validates Stage 2: n_sec distribution
-    (+ classification accuracy), secondary species distribution, and
-    per-slot energy-fraction marginals — restricted to each side's own valid
-    slots (real: `n_sec`; generated: the Stage-1 head's argmax), since the
-    two need not agree on how many slots are valid. Adds
-    {"n_sec_real", "n_sec_pred", "n_sec_accuracy", "species_real",
-    "species_generated", "energy_fraction_kl"} to the returned dict.
+    (+ classification accuracy), predicted secondary physical-identity
+    (log_mass, charge) marginals, and per-slot energy-fraction marginals —
+    restricted to each side's own valid slots (real: `n_sec`; generated: the
+    Stage-1 head's argmax), since the two need not agree on how many slots
+    are valid. Compared directly in normalised space (no denormalising —
+    KL estimated from a shared per-sample histogram is invariant to a shared
+    affine rescaling of both sides). Adds {"n_sec_real", "n_sec_pred",
+    "n_sec_accuracy", "phys_real", "phys_generated", "phys_kl",
+    "energy_fraction_kl"} to the returned dict.
     """
     if device is None:
         device = next(model.parameters()).device
@@ -78,15 +82,15 @@ def validate_marginals(
 
     all_real, all_gen = [], []
     all_n_sec_real, all_n_sec_pred = [], []
-    all_species_real, all_species_gen = [], []
+    all_phys_real, all_phys_gen = [], []
     all_frac_real: list[list[np.ndarray]] = [[] for _ in range(K_MAX)]
     all_frac_gen: list[list[np.ndarray]] = [[] for _ in range(K_MAX)]
 
     for i, batch in enumerate(val_loader):
         if n_batches is not None and i >= n_batches:
             break
-        # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx).
-        cond_cont, cond_cat, x1, n_sec, sec_cont, sec_pdg_idx, _proc_idx = batch
+        # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx).
+        cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx = batch
         cond_cont = cond_cont.to(device)
         cond_cat = cond_cat.to(device)
 
@@ -112,9 +116,9 @@ def validate_marginals(
 
         real_valid = np.arange(K_MAX)[None, :] < n_sec_np[:, None]  # (B, K_MAX)
         real_frac = 1.0 / (1.0 + np.exp(-sec_cont[:, :, 0].numpy().astype(np.float64)))
-        real_species = sec_pdg_idx.numpy()
+        real_phys = sec_cont[:, :, 4:6].numpy()  # (B, K_MAX, 2) [log_mass, charge]
 
-        sec_cont_pred, sec_type_emb, sec_valid_pred = sample_secondaries(
+        sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries(
             sec_decoder,
             cond_cont,
             cond_cat,
@@ -122,15 +126,14 @@ def validate_marginals(
             n_sec_pred,
             steps=steps if steps is not None else 10,
         )
-        sec_pdg_pred = snap_type_to_pdg_idx(sec_type_emb, model.pdg_embedding_weight())
         gen_frac = 1.0 / (
             1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64))
         )
-        gen_species = sec_pdg_pred.cpu().numpy()
+        gen_phys = sec_phys_pred.cpu().numpy()
         gen_valid = sec_valid_pred.cpu().numpy()
 
-        all_species_real.append(real_species[real_valid])
-        all_species_gen.append(gen_species[gen_valid])
+        all_phys_real.append(real_phys[real_valid])
+        all_phys_gen.append(gen_phys[gen_valid])
         for j in range(K_MAX):
             all_frac_real[j].append(real_frac[real_valid[:, j], j])
             all_frac_gen[j].append(gen_frac[gen_valid[:, j], j])
@@ -169,14 +172,12 @@ def validate_marginals(
     n_sec_real = np.concatenate(all_n_sec_real, axis=0)
     n_sec_pred_all = np.concatenate(all_n_sec_pred, axis=0)
     n_sec_accuracy = float((n_sec_real == n_sec_pred_all).mean())
-    species_real = np.concatenate(all_species_real, axis=0)
-    species_gen = np.concatenate(all_species_gen, axis=0)
+    phys_real = np.concatenate(all_phys_real, axis=0)  # (M, 2)
+    phys_gen = np.concatenate(all_phys_gen, axis=0)  # (M, 2)
 
-    n_classes = (
-        max(int(species_real.max(initial=0)), int(species_gen.max(initial=0))) + 1
+    phys_kl = np.array(
+        [_histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins) for j in range(2)]
     )
-    species_real_dist = _bincount_frac(species_real, n_classes)
-    species_gen_dist = _bincount_frac(species_gen, n_classes)
 
     energy_fraction_kl = np.full(K_MAX, np.nan)
     print(
@@ -192,10 +193,17 @@ def validate_marginals(
     for v in range(max_n_sec):
         print(f"{v:<20} {real_n_sec_dist[v]:>10.4f} {gen_n_sec_dist[v]:>10.4f}")
 
-    print(f"\n{'pdg model-index':<20} {'real_frac':>10} {'gen_frac':>10}")
-    print("-" * 42)
-    for c in range(n_classes):
-        print(f"{c:<20} {species_real_dist[c]:>10.4f} {species_gen_dist[c]:>10.4f}")
+    print(
+        f"\n{'sec phys (normalised)':<20} {'real_mean':>10} {'gen_mean':>10} "
+        f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
+    )
+    print("-" * 68)
+    for j, name in enumerate(_SEC_PHYS_NAMES):
+        r, g = phys_real[:, j], phys_gen[:, j]
+        print(
+            f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} "
+            f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
+        )
 
     print(
         f"\n{'sec slot (energy frac.)':<24} {'real_mean':>10} {'gen_mean':>10} "
@@ -219,8 +227,9 @@ def validate_marginals(
             "n_sec_real": n_sec_real,
             "n_sec_pred": n_sec_pred_all,
             "n_sec_accuracy": n_sec_accuracy,
-            "species_real": species_real,
-            "species_generated": species_gen,
+            "phys_real": phys_real,
+            "phys_generated": phys_gen,
+            "phys_kl": phys_kl,
             "energy_fraction_kl": energy_fraction_kl,
         }
     )