lars
06c9ad8e5f
- make_seed_frontier only resolves particle mass/charge in "physical" mode, so "embedding"-mode rollouts no longer crash on a seed PDG code giant.particles can't resolve (the TERM_UNKNOWN_PDG gate now handles it). - nearest_known_pdg skips unresolvable candidate PDG codes instead of raising and killing the whole rollout/predict run. - predict/rollout fail with a clear message when a checkpoint predates the sec_phys normalizer, instead of a bare KeyError. - validate_marginals' phys_kl degrades to NaN (matching the energy_fraction_kl pattern) instead of crashing when a validated batch has zero secondaries on either side. - Correct CLAUDE.md's stale claim that the materials table is unfilled. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
114 lines
4.8 KiB
Python
114 lines
4.8 KiB
Python
"""Particle physical-property lookup (mass, charge) for "physical" conditioning.
|
|
|
|
Uses the scikit-HEP `particle` package (PDG data tables) for standard particles
|
|
and ground-state nuclei; falls back to the Z/A decode formula (PDG's 10-digit
|
|
ion scheme `10LZZZAAAI`: Z and A decoded straight from the digits, no lookup
|
|
table involved) for isomer/excited nuclear codes the package's ground-state-only
|
|
nuclide table doesn't cover — confirmed necessary for ~32% of the nuclear codes
|
|
actually present in the multi-material dataset
|
|
(`0932fb02-f2ce-43ca-a4ef-60a2b1221bbc.parquet`).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from functools import lru_cache
|
|
|
|
import numpy as np
|
|
from particle import InvalidParticle, Particle, ParticleNotFound
|
|
from particle import pdgid as _pdgid
|
|
|
|
# First-pass nuclear mass approximation (A * atomic mass unit); no
|
|
# binding-energy correction. Only used for codes missing from `particle`'s
|
|
# ground-state nuclide table -- ground-state codes get the package's real
|
|
# (binding-energy-corrected) mass.
|
|
_AMU_MEV = 931.494
|
|
|
|
# Nearest-neighbour distance weight for `nearest_known_pdg`: charge is a
|
|
# small conserved quantum number and should usually match exactly, so it's
|
|
# weighted far more heavily than the (already log-scaled) mass term.
|
|
_CHARGE_WEIGHT = 50.0
|
|
_LOG_EPS = 1e-8
|
|
|
|
|
|
@lru_cache(maxsize=None)
|
|
def particle_mass_charge(pdg: int) -> tuple[float, float]:
|
|
"""Return (mass_MeV, charge_e) for a raw PDG code.
|
|
|
|
Cached per unique code: the training vocabulary is typically O(100)
|
|
unique codes while a dataset can have O(1e8) rows, and each entry's
|
|
lookup (package query + possible ion decode) is nontrivial enough to be
|
|
worth memoizing rather than repeating per row.
|
|
"""
|
|
pdg = int(pdg)
|
|
try:
|
|
p = Particle.from_pdgid(pdg)
|
|
except (ParticleNotFound, InvalidParticle):
|
|
if _pdgid.is_nucleus(pdg):
|
|
z, a = _pdgid.Z(pdg), _pdgid.A(pdg)
|
|
if z is None or a is None:
|
|
raise ValueError(
|
|
f"PDG {pdg}: is_nucleus but Z/A decode failed"
|
|
) from None
|
|
return float(a) * _AMU_MEV, float(z)
|
|
raise ValueError(
|
|
f"PDG code {pdg} could not be resolved via the `particle` package "
|
|
"and is not a nuclear/ion code (is_nucleus=False) -- no fallback "
|
|
"available; add explicit handling if this is a legitimate code"
|
|
) from None
|
|
# Neutrinos have unmeasured mass in the PDG tables (Particle.mass is
|
|
# None) -- treat as exactly 0, same physical treatment as the photon.
|
|
mass = 0.0 if p.mass is None else float(p.mass)
|
|
charge = 0.0 if p.charge is None else float(p.charge)
|
|
return mass, charge
|
|
|
|
|
|
def particle_phys_array(pdg_codes: np.ndarray) -> np.ndarray:
|
|
"""(N,) int PDG codes -> (N, 2) float32 [mass_MeV, charge_e]."""
|
|
out = np.array(
|
|
[particle_mass_charge(int(p)) for p in np.asarray(pdg_codes)],
|
|
dtype=np.float32,
|
|
)
|
|
return out.reshape(-1, 2)
|
|
|
|
|
|
def nearest_known_pdg(mass: np.ndarray, charge: np.ndarray, candidates) -> np.ndarray:
|
|
"""Reporting-only nearest-PDG label for predicted (mass, charge) pairs.
|
|
|
|
Never used in the inference/training path -- a Stage-2 secondary's
|
|
physical identity is always its raw predicted (mass, charge). This is
|
|
only for populating an output row's nominal "pdg" column and as an
|
|
"embedding" conditioning-mode fallback lookup key for tracks with no
|
|
real PDG code (see giant/rollout.py). Nearest neighbour in
|
|
(log_mass, charge) space over `candidates` (an iterable of PDG codes,
|
|
typically a `pdg_map`'s keys — the training vocabulary), weighting
|
|
charge heavily since it's a small conserved quantum number that should
|
|
usually match exactly.
|
|
"""
|
|
# Resolve each candidate individually and skip ones giant.particles can't
|
|
# resolve, rather than letting one bad code in the training vocabulary
|
|
# crash every rollout/predict run over this reporting-only lookup — an
|
|
# unresolvable code was never a valid label to begin with.
|
|
resolved = []
|
|
for c in sorted({int(c) for c in candidates}):
|
|
try:
|
|
resolved.append((c, *particle_mass_charge(c)))
|
|
except ValueError:
|
|
continue
|
|
if len(resolved) == 0:
|
|
raise ValueError("nearest_known_pdg: no resolvable candidates")
|
|
codes = np.array([r[0] for r in resolved], dtype=np.int64)
|
|
table_log_mass = np.log(
|
|
np.array([r[1] for r in resolved], dtype=np.float64) + _LOG_EPS
|
|
)
|
|
table_charge = np.array([r[2] for r in resolved], dtype=np.float64)
|
|
|
|
mass = np.asarray(mass, dtype=np.float64)
|
|
charge = np.asarray(charge, dtype=np.float64)
|
|
query_log_mass = np.log(np.maximum(mass, 0.0) + _LOG_EPS)
|
|
|
|
d2 = (query_log_mass[:, None] - table_log_mass[None, :]) ** 2 + _CHARGE_WEIGHT * (
|
|
charge[:, None] - table_charge[None, :]
|
|
) ** 2
|
|
idx = d2.argmin(axis=1)
|
|
return codes[idx]
|