giant/giant/particles.py

"""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`).

Also holds the v0.3.0 stage-2 categorical-type rollout decode:
`decode_topn_class`/`decode_embedding_nearest` turn
`Stage2Autoregressive`/`Stage2OneShot`'s `"onehot"`/`"embedding"` type
predictions back into concrete PDG codes, the one place a secondary's
categorical/continuous type representation is ever discretized (its
free-running history representation stays unsnapped — see
`giant/sample.py`'s AR loop).
"""

from __future__ import annotations

from functools import lru_cache
from typing import TYPE_CHECKING

import numpy as np
from particle import InvalidParticle, Particle, ParticleNotFound
from particle import pdgid as _pdgid

if TYPE_CHECKING:
    from giant.data.loader import TopNMap

# 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]


def invert_dense_map(m: dict[int, int]) -> dict[int, int]:
    """index -> key, inverting a dense, bijective value->index map (`pdg_map`,
    or a `TopNMap.class_map`'s non-"other" entries — see `decode_topn_class`,
    which needs a *partial* inverse, not this general one, because its "other"
    index isn't unique-preimage). `pdg_map` itself is always a true bijection
    (`giant.data.loader.build_index_maps_from_files` enumerates the vocab), so
    a plain dict-comprehension inversion is exact here — used for
    `stage2_model.particle_type.target = "embedding"` decode, whose vocabulary
    is the full dense `pdg_map`, not a top-N-plus-other map."""
    return {v: k for k, v in m.items()}


def decode_topn_class(
    class_idx: np.ndarray,
    topn_map: "TopNMap",
    n_classes: int,
    other_policy: str = "sample",
    rng: np.random.Generator | None = None,
) -> np.ndarray:
    """`conditioning.particle.type` / `stage2_model.particle_type.target =
    "onehot"` inference decode: per-row top-N class index -> concrete PDG
    code.

    class_idx: int array, any shape, values in `[0, n_classes)`.
    topn_map: the `TopNMap` (`giant.data.loader.build_pdg_topn_map_from_files`)
        this class index was built from — `class_map` (PDG -> class, injective
        except at the shared "other" index) plus `other_members` (the
        empirical within-"other" distribution, needed for `other_policy =
        "sample"`/`"modal"`).
    n_classes: `conditioning.particle.emb_dim` — the class count; the "other"
        bucket is index `n_classes - 1` by construction
        (`giant.data.loader._topn_plus_other_map`).
    other_policy: `"sample"` draws from `other_members`' empirical frequency;
        `"modal"` always the single most common "other" member; `"drop"`
        returns PDG `0` for those rows (not a valid PDG code — the caller
        must treat it as "no secondary", the same convention as
        `TERM_UNKNOWN_PDG` elsewhere in the rollout driver).

    Every non-"other" class index has a unique inverse (the top `n_classes -
    1` keys each got their own index in `_topn_plus_other_map`), so those
    rows decode exactly; only "other" rows need `other_policy`.
    """
    other_idx = n_classes - 1
    inv = np.zeros(n_classes, dtype=np.int64)
    for pdg, idx in topn_map.class_map.items():
        if idx != other_idx:
            inv[idx] = pdg

    flat = np.asarray(class_idx, dtype=np.int64).reshape(-1)
    out = inv[np.clip(flat, 0, n_classes - 1)]

    other_mask = flat == other_idx
    n_other = int(other_mask.sum())
    if n_other:
        if not topn_map.other_members:
            raise ValueError("decode_topn_class: 'other' class predicted but topn_map.other_members is empty")
        members = np.array(list(topn_map.other_members.keys()), dtype=np.int64)
        counts = np.array(list(topn_map.other_members.values()), dtype=np.float64)
        if other_policy == "drop":
            out[other_mask] = 0
        elif other_policy == "modal":
            out[other_mask] = members[counts.argmax()]
        elif other_policy == "sample":
            rng = rng if rng is not None else np.random.default_rng()
            probs = counts / counts.sum()
            out[other_mask] = rng.choice(members, size=n_other, p=probs)
        else:
            raise ValueError(f"unknown other_policy {other_policy!r}")

    return out.reshape(np.asarray(class_idx).shape)


def decode_embedding_nearest(
    vectors: np.ndarray,
    emb_weight: np.ndarray,
    idx_to_pdg: dict[int, int],
) -> tuple[np.ndarray, np.ndarray]:
    """`stage2_model.particle_type.target = "embedding"` inference decode:
    L1-nearest row of the conditioning's own particle embedding table, since
    a generative model's continuous output
    essentially never lands within float tolerance of a table row (the exact-
    match form is only valid as a round-trip test assertion, never here).

    vectors: `(..., emb_dim)` raw predicted vectors, any leading shape.
    emb_weight: `(vocab, emb_dim)` — `ConditionEncoder.pdg_emb.weight`,
        detached and moved to numpy by the caller. This is the SAME table
        `particle_type.target = "embedding"` was regressed against
        (`validate_config` requires `conditioning.particle.type =
        "embedding"` whenever this target is used — one table, not two).
    idx_to_pdg: `invert_dense_map(pdg_map)` — embedding row index -> PDG.

    Returns `(pdg, l1_dist)`, both shaped like `vectors.shape[:-1]`. `l1_dist`
    is a diagnostic: a heavy tail means the decoder is emitting vectors off
    the embedding manifold, the direct analogue of the species-collapse
    symptom this redesign exists to fix.
    """
    emb_dim = vectors.shape[-1]
    flat = np.asarray(vectors, dtype=np.float64).reshape(-1, emb_dim)
    table = np.asarray(emb_weight, dtype=np.float64)
    d = np.abs(flat[:, None, :] - table[None, :, :]).sum(axis=-1)  # (N, vocab)
    nearest = d.argmin(axis=1)
    dist = d[np.arange(len(nearest)), nearest]
    pdg = np.array([idx_to_pdg[int(i)] for i in nearest], dtype=np.int64)
    lead_shape = vectors.shape[:-1]
    return pdg.reshape(lead_shape), dist.reshape(lead_shape).astype(np.float32)