lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Condition on material/particle physical properties instead of learned embeddings

Browse Source 
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:
Lars Bogner
2026-07-17 15:12:54 +02:00
parent a6bb142a40
commit 68fb99bed8
23 changed files with 1252 additions and 304 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CLAUDE.md
+9 -3
View File
@@ -40,11 +40,15 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
**Stage-1 output space (9D, `giant/constants.py:LOCAL_TARGET_NAMES`):** `log_step_length`, two additive-log-ratio (ALR) coordinates `edep_logit`/`sec_logit` of a **deposit / secondary / post-energy simplex**, `post_dir` (post-scattering momentum direction, unit vector in the local frame), and `travel_dir` (direction of `post_pos - pre_pos`, unit vector in the local frame). The energy simplex decodes via softmax over `[edep_logit, sec_logit, 0]` × `pre_E` so `edep + e_sec + post_E == pre_E` holds by construction — energy conservation is architectural, not learned (see `energy_simplex_decode`). `post_pos` is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, since `step_length` already encodes that displacement's magnitude and duplicating it would let the two become inconsistent.
**Conditioning vector (8D continuous, `COND_DIM`):** pre-step position, log(pre-energy), pre-step direction, layer ID — plus PDG code and material as embeddings. `n_sec` and `e_sec` are **no longer conditioning inputs** (that was Phase 1 / the energy-conservation PoC); the model now predicts them.
**Conditioning vector (15D continuous, `COND_DIM`):** pre-step position, log(pre-energy), pre-step direction, layer ID (`COND_DIM_BASE=8`) — plus, since particle/material physical-property conditioning (`model.conditioning`, see below), 7 more columns: particle `log(mass)`/`charge` (`PARTICLE_PHYS_DIM=2`, `giant/particles.py`) and material `Z_eff`/`A_eff`/`log(density)`/`log(X0)`/`log(λ_int)` (`MATERIAL_PHYS_DIM=5`, `giant/materials.py`). `n_sec` and `e_sec` are **not conditioning inputs** (that was Phase 1 / the energy-conservation PoC); the model predicts them.
`ConditionEncoder`/`SecondaryConditionEncoder` (`giant/model/network.py`) support two mutually exclusive `conditioning` modes, selected per-checkpoint (`model_config["conditioning"]`, defaulting to `"embedding"` for old checkpoints without the key, `"physical"` for new `giant train` runs — see `--conditioning`):
- **`"embedding"`** (original Phase 2 design): a learned `nn.Embedding` per PDG code / material name, indexed by a dataset-scoped dense vocab (`pdg_map`/`mat_map`). Memorizes the training menu.
- **`"physical"`** (default): the 7 physical-property columns above are each routed through a small MLP (`particle_mlp`/`material_mlp`) to the same `emb_dim` width the embedding tables would have produced — a drop-in replacement computable for any PDG code / material name, not just ones seen in training, which is what lets the surrogate generalize to a held-out material or species. `giant/particles.py` decodes nuclear/ion PDG codes (the `10LZZZAAAI` scheme) via the scikit-HEP `particle` package with a Z/A-digit-decode fallback for isomer codes the package's ground-state-only table misses. `giant/materials.py` ships as an intentionally-unfilled stub (`MaterialProperties(None, ...)` per material) that raises loudly (`MaterialPropertiesNotFilledError`) rather than silently defaulting — a physicist must populate real values before `"physical"` mode can train.
**Model** (`giant/model/network.py`): a two-stage model, both checkpointed together.
- **Stage 1 — `DenoisingMLP`:** `ResBlock` stack with a `SinusoidalEmbedding` for the flow/diffusion time variable and a `ConditionEncoder` fusing the conditioning. Predicts the 9D primary vector field, plus an `n_sec_head` classifier over `{0..K_MAX}` (`K_MAX=15`) that runs on the condition encoding alone (no diffusion noise), callable via `predict_n_sec`.
- **Stage 2 — `SecondaryDecoder`:** a second flow-matching net (`SecondaryConditionEncoder` fuses the pre-step conditioning with the Stage-1 outcome) that generates all `K_MAX` secondary slots at once. Each slot is `(stick-breaking energy logit, local-frame direction 3D, continuous type embedding 16D)` = `SEC_SLOT_DIM=20`, ordered by descending energy; slots beyond the predicted `n_sec` are masked. Secondary energies are a **stick-breaking partition of the `e_sec` budget** from Stage 1 (they sum to it), so the whole chain conserves energy. The type embedding is trained against a detached PDG-embedding target (stops self-referential collapse) and snapped to the nearest PDG at inference (`snap_type_to_pdg_idx`).
- **Stage 2 — `SecondaryDecoder`:** a second flow-matching net (`SecondaryConditionEncoder` fuses the pre-step conditioning with the Stage-1 outcome) that generates all `K_MAX` secondary slots at once. Each slot is `(stick-breaking energy logit, local-frame direction 3D, log-mass, charge)` = `SEC_SLOT_DIM=6`, ordered by descending energy; slots beyond the predicted `n_sec` are masked. Secondary energies are a **stick-breaking partition of the `e_sec` budget** from Stage 1 (they sum to it), so the whole chain conserves energy. A secondary's mass/charge are regressed directly against a fixed physics-derived target (its ground-truth PDG code's `giant.particles.particle_mass_charge`) — not a learned/moving embedding target, so nothing needs detaching. **No snapping at inference**: the predicted (mass, charge) are used as-is as the secondary's physical identity, including for its own future conditioning if it goes on to take further steps in a rollout. A separate, reporting-only nearest-known-PDG lookup (`giant.particles.nearest_known_pdg`) is used purely to populate a nominal `pdg` label for output rows / `"embedding"`-mode fallback conditioning — it never feeds back into the model.
`schedule.py` provides both a `CosineSchedule` for DDPM and the flow matching loss utilities (Lipman et al. 2022 conditional flow matching).
@@ -60,4 +64,6 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
**Phase 2 (implemented — baseline):** the two-stage model above jointly predicts `n_sec`, the energy simplex (`e_sec` falls out of it), and each secondary's energy/direction/species, so a rollout is self-contained (no ground-truth secondary counts injected). This is the "get a baseline out" track agreed with Jan & Tobias (2026-07-07).
**Next directions** (parallel, not yet built): faster-eval architectures measured against a ~10× native-Geant4 budget — a Wasserstein-GAN throwaway (single-pass eval) and a mixture-of-experts / routing tree of small nets selected per call (pdg / energy / process), with soft/differentiable gating on continuous routing axes; a sampling-calorimeter (multi-material) dataset; and preferring **material + particle physical properties** over learned embeddings for conditioning. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
**Physical-property conditioning (implemented):** `model.conditioning = "physical" | "embedding"` (see above) replaces the learned PDG/material embeddings with a small MLP over particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, and Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. `"embedding"` stays available as the generalization-comparison baseline. **Not yet done:** `giant/materials.py`'s table needs real physicist-supplied values before `"physical"` mode can train (currently unfilled, fails loudly if used); once filled, the actual held-out-material/species generalization comparison against the `"embedding"` baseline is unrun — the 34GB multi-material dataset at the repo root (6 materials, 237 PDG codes including nuclear/ion codes) is the natural dataset for that experiment.
**Next directions** (parallel, not yet built): faster-eval architectures measured against a ~10× native-Geant4 budget — a Wasserstein-GAN throwaway (single-pass eval) and a mixture-of-experts / routing tree of small nets selected per call (pdg / energy / process), with soft/differentiable gating on continuous routing axes; a sampling-calorimeter (multi-material) dataset. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
giant/cli.py
+38 -15
View File
@@ -42,9 +42,10 @@ from giant.data.transforms import (
)
from giant.geometry import GeometryOracle
from giant.model.network import build_models
from giant.particles import nearest_known_pdg
from giant.pipeline import run_train_job
from giant.rollout import rollout as run_rollout
from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx
from giant.sample import sample_flow, sample_secondaries
app = typer.Typer(no_args_is_help=True)
@@ -178,6 +179,11 @@ class Mode(str, Enum):
ddpm = "ddpm"
class Conditioning(str, Enum):
physical = "physical"
embedding = "embedding"
class Coord(str, Enum):
global_ = "global"
local = "local"
@@ -266,6 +272,15 @@ def train(
"--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"
),
] = None,
conditioning: Annotated[
Optional[Conditioning],
typer.Option(
"--conditioning",
help="Input conditioning: continuous physical properties "
"(mass/charge/Z_eff/A_eff/density/X0/lambda_int, default) or the "
"original learned PDG/material embeddings",
),
] = None,
router: Annotated[
Optional[bool],
typer.Option(
@@ -390,6 +405,7 @@ def train(
"n_blocks": n_blocks,
"emb_dim": emb_dim,
"dropout": dropout,
"conditioning": conditioning.value if conditioning is not None else None,
}.items()
if v is not None
}
@@ -432,6 +448,7 @@ def train(
f"_h{m['hidden_dim']}"
f"_b{m['n_blocks']}"
f"_e{m['emb_dim']}"
f"_c{m['conditioning']}"
f"_lr{t['lr']}"
f"_bs{t['batch_size']}"
)
@@ -573,11 +590,12 @@ def predict(
assert batch_size_value is not None
bs = batch_size_value
conditioning = model_cfg.get("conditioning", "embedding")
pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
pdg_map_inv = {v: k for k, v in pdg_map.items()}
mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"])
model, sec_decoder = build_models(model_cfg)
_load_model_weights(model, sec_decoder, ckpt, weights, checkpoint)
@@ -612,13 +630,13 @@ def predict(
nonlocal writer, total
if coord == Coord.local:
cond_cont, cond_cat, target_raw, _, _, _, _, _, _ = build_features(
piece, pdg_map, mat_map
cond_cont, cond_cat, target_raw, _, _, _, _, _ = build_features(
piece, pdg_map, mat_map, conditioning=conditioning
)
cond_cont = cond_norm.transform(cond_cont)
else:
cond_cont, cond_cat = build_cond_features(
piece, pdg_map, mat_map, cond_norm
piece, pdg_map, mat_map, cond_norm, conditioning=conditioning
)
cc = torch.from_numpy(cond_cont).float().to(_device)
@@ -626,14 +644,10 @@ def predict(
stage1_norm, n_sec_pred = sample_flow(model, cc, ck, steps=steps)
if coord == Coord.global_:
sec_cont, sec_type_emb, _sec_valid_pred = sample_secondaries(
sec_cont, sec_phys, _sec_valid_pred = sample_secondaries(
sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
)
sec_pdg_idx = snap_type_to_pdg_idx(
sec_type_emb, model.pdg_embedding_weight()
)
sec_cont_np = sec_cont.cpu().numpy()
sec_pdg_idx_np = sec_pdg_idx.cpu().numpy()
sec_full_np = torch.cat([sec_cont, sec_phys], dim=-1).cpu().numpy()
n_sec_pred_np = n_sec_pred.cpu().numpy()
pred = stage1_norm.cpu().numpy() # normalised
@@ -692,14 +706,19 @@ def predict(
piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local
)
sec_E, sec_dir_world, sec_pdg_code, _sec_valid = decode_secondaries(
sec_cont_np,
sec_pdg_idx_np,
sec_E, sec_dir_world, sec_mass, sec_charge, _sec_valid = decode_secondaries(
sec_full_np,
n_sec_pred_np,
e_sec_pred,
piece["pre_dir"],
pdg_map_inv,
sec_phys_normalizer=sec_phys_norm,
)
# Reporting-only nearest-known-PDG label (never fed back into the
# model) for the sec_pdg_list output column — see
# giant/particles.py and the "no snapping at inference" design.
sec_pdg_code = nearest_known_pdg(
sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()
).reshape(sec_mass.shape)
sec_pdg_list = [
sec_pdg_code[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)
]
@@ -931,10 +950,12 @@ def rollout(
raise typer.Exit(1)
model_cfg = ckpt["model_config"]
conditioning = model_cfg.get("conditioning", "embedding")
pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"])
model, sec_decoder = build_models(model_cfg)
_load_model_weights(model, sec_decoder, ckpt, weights, checkpoint)
@@ -981,6 +1002,7 @@ def rollout(
seeds,
cond_norm,
tgt_norm,
sec_phys_norm,
pdg_map,
mat_map,
energy_cutoff=energy_cutoff,
@@ -991,6 +1013,7 @@ def rollout(
max_tracks_per_event=max_tracks_per_event,
escape_threshold=escape_threshold,
on_chunk=_write_chunk,
conditioning=conditioning,
)
if writer is not None:
writer.close()
giant/config.py
+7
View File
@@ -33,6 +33,13 @@ DEFAULT_CONFIG: dict = {
"n_blocks": 6,
"emb_dim": 16,
"dropout": 0.1,
# "physical" conditions on material/particle physical properties via
# a small MLP (giant.model.network.ConditionEncoder); "embedding"
# keeps the original learned pdg/material embedding tables — kept
# available as the generalization-comparison baseline. Checkpoints
# from before this option existed have no "conditioning" key and
# load as "embedding" (see giant.model.network.build_models).
"conditioning": "physical",
"router": {
"enabled": False,
"type": "energy", # selects the Router impl from ROUTER_REGISTRY
giant/constants.py
+37 -12
View File
@@ -1,25 +1,50 @@
X_DIM = 9
# Conditioning continuous-feature width (Phase 2): pre_pos(3), log(pre_E)(1),
# pre_dir(3), layer_id(1). n_sec and log(e_sec) are removed — they are now
# *outputs* predicted by Stage 1, not conditioning inputs.
COND_DIM = 8
# Original Phase-2 continuous conditioning: pre_pos(3), log(pre_E)(1),
# pre_dir(3), layer_id(1). This is the slice ConditionEncoder's "embedding"
# mode reads from cond_cont (see giant/model/network.py); n_sec and
# log(e_sec) are not part of it — they are *outputs* predicted by Stage 1,
# not conditioning inputs.
COND_DIM_BASE = 8
# Particle physical-property conditioning: log(mass)(1), charge(1). See
# giant/particles.py.
PARTICLE_PHYS_DIM = 2
# Material physical-property conditioning: Z_eff(1), A_eff(1),
# log(density)(1), log(X0)(1), log(lambda_int)(1). See giant/materials.py.
MATERIAL_PHYS_DIM = 5
# Conditioning continuous-feature width. cond_cont is unconditionally this
# wide regardless of model_config["conditioning"]: "physical" mode computes
# the trailing PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM columns for real,
# "embedding" mode zero-fills them (and never reads them) — see
# giant.data.transforms.build_features/build_cond_features. Bumped 8->15 for
# physical-property conditioning, the same kind of breaking bump as Phase 1
# (10) -> Phase 2 (8).
COND_DIM = COND_DIM_BASE + PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM # 15
# Maximum number of secondary slots. From data: max(n_sec)=14 in PbWO4 dataset;
# K_MAX=15 covers it with one spare slot.
K_MAX = 15
# Per-slot secondary target dimension: 1 (stick-breaking logit) + 3 (local dir) +
# EMB_DIM (continuous type embedding). EMB_DIM must match DenoisingMLP.emb_dim.
# Default emb_dim=16 → SEC_SLOT_DIM=20.
SEC_SLOT_DIM = 20 # 1 + 3 + 16
EMB_DIM = 16 # must match model emb_dim default
# Per-slot continuous width: stick-breaking logit(1) + local dir(3).
CONT_SLOT_DIM = 4
# Per-slot continuous (non-embedding) width: stick-breaking logit + local dir.
CONT_SLOT_DIM = SEC_SLOT_DIM - EMB_DIM # 4
# Per-slot secondary target dimension: CONT_SLOT_DIM (stick-breaking logit +
# local dir) + PARTICLE_PHYS_DIM (log(mass), charge — the secondary's
# predicted physical identity, regressed directly against real physics
# targets rather than a learned/snapped embedding).
SEC_SLOT_DIM = CONT_SLOT_DIM + PARTICLE_PHYS_DIM # 6
# ConditionEncoder's physical-property sub-MLP output width (see
# giant/model/network.py) and "embedding" mode's pdg_emb/mat_emb width.
# Independent of SEC_SLOT_DIM — unlike Phase 2, Stage 2's per-slot physical
# output width is fixed by PARTICLE_PHYS_DIM, not by this.
EMB_DIM = 16
# Flattened Stage-2 target dimension
SEC_DIM = K_MAX * SEC_SLOT_DIM # 15 * 20 = 300
SEC_DIM = K_MAX * SEC_SLOT_DIM # 15 * 6 = 90
# Stage-1 9D target names (unchanged from energy-conservation PoC)
LOCAL_TARGET_NAMES = [
giant/data/dataset.py
+18 -21
View File
@@ -36,16 +36,17 @@ class StreamingStepsDataset(IterableDataset):
numpy slicing instead of a per-row Python loop in the default collate.
Each batch is a tuple:
(cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx)
(cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx)
where:
cond_cont: (B, COND_DIM) float32
cond_cat: (B, 2) int64
target_s1: (B, 9) float32 — normalised Stage-1 primary target
n_sec: (B,) int64 — true secondary count per step
sec_cont: (B, K_MAX, 4) float32 — [stick_logit, local_dir] per slot
sec_pdg_idx: (B, K_MAX) int64 — PDG model-index per secondary slot
proc_idx: (B,) int64 — process-class label (ProcessRouter supervision
only; zeros when `proc_map` is None)
cond_cont: (B, COND_DIM) float32
cond_cat: (B, 2) int64
target_s1: (B, 9) float32 — normalised Stage-1 primary target
n_sec: (B,) int64 — true secondary count per step
sec_cont: (B, K_MAX, SEC_SLOT_DIM) float32 — [stick_logit,
local_dir, log_mass, charge] per slot (mass/charge
normalised iff `sec_phys_normalizer` was given)
proc_idx: (B,) int64 — process-class label (ProcessRouter supervision
only; zeros when `proc_map` is None)
"""
def __init__(
@@ -60,6 +61,8 @@ class StreamingStepsDataset(IterableDataset):
shuffle_buffer: int = 65536,
shuffle: bool = True,
proc_map: dict[str, int] | None = None,
conditioning: str = "embedding",
sec_phys_normalizer: Normalizer | None = None,
) -> None:
self.files = list(files)
self.split_events = split_events
@@ -72,6 +75,8 @@ class StreamingStepsDataset(IterableDataset):
self.shuffle_buffer = max(shuffle_buffer, batch_size)
self.shuffle = shuffle
self.proc_map = proc_map
self.conditioning = conditioning
self.sec_phys_normalizer = sec_phys_normalizer
def __iter__(self):
worker_info = torch.utils.data.get_worker_info()
@@ -88,7 +93,6 @@ class StreamingStepsDataset(IterableDataset):
buf_tgt: list[np.ndarray] = []
buf_nsec: list[np.ndarray] = []
buf_sec: list[np.ndarray] = []
buf_spdg: list[np.ndarray] = []
buf_proc: list[np.ndarray] = []
buf_n = 0
@@ -105,7 +109,6 @@ class StreamingStepsDataset(IterableDataset):
target_s1,
n_sec,
sec_cont,
sec_pdg_idx,
proc_idx,
_,
_,
@@ -115,15 +118,16 @@ class StreamingStepsDataset(IterableDataset):
self.mat_map,
cond_normalizer=self.cond_normalizer,
target_normalizer=self.target_normalizer,
sec_phys_normalizer=self.sec_phys_normalizer,
proc_map=self.proc_map,
require_secondaries=True,
conditioning=self.conditioning,
)
buf_cont.append(cond_cont)
buf_cat.append(cond_cat)
buf_tgt.append(target_s1)
buf_nsec.append(n_sec)
buf_sec.append(sec_cont)
buf_spdg.append(sec_pdg_idx)
buf_proc.append(proc_idx)
buf_n += len(cond_cont)
@@ -134,7 +138,6 @@ class StreamingStepsDataset(IterableDataset):
buf_tgt,
buf_nsec,
buf_sec,
buf_spdg,
buf_proc,
buf_n,
) = yield from self._flush(
@@ -143,7 +146,6 @@ class StreamingStepsDataset(IterableDataset):
buf_tgt,
buf_nsec,
buf_sec,
buf_spdg,
buf_proc,
final=False,
)
@@ -155,7 +157,6 @@ class StreamingStepsDataset(IterableDataset):
buf_tgt,
buf_nsec,
buf_sec,
buf_spdg,
buf_proc,
final=True,
)
@@ -167,7 +168,6 @@ class StreamingStepsDataset(IterableDataset):
buf_tgt: list[np.ndarray],
buf_nsec: list[np.ndarray],
buf_sec: list[np.ndarray],
buf_spdg: list[np.ndarray],
buf_proc: list[np.ndarray],
final: bool,
):
@@ -176,13 +176,12 @@ class StreamingStepsDataset(IterableDataset):
tgt = np.concatenate(buf_tgt)
nsec = np.concatenate(buf_nsec)
sec = np.concatenate(buf_sec)
spdg = np.concatenate(buf_spdg)
proc = np.concatenate(buf_proc)
if self.shuffle:
idx = np.random.permutation(len(cont))
cont, cat, tgt = cont[idx], cat[idx], tgt[idx]
nsec, sec, spdg, proc = nsec[idx], sec[idx], spdg[idx], proc[idx]
nsec, sec, proc = nsec[idx], sec[idx], proc[idx]
bs = self.batch_size
n = len(cont)
@@ -195,12 +194,11 @@ class StreamingStepsDataset(IterableDataset):
torch.from_numpy(tgt[start:end]).float(),
torch.from_numpy(nsec[start:end]).long(),
torch.from_numpy(sec[start:end]).float(),
torch.from_numpy(spdg[start:end]).long(),
torch.from_numpy(proc[start:end]).long(),
)
if final:
return [], [], [], [], [], [], [], 0
return [], [], [], [], [], [], 0
rem = n_full * bs
return (
[cont[rem:]],
@@ -208,7 +206,6 @@ class StreamingStepsDataset(IterableDataset):
[tgt[rem:]],
[nsec[rem:]],
[sec[rem:]],
[spdg[rem:]],
[proc[rem:]],
n - rem,
)
giant/data/transforms.py
+140 -44
View File
@@ -289,20 +289,28 @@ def encode_secondaries(
sec_valid: np.ndarray,
e_sec: np.ndarray,
pre_dir: np.ndarray,
sec_pdg_list: np.ndarray | None = None,
) -> np.ndarray:
"""Encode per-secondary attributes into continuous per-slot targets.
Secondaries must already be sorted descending by energy (as stored in the
parquet). Returns sec_cont of shape (N, K_MAX, 4):
slot[i] = [stick_break_logit, local_dir_x, local_dir_y, local_dir_z]
parquet). Returns sec_cont of shape (N, K_MAX, SEC_SLOT_DIM=6):
slot[i] = [stick_break_logit, local_dir_x, local_dir_y, local_dir_z,
log_mass, charge]
Stick-breaking logit: for slot i, f_i = E_i / remaining_budget, where
remaining_budget = e_sec - sum(E_0..E_{i-1}). The logit is log(f/(1-f)),
clipped to ±_STICK_LOGIT_CLIP. The last valid slot gets +_STICK_LOGIT_CLIP
(takes the full remaining budget). Padding slots get 0.
sec_pdg_idx (integer) is not processed here — kept separate so the loss
function can look up the embedding table at training time.
log_mass/charge are the secondary's real physical identity, looked up
from its ground-truth PDG code (`sec_pdg_list`) via
`giant.particles.particle_phys_array` — a fixed physics-derived
regression target, not a learned/moving one (unlike the embedding-table
target this replaced), so nothing needs to be detached at training time.
`sec_pdg_list` is optional so callers that only need the continuous
stick/dir block (e.g. inference-time re-encoding) can omit it; omitting
it zero-fills the last two columns, matching the padding-slot convention.
"""
N, K = sec_E_list.shape
e_sec = np.asarray(e_sec, dtype=np.float64)
@@ -336,36 +344,66 @@ def encode_secondaries(
pre_dir[valid_mask], sec_dir_list[valid_mask, i]
)
if sec_pdg_list is not None:
from giant.particles import particle_phys_array
# Padding slots carry sentinel pdg 0 (see loader._pad_list_col_int),
# which isn't a resolvable particle — substitute a dummy resolvable
# code (22, photon) there, since the result is discarded below by
# the sec_valid mask regardless.
safe_pdg = np.where(sec_valid, sec_pdg_list, 22)
flat_mass_charge = particle_phys_array(safe_pdg.reshape(-1)) # (N*K, 2)
mass = flat_mass_charge[:, 0].reshape(N, K)
charge = flat_mass_charge[:, 1].reshape(N, K)
log_mass = np.where(sec_valid, log_transform(mass), 0.0).astype(np.float32)
charge = np.where(sec_valid, charge, 0.0).astype(np.float32)
else:
log_mass = np.zeros((N, K), dtype=np.float32)
charge = np.zeros((N, K), dtype=np.float32)
sec_cont = np.concatenate(
[stick_logits[:, :, None], dir_local], axis=-1
) # (N, K, 4)
[stick_logits[:, :, None], dir_local, log_mass[:, :, None], charge[:, :, None]],
axis=-1,
) # (N, K, 6)
return sec_cont.astype(np.float32)
def decode_secondaries(
sec_cont: np.ndarray,
sec_pdg_pred: np.ndarray,
n_sec: np.ndarray,
e_sec: np.ndarray,
pre_dir: np.ndarray,
pdg_map_inv: dict[int, int],
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
sec_phys_normalizer: "Normalizer | None" = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Inverse of encode_secondaries: continuous targets → physical secondary attrs.
sec_cont: (N, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z]
sec_pdg_pred: (N, K_MAX) integer PDG indices (from nearest-neighbor snap)
sec_cont: (N, K_MAX, 6) — [stick_logit, local_dir_x, local_dir_y,
local_dir_z, log_mass, charge] (log_mass/charge normalised iff
`sec_phys_normalizer` was applied when this was produced — e.g. a
raw model prediction; pass the same normalizer here to invert it)
n_sec: (N,) integer secondary counts
e_sec: (N,) total secondary energy budget [MeV]
pre_dir: (N, 3) pre-step world-frame direction
pdg_map_inv: maps model index → PDG code
Returns (sec_E, sec_dir_world, sec_pdg_code, sec_valid) each shape (N, K_MAX).
The valid slots' energies (`sec_E[sec_valid]`, per row) always sum to
exactly `e_sec` — see the rescaling below.
Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid) each
shape (N, K_MAX). The valid slots' energies (`sec_E[sec_valid]`, per row)
always sum to exactly `e_sec` — see the rescaling below. mass/charge are
the model's raw predicted physical identity for each secondary, used
as-is (no snapping to a discrete PDG code) — see giant/particles.py for
the separate, reporting-only nearest-PDG lookup callers may apply on top
of this for display/bookkeeping purposes.
"""
if sec_phys_normalizer is not None:
N_, K_, _ = sec_cont.shape
phys = sec_phys_normalizer.inverse_transform(sec_cont[:, :, 4:6].reshape(-1, 2))
sec_cont = sec_cont.copy()
sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2)
N, K, _ = sec_cont.shape
stick_logits = sec_cont[:, :, 0] # (N, K)
dir_local = sec_cont[:, :, 1:].copy() # (N, K, 3)
dir_local = sec_cont[:, :, 1:4].copy() # (N, K, 3)
log_mass = sec_cont[:, :, 4] # (N, K)
charge = sec_cont[:, :, 5] # (N, K)
# Flow-matching output isn't guaranteed unit norm; normalise before the
# rotation below, which preserves magnitude rather than fixing it up.
@@ -409,15 +447,62 @@ def decode_secondaries(
pre_dir[valid], dir_local[valid, i]
)
sec_pdg_code = np.array(
[
[pdg_map_inv.get(int(sec_pdg_pred[n, i]), 0) for i in range(K)]
for n in range(N)
],
dtype=np.int32,
)
# mass is non-negative by construction (inv_log_transform of a real
# number is always > 0); clip to 0 for padded/invalid slots rather than
# leaving a spurious small positive floor from the log inverse.
sec_mass = np.where(sec_valid, inv_log_transform(log_mass), 0.0).astype(np.float32)
sec_charge = np.where(sec_valid, charge, 0.0).astype(np.float32)
return sec_E, sec_dir_world, sec_pdg_code, sec_valid
return sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid
def _physical_cond_columns(
data: dict[str, np.ndarray], conditioning: str
) -> np.ndarray:
"""(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns.
"embedding" mode zero-fills (cheap, and ConditionEncoder never reads
these columns in that mode — so an unfilled giant.materials table can
never crash an "embedding"-mode run). "physical" mode computes them for
real: particle columns come from `data["mass"]`/`data["charge"]` when the
caller already knows them directly (rollout.py, for a track descended
from a model-predicted secondary — see giant/rollout.py's "no snapping"
design), else derived from `data["pdg"]` via giant.particles; material
columns always come from `data["material"]` via giant.materials, since
material is never itself a model prediction.
"""
from giant.constants import MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
if conditioning == "embedding":
n = len(next(iter(data.values())))
return np.zeros((n, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM), dtype=np.float32)
if conditioning != "physical":
raise ValueError(f"unknown conditioning mode {conditioning!r}")
from giant.materials import material_properties_array
from giant.particles import particle_phys_array
if "mass" in data and "charge" in data:
mass = np.asarray(data["mass"], dtype=np.float32)
charge = np.asarray(data["charge"], dtype=np.float32)
else:
mass, charge = particle_phys_array(data["pdg"]).T
z_eff, a_eff, density, x0, lambda_int = material_properties_array(
data["material"]
).T
return np.column_stack(
[
log_transform(mass),
charge,
z_eff,
a_eff,
log_transform(density),
log_transform(x0),
log_transform(lambda_int),
]
).astype(np.float32)
def build_cond_features(
@@ -425,6 +510,7 @@ def build_cond_features(
pdg_map: dict[int, int],
mat_map: dict[str, int],
cond_normalizer: "Normalizer | None" = None,
conditioning: str = "embedding",
) -> tuple[np.ndarray, np.ndarray]:
"""Build conditioning arrays only — no target, no post-step variables."""
cond_cont = np.column_stack(
@@ -435,6 +521,9 @@ def build_cond_features(
data["layer_id"].astype(np.float32),
]
).astype(np.float32)
cond_cont = np.column_stack(
[cond_cont, _physical_cond_columns(data, conditioning)]
).astype(np.float32)
pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64)
@@ -452,9 +541,11 @@ def build_features(
mat_map: dict[str, int],
cond_normalizer: Normalizer | None = None,
target_normalizer: Normalizer | None = None,
sec_phys_normalizer: Normalizer | None = None,
fit: bool = False,
proc_map: dict[str, int] | None = None,
require_secondaries: bool = False,
conditioning: str = "embedding",
) -> tuple[
np.ndarray,
np.ndarray,
@@ -462,17 +553,17 @@ def build_features(
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
Normalizer | None,
Normalizer | None,
]:
"""Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx) arrays.
"""Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx) arrays.
target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1)
n_sec: (N,) integer secondary counts (target for n_sec head)
sec_cont: (N, K_MAX, 4) continuous secondary targets [stick_logit, dir_local]
sec_pdg_idx: (N, K_MAX) integer PDG model-indices; used to look up embedding
targets in the training loop
sec_cont: (N, K_MAX, SEC_SLOT_DIM=6) continuous secondary targets
[stick_logit, dir_local, log_mass, charge] — mass/charge are
the secondary's real physical identity (from its ground-truth
PDG code), a fixed regression target, not a learned/snapped one.
proc_idx: (N,) integer process-class label (ProcessRouter supervision only —
never conditioning). Zeros when `proc_map` is None or the loaded
data has no "process" column (e.g. pre-conversion parquet files).
@@ -510,7 +601,10 @@ def build_features(
data["pre_dir"],
data["layer_id"].astype(np.float32),
]
).astype(np.float32) # (N, COND_DIM=8)
).astype(np.float32) # (N, COND_DIM_BASE=8)
cond_cont = np.column_stack(
[cond_cont, _physical_cond_columns(data, conditioning)]
).astype(np.float32) # (N, COND_DIM=15)
pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64)
@@ -521,9 +615,9 @@ def build_features(
) # (N,) unclamped, for the valid-slot mask
# Clamp the classification label to K_MAX: the head only has K_MAX+1 classes
# (0..K_MAX), and truncating here mirrors the K_MAX-slot truncation already
# applied to sec_cont/sec_pdg_idx by the loader's list padding. Without this,
# a rare high-multiplicity step (real data goes up to ~37) hands
# cross_entropy an out-of-range target and CUDA asserts.
# applied to sec_cont by the loader's list padding. Without this, a rare
# high-multiplicity step (real data goes up to ~37) hands cross_entropy
# an out-of-range target and CUDA asserts.
n_sec = np.minimum(n_sec_raw, K_MAX).astype(np.int64) # (N,)
# Secondary continuous targets
@@ -534,14 +628,13 @@ def build_features(
if sec_E_list is not None and sec_dir_list is not None and sec_pdg_list is not None:
sec_valid = np.arange(K_MAX)[None, :] < n_sec_raw[:, None] # (N, K_MAX)
sec_cont = encode_secondaries(
sec_E_list, sec_dir_list, sec_valid, data["e_sec"], data["pre_dir"]
) # (N, K_MAX, 4)
# Padding slots carry sentinel pdg 0 (see loader._pad_list_col_int),
# which is never a real PDG code, so `.get(..., 0)` naturally maps
# both real unknown codes and padding to the same masked-out index.
sec_pdg_idx = np.vectorize(lambda p: pdg_map.get(int(p), 0))(
sec_pdg_list
).astype(np.int64)
sec_E_list,
sec_dir_list,
sec_valid,
data["e_sec"],
data["pre_dir"],
sec_pdg_list=sec_pdg_list,
) # (N, K_MAX, 6)
else:
# Guard against silently training Stage 2 on zeroed targets: if any step
# actually spawned secondaries (n_sec > 0, from child_track_ids) but the
@@ -563,8 +656,7 @@ def build_features(
"require_secondaries=False for Stage-1-only use."
)
N = len(n_sec)
sec_cont = np.zeros((N, K_MAX, 4), dtype=np.float32)
sec_pdg_idx = np.zeros((N, K_MAX), dtype=np.int64)
sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
if fit:
cond_normalizer = Normalizer().fit(cond_cont)
@@ -574,6 +666,11 @@ def build_features(
cond_cont = cond_normalizer.transform(cond_cont)
if target_normalizer is not None:
target_s1 = target_normalizer.transform(target_s1)
if sec_phys_normalizer is not None:
N_, K_, _ = sec_cont.shape
phys = sec_phys_normalizer.transform(sec_cont[:, :, 4:6].reshape(-1, 2))
sec_cont = sec_cont.copy()
sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2)
process = data.get("process")
if proc_map is not None and process is not None:
@@ -587,7 +684,6 @@ def build_features(
target_s1,
n_sec,
sec_cont,
sec_pdg_idx,
proc_idx,
cond_normalizer,
target_normalizer,
giant/materials.py
+145
View File
@@ -0,0 +1,145 @@
"""Material physical-property table for the "physical" conditioning mode.
Values are Geant4's own built-in NIST material constants, not hand-typed
literature numbers -- extracted directly from a Geant4 11.4.1 build (the one
vendored in /home/lars/Programming/minicalosim/lib/geant4, built at
minicalosim/build/geant4-install) via a small standalone C++ program linked
against that build (G4NistManager::FindOrBuildMaterial + G4Material::
GetDensity/GetRadlen/GetNuclearInterLength + G4IonisParamMat::GetZeffective).
`a_eff` isn't directly exposed by Geant4, so it's computed with the same
atomic-number-density-weighted-average formula Geant4 itself uses for Zeff
(see G4IonisParamMat::BuildFluctModel in
lib/geant4/source/materials/src/G4IonisParamMat.cc), just applied to A
instead of Z -- for a single-element material this is exact; for a compound
it matches Geant4's own effective-Z convention rather than a different
weighting scheme.
Never silently substitute a default for a material missing from this table
(see UnknownMaterialError/MaterialPropertiesNotFilledError below) -- a wrong
material property would corrupt a whole conditioning axis without any
visible symptom until deep into training.
"""
from __future__ import annotations
from typing import NamedTuple
import numpy as np
class MaterialProperties(NamedTuple):
z_eff: float | None # effective atomic number
a_eff: float | None # effective atomic mass [g/mol]
density: float | None # [g/cm^3]
x0: float | None # radiation length [cm]
lambda_int: float | None # nuclear interaction length [cm]
class UnknownMaterialError(KeyError):
pass
class MaterialPropertiesNotFilledError(NotImplementedError):
pass
# Keys: every NIST material name seen in
# physics/detector-design/minicalosim-geometry.md, plus G4_AIR/G4_lAr which
# appear in dataset parquet files but not that doc. Values from Geant4's
# built-in NIST database (see module docstring) -- all present except
# G4_LYSO, which is not actually a stock Geant4 NIST material (confirmed:
# G4NistManager::FindOrBuildMaterial("G4_LYSO") fails to build in the
# vendored Geant4 11.4.1; it only appears as a plotting-color key in
# minicalosim/bind/G4Calo.py, never constructed in DetectorConstruction.cc)
# -- left unfilled until it's either built as a custom material (e.g.
# Lu1.8Y0.2SiO5:Ce) or dropped from the geometry menu.
MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
"G4_PbWO4": MaterialProperties(
z_eff=31.333333,
a_eff=75.843426,
density=8.28,
x0=0.892453,
lambda_int=20.739740,
),
"G4_CESIUM_IODIDE": MaterialProperties(
z_eff=54.0, a_eff=129.904539, density=4.51, x0=1.860288, lambda_int=39.305990
),
"G4_Pb": MaterialProperties(
z_eff=82.0, a_eff=207.216962, density=11.35, x0=0.561253, lambda_int=18.247950
),
"G4_W": MaterialProperties(
z_eff=74.0, a_eff=183.841648, density=19.30, x0=0.350418, lambda_int=10.311580
),
"G4_Cu": MaterialProperties(
z_eff=29.0, a_eff=63.545648, density=8.96, x0=1.435578, lambda_int=15.587940
),
"G4_Fe": MaterialProperties(
z_eff=26.0, a_eff=55.845113, density=7.874, x0=1.757493, lambda_int=16.990300
),
"G4_BRASS": MaterialProperties(
z_eff=30.939130,
a_eff=68.500857,
density=8.52,
x0=1.367465,
lambda_int=16.947420,
),
"G4_POLYSTYRENE": MaterialProperties(
z_eff=3.5, a_eff=6.509339, density=1.06, x0=41.312510, lambda_int=68.749880
),
"G4_PLASTIC_SC_VINYLTOLUENE": MaterialProperties(
z_eff=3.368421,
a_eff=6.219791,
density=1.032,
x0=42.544200,
lambda_int=69.969390,
),
"G4_BGO": MaterialProperties(
z_eff=27.578947,
a_eff=65.565839,
density=7.13,
x0=1.118030,
lambda_int=22.710130,
),
"G4_LYSO": MaterialProperties(None, None, None, None, None),
"G4_AIR": MaterialProperties(
z_eff=7.261982,
a_eff=14.547593,
density=1.204790e-3,
x0=30392.070000,
lambda_int=71009.500000,
),
"G4_lAr": MaterialProperties(
z_eff=18.0, a_eff=39.947692, density=1.396, x0=14.003440, lambda_int=85.706400
),
}
def get_material_properties(
name: str, table: dict[str, MaterialProperties] | None = None
) -> MaterialProperties:
t = MATERIAL_PROPERTIES if table is None else table
if name not in t:
raise UnknownMaterialError(
f"material {name!r} is not in giant.materials.MATERIAL_PROPERTIES "
f"-- add it (known: {sorted(t)})"
)
props = t[name]
if any(v is None for v in props):
raise MaterialPropertiesNotFilledError(
f"material {name!r} has un-filled physical properties in "
"giant/materials.py -- a physicist must populate real "
"z_eff/a_eff/density/x0/lambda_int values before "
"conditioning='physical' can be used with this material"
)
return props
def material_properties_array(
names: np.ndarray, table: dict[str, MaterialProperties] | None = None
) -> np.ndarray:
"""(N,) str material names -> (N, 5) float32 [z_eff, a_eff, density, x0, lambda_int]."""
out = np.array(
[get_material_properties(str(m), table) for m in np.asarray(names)],
dtype=np.float32,
)
return out.reshape(-1, 5)
giant/model/network.py
+85 -19
View File
@@ -6,7 +6,16 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_DIM, X_DIM
from giant.constants import (
COND_DIM,
COND_DIM_BASE,
EMB_DIM,
K_MAX,
MATERIAL_PHYS_DIM,
PARTICLE_PHYS_DIM,
SEC_DIM,
X_DIM,
)
class SinusoidalEmbedding(nn.Module):
@@ -28,6 +37,23 @@ class SinusoidalEmbedding(nn.Module):
class ConditionEncoder(nn.Module):
"""Fuses continuous conditioning with particle/material identity.
Two mutually exclusive ways to turn (pdg, material) identity into the
two `emb_dim`-wide vectors concatenated with the base continuous
conditioning before the fusion MLP:
- "embedding": a learned `nn.Embedding` lookup table per axis, indexed
by `cond_cat`'s dense training-vocab index. Memorizes the training
menu; the original Phase-2 design.
- "physical": a small MLP per axis, mapping the axis's raw physical
properties (already present in `cond_cont[:, COND_DIM_BASE:]` — see
giant.data.transforms.build_features) to an `emb_dim`-wide vector —
a drop-in replacement for the embedding lookup, computable for any
PDG code / material name rather than only ones seen in training.
Both modes produce the same `in_dim = COND_DIM_BASE + 2*emb_dim` for the
fusion MLP, so only how the two vectors are produced differs.
"""
def __init__(
self,
pdg_vocab: int,
@@ -35,11 +61,27 @@ class ConditionEncoder(nn.Module):
cont_dim: int = COND_DIM,
emb_dim: int = 16,
out_dim: int = 128,
conditioning: str = "embedding",
) -> None:
super().__init__()
self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
in_dim = cont_dim + 2 * emb_dim
if conditioning not in ("embedding", "physical"):
raise ValueError(f"unknown conditioning mode {conditioning!r}")
self.conditioning = conditioning
if conditioning == "embedding":
self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
else:
self.particle_mlp = nn.Sequential(
nn.Linear(PARTICLE_PHYS_DIM, emb_dim),
nn.SiLU(),
nn.Linear(emb_dim, emb_dim),
)
self.material_mlp = nn.Sequential(
nn.Linear(MATERIAL_PHYS_DIM, emb_dim),
nn.SiLU(),
nn.Linear(emb_dim, emb_dim),
)
in_dim = COND_DIM_BASE + 2 * emb_dim
self.mlp = nn.Sequential(
nn.Linear(in_dim, out_dim),
nn.SiLU(),
@@ -47,9 +89,17 @@ class ConditionEncoder(nn.Module):
)
def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
pdg_e = self.pdg_emb(cond_cat[:, 0])
mat_e = self.mat_emb(cond_cat[:, 1])
x = torch.cat([cond_cont, pdg_e, mat_e], dim=-1)
if self.conditioning == "embedding":
pdg_e = self.pdg_emb(cond_cat[:, 0])
mat_e = self.mat_emb(cond_cat[:, 1])
else:
particle_phys = cond_cont[
:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM
]
material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :]
pdg_e = self.particle_mlp(particle_phys)
mat_e = self.material_mlp(material_phys)
x = torch.cat([cond_cont[:, :COND_DIM_BASE], pdg_e, mat_e], dim=-1)
return self.mlp(x)
@@ -91,6 +141,7 @@ class DenoisingMLP(nn.Module):
x_dim: int = X_DIM,
dropout: float = 0.1,
k_max: int = K_MAX,
conditioning: str = "embedding",
) -> None:
super().__init__()
self.time_emb = SinusoidalEmbedding(time_dim)
@@ -99,6 +150,7 @@ class DenoisingMLP(nn.Module):
mat_vocab=mat_vocab,
emb_dim=emb_dim,
out_dim=cond_out_dim,
conditioning=conditioning,
)
merged_cond_dim = time_dim + cond_out_dim
self.input_proj = nn.Linear(x_dim, hidden_dim)
@@ -141,10 +193,6 @@ class DenoisingMLP(nn.Module):
c_emb = self.cond_enc(cond_cont, cond_cat)
return self.n_sec_head(c_emb)
def pdg_embedding_weight(self) -> torch.Tensor:
"""Return the PDG embedding table weights for secondary type targets."""
return self.cond_enc.pdg_emb.weight
class SecondaryConditionEncoder(nn.Module):
"""Encodes pre-step conditioning + Stage-1 output for the secondary decoder."""
@@ -158,6 +206,7 @@ class SecondaryConditionEncoder(nn.Module):
stage1_dim: int = X_DIM,
stage1_proj_dim: int = 64,
out_dim: int = 128,
conditioning: str = "embedding",
) -> None:
super().__init__()
self.base = ConditionEncoder(
@@ -165,6 +214,7 @@ class SecondaryConditionEncoder(nn.Module):
mat_vocab=mat_vocab,
emb_dim=emb_dim,
out_dim=cond_out_dim,
conditioning=conditioning,
)
self.stage1_proj = nn.Linear(stage1_dim, stage1_proj_dim)
fused_dim = cond_out_dim + stage1_proj_dim
@@ -187,8 +237,12 @@ class SecondaryConditionEncoder(nn.Module):
class SecondaryDecoder(nn.Module):
"""Stage-2 model: predicts vector field over K_MAX secondary slots simultaneously.
Each slot encodes (stick_break_logit, local_dir_3D, type_emb) for one
secondary ordered by descending energy. Padded slots are masked from loss.
Each slot encodes (stick_break_logit, local_dir_3D, log_mass, charge) for
one secondary ordered by descending energy — mass/charge are the
secondary's predicted physical identity, regressed directly against real
physics targets (see giant.data.transforms.encode_secondaries), used
as-is with no snapping to a discrete PDG code. Padded slots are masked
from loss.
"""
def __init__(
@@ -203,6 +257,7 @@ class SecondaryDecoder(nn.Module):
stage1_proj_dim: int = 64,
sec_dim: int = SEC_DIM,
dropout: float = 0.1,
conditioning: str = "embedding",
) -> None:
super().__init__()
self.time_emb = SinusoidalEmbedding(time_dim)
@@ -213,6 +268,7 @@ class SecondaryDecoder(nn.Module):
cond_out_dim=cond_out_dim,
stage1_proj_dim=stage1_proj_dim,
out_dim=cond_out_dim,
conditioning=conditioning,
)
merged_cond_dim = time_dim + cond_out_dim
self.input_proj = nn.Linear(sec_dim, hidden_dim)
@@ -587,8 +643,8 @@ class RoutedDenoisingMLP(nn.Module):
Shares the time embedding, `ConditionEncoder`, and `n_sec_head` (all
tiny) across experts and routes only the trunk (where the FLOPs are).
Same `forward`/`predict_n_sec`/`pdg_embedding_weight` signatures as
`DenoisingMLP`, so sample.py/rollout.py/validate.py need no changes.
Same `forward`/`predict_n_sec` signatures as `DenoisingMLP`, so
sample.py/rollout.py/validate.py need no changes.
"""
def __init__(
@@ -604,6 +660,7 @@ class RoutedDenoisingMLP(nn.Module):
x_dim: int = X_DIM,
dropout: float = 0.1,
k_max: int = K_MAX,
conditioning: str = "embedding",
) -> None:
super().__init__()
self.router = router
@@ -613,6 +670,7 @@ class RoutedDenoisingMLP(nn.Module):
mat_vocab=mat_vocab,
emb_dim=emb_dim,
out_dim=cond_out_dim,
conditioning=conditioning,
)
merged_cond_dim = time_dim + cond_out_dim
self.experts = nn.ModuleList(
@@ -656,10 +714,6 @@ class RoutedDenoisingMLP(nn.Module):
c_emb = self.cond_enc(cond_cont, cond_cat)
return self.n_sec_head(c_emb)
def pdg_embedding_weight(self) -> torch.Tensor:
"""Return the PDG embedding table weights for secondary type targets."""
return self.cond_enc.pdg_emb.weight
class RoutedSecondaryDecoder(nn.Module):
"""Routed drop-in for `SecondaryDecoder`.
@@ -682,6 +736,7 @@ class RoutedSecondaryDecoder(nn.Module):
stage1_proj_dim: int = 64,
sec_dim: int = SEC_DIM,
dropout: float = 0.1,
conditioning: str = "embedding",
) -> None:
super().__init__()
self.router = router
@@ -693,6 +748,7 @@ class RoutedSecondaryDecoder(nn.Module):
cond_out_dim=cond_out_dim,
stage1_proj_dim=stage1_proj_dim,
out_dim=cond_out_dim,
conditioning=conditioning,
)
merged_cond_dim = time_dim + cond_out_dim
self.experts = nn.ModuleList(
@@ -732,6 +788,7 @@ _STAGE1_MODEL_KEYS = {
"emb_dim",
"dropout",
"k_max",
"conditioning",
}
_SEC_DECODER_MODEL_KEYS = {
"pdg_vocab",
@@ -740,6 +797,7 @@ _SEC_DECODER_MODEL_KEYS = {
"n_blocks",
"emb_dim",
"dropout",
"conditioning",
}
@@ -798,6 +856,13 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
is truthy; a missing/absent "router" key (pre-routing checkpoints)
falls back to the monolithic pair unchanged, so this is a drop-in
replacement for the ad-hoc constructions it replaces.
`model_config.get("conditioning", "embedding")` — old checkpoints have no
"conditioning" key and must keep loading with their original embedding
tables, so the default here is "embedding", not the training-time
default (which is "physical" — see giant.config.DEFAULT_CONFIG). Read
once and passed to both stage1/sec_decoder, so they structurally always
share one mode.
"""
router_cfg = model_config.get("router")
if router_cfg and router_cfg.get("enabled"):
@@ -810,6 +875,7 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
expert_n_blocks=model_config.get("expert_n_blocks", 3),
emb_dim=model_config.get("emb_dim", EMB_DIM),
dropout=model_config.get("dropout", 0.1),
conditioning=model_config.get("conditioning", "embedding"),
)
stage1 = RoutedDenoisingMLP(
router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab),
giant/model/schedule.py
+12 -9
View File
@@ -81,19 +81,22 @@ def flow_matching_loss_secondary(
) -> torch.Tensor:
"""Flow matching loss for the secondary decoder with per-slot masking.
x1: (B, SEC_DIM) — flattened secondary target (stick_logit, dir, type_emb)
x1: (B, SEC_DIM) — flattened secondary target (stick_logit, dir, log_mass, charge)
sec_mask: (B, K_MAX) bool — True for valid secondary slots
Only valid-slot dimensions contribute to the loss; padded slots are zeroed
before averaging, so the loss is not diluted by empty slots.
Each slot packs CONT_SLOT_DIM continuous dims (stick_logit, dir) followed
by EMB_DIM type-embedding dims. A flat per-dimension mean would let the
16 embedding dims outvote the 4 physically-interesting ones, so the two
blocks are each averaged over their own width first and then combined
with equal weight — this stays correct if EMB_DIM/CONT_SLOT_DIM change.
by PARTICLE_PHYS_DIM physical-identity dims (log_mass, charge) — the
secondary's predicted physical identity, a fixed regression target (see
giant.data.transforms.encode_secondaries). Even though the two blocks are
the same order of magnitude now (unlike the 16-wide learned embedding
block this replaced), they're still on different physical scales, so
they're each averaged over their own width first and then combined with
equal weight — this stays correct if PARTICLE_PHYS_DIM/CONT_SLOT_DIM change.
"""
from giant.constants import CONT_SLOT_DIM, EMB_DIM, K_MAX, SEC_SLOT_DIM
from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
B = x1.size(0)
t = torch.rand(B, device=x1.device)
@@ -104,10 +107,10 @@ def flow_matching_loss_secondary(
err = ((v_t - u_t) ** 2).view(B, K_MAX, SEC_SLOT_DIM)
cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1) # (B, K_MAX)
emb_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + EMB_DIM].mean(dim=-1)
phys_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM].mean(dim=-1)
mask = sec_mask.float()
denom = mask.sum().clamp(min=1)
cont_loss = (cont_err * mask).sum() / denom
emb_loss = (emb_err * mask).sum() / denom
return cont_loss + emb_loss
phys_loss = (phys_err * mask).sum() / denom
return cont_loss + phys_loss
giant/particles.py
+102
View File
@@ -0,0 +1,102 @@
"""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.
"""
codes = np.array(sorted({int(c) for c in candidates}), dtype=np.int64)
if len(codes) == 0:
raise ValueError("nearest_known_pdg: candidates is empty")
table = particle_phys_array(codes) # (C, 2)
table_log_mass = np.log(table[:, 0].astype(np.float64) + _LOG_EPS)
table_charge = table[:, 1].astype(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]
giant/pipeline.py
+32 -15
View File
@@ -5,7 +5,14 @@ import torch
from torch.utils.data import DataLoader
from giant import config
from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_SLOT_DIM, X_DIM
from giant.constants import (
COND_DIM,
EMB_DIM,
K_MAX,
PARTICLE_PHYS_DIM,
SEC_SLOT_DIM,
X_DIM,
)
from giant.data.loader import (
find_parquet_files,
load_event_ids,
@@ -67,27 +74,33 @@ def run_train_job(
)
echo("fitting normalizer (streaming) …")
conditioning = m["conditioning"]
cond_acc = _WelfordAccumulator(COND_DIM)
tgt_acc = _WelfordAccumulator(X_DIM)
sec_phys_acc = _WelfordAccumulator(PARTICLE_PHYS_DIM)
for path in files:
for chunk in iter_file_chunks(path):
mask = np.isin(chunk["event_id"], events_arr)
if not mask.any():
continue
chunk_tr = {k: v[mask] for k, v in chunk.items()}
cond_cont, _, target_s1, _n_sec, _sec_cont, _sec_pdg, _proc, _, _ = (
build_features(
chunk_tr,
pdg_map,
mat_map,
proc_map=proc_map,
require_secondaries=True,
)
cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _ = build_features(
chunk_tr,
pdg_map,
mat_map,
proc_map=proc_map,
require_secondaries=True,
conditioning=conditioning,
)
cond_acc.update(cond_cont)
tgt_acc.update(target_s1)
sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None]
sec_phys = sec_cont[:, :, 4:6][sec_valid]
if len(sec_phys) > 0:
sec_phys_acc.update(sec_phys)
cond_norm = cond_acc.to_normalizer()
tgt_norm = tgt_acc.to_normalizer()
sec_phys_norm = sec_phys_acc.to_normalizer()
train_ds = StreamingStepsDataset(
files=files,
@@ -100,6 +113,8 @@ def run_train_job(
shuffle_buffer=shuffle_buffer,
shuffle=True,
proc_map=proc_map,
conditioning=conditioning,
sec_phys_normalizer=sec_phys_norm,
)
val_ds = StreamingStepsDataset(
files=files,
@@ -111,6 +126,8 @@ def run_train_job(
batch_size=t["batch_size"],
shuffle=False,
proc_map=proc_map,
conditioning=conditioning,
sec_phys_normalizer=sec_phys_norm,
)
pin = device.type == "cuda"
@@ -128,11 +145,6 @@ def run_train_job(
)
emb_dim = m.get("emb_dim", EMB_DIM)
# SEC_SLOT_DIM must match constants (1 stick + 3 dir + emb_dim)
assert SEC_SLOT_DIM == 1 + 3 + emb_dim, (
f"SEC_SLOT_DIM={SEC_SLOT_DIM} must equal 1+3+emb_dim={1 + 3 + emb_dim}; "
"update giant/constants.py if emb_dim changed"
)
model_config = {
"pdg_vocab": len(pdg_map),
@@ -143,6 +155,7 @@ def run_train_job(
"dropout": m["dropout"],
"k_max": K_MAX,
"sec_slot_dim": SEC_SLOT_DIM,
"conditioning": conditioning,
"router": dict(router_cfg),
"expert_hidden_dim": router_cfg["expert_hidden_dim"],
"expert_n_blocks": router_cfg["expert_n_blocks"],
@@ -183,7 +196,11 @@ def run_train_job(
lambda_s2=t.get("lambda_s2", 1.0),
lambda_balance=router_cfg.get("lambda_balance", 0.0),
lambda_proc=router_cfg.get("lambda_proc", 0.0),
normalizer_dict={"cond": cond_norm.to_dict(), "target": tgt_norm.to_dict()},
normalizer_dict={
"cond": cond_norm.to_dict(),
"target": tgt_norm.to_dict(),
"sec_phys": sec_phys_norm.to_dict(),
},
pdg_map={str(k): v for k, v in pdg_map.items()},
mat_map={str(k): v for k, v in mat_map.items()},
proc_map=proc_map,
giant/rollout.py
+52 -16
View File
@@ -38,7 +38,8 @@ from giant.data.transforms import (
inv_log_transform,
reconstruct_post_pos,
)
from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx
from giant.particles import nearest_known_pdg, particle_phys_array
from giant.sample import sample_flow, sample_secondaries
# Record columns produced per step / per terminal marker.
_RECORD_KEYS = [
@@ -82,6 +83,8 @@ def _empty_frontier() -> dict[str, np.ndarray]:
"pre_pos": np.empty((0, 3), dtype=np.float64),
"pre_E": np.empty(0, dtype=np.float64),
"pre_dir": np.empty((0, 3), dtype=np.float64),
"mass": np.empty(0, dtype=np.float64),
"charge": np.empty(0, dtype=np.float64),
}
@@ -215,16 +218,24 @@ def make_seed_frontier(
dir_ = np.asarray(pre_dir, dtype=np.float64)
dir_ = dir_ / np.clip(np.linalg.norm(dir_, axis=1, keepdims=True), 1e-12, None)
pdg_arr = np.asarray(pdg, dtype=np.int64)
# Real primaries always have a genuine ground-truth PDG code, looked up
# once here and carried forward unchanged for the track's lifetime (its
# species never changes mid-track) — same lifecycle as "pdg" itself.
mass, charge = particle_phys_array(pdg_arr).T
frontier = {
"event_id": event_id,
"track_id": track_id,
"parent_id": np.full(n, -1, dtype=np.int64),
"generation": np.zeros(n, dtype=np.int64),
"step_in_track": np.zeros(n, dtype=np.int64),
"pdg": np.asarray(pdg, dtype=np.int64),
"pdg": pdg_arr,
"pre_pos": np.asarray(pre_pos, dtype=np.float64),
"pre_E": np.asarray(pre_E, dtype=np.float64),
"pre_dir": dir_,
"mass": mass.astype(np.float64),
"charge": charge.astype(np.float64),
}
return frontier, counts
@@ -272,6 +283,7 @@ def rollout(
seeds: dict[str, np.ndarray],
cond_norm: Normalizer,
tgt_norm: Normalizer,
sec_phys_norm: Normalizer,
pdg_map: dict[int, int],
mat_map: dict[str, int],
*,
@@ -283,6 +295,7 @@ def rollout(
max_tracks_per_event: int | None = None,
escape_threshold: float | None = None,
on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
conditioning: str = "embedding",
) -> dict[str, np.ndarray] | RolloutSummary:
"""Run showers to completion.
@@ -302,9 +315,6 @@ def rollout(
if escape_threshold is not None:
oracle.escape_threshold = float(escape_threshold)
pdg_map_inv = {v: k for k, v in pdg_map.items()}
pdg_emb_weight = stage1_model.pdg_embedding_weight()
frontier, counts = make_seed_frontier(
seeds["event_id"],
seeds["pdg"],
@@ -327,10 +337,9 @@ def rollout(
oracle,
cond_norm,
tgt_norm,
sec_phys_norm,
pdg_map,
mat_map,
pdg_map_inv,
pdg_emb_weight,
rec,
counts,
energy_cutoff,
@@ -338,6 +347,7 @@ def rollout(
steps,
device,
max_tracks_per_event,
conditioning,
)
)
frontier = _concat_frontiers(next_parts)
@@ -357,10 +367,9 @@ def _step_chunk(
oracle,
cond_norm,
tgt_norm,
sec_phys_norm,
pdg_map,
mat_map,
pdg_map_inv,
pdg_emb_weight,
rec,
counts,
energy_cutoff,
@@ -368,6 +377,7 @@ def _step_chunk(
steps,
device,
max_tracks_per_event,
conditioning,
) -> dict[str, np.ndarray]:
"""Advance one chunk of tracks by a single step; return the next frontier."""
n = len(tr["event_id"])
@@ -421,6 +431,11 @@ def _step_chunk(
layer_id = tr["_layer_id"]
# --- Build conditioning and run the two stages ---
# "mass"/"charge" are the track's own already-resolved physical identity
# (real for a primary, the model's raw predicted values with no snapping
# for a track descended from a secondary — see _spawn_secondaries), used
# directly instead of re-deriving via a pdg lookup. "pdg" still flows
# through for cond_cat's embedding-mode index and the known_pdg gate.
cond_dict = {
"pre_pos": tr["pre_pos"],
"pre_E": tr["pre_E"],
@@ -428,8 +443,12 @@ def _step_chunk(
"layer_id": layer_id,
"material": material,
"pdg": tr["pdg"],
"mass": tr["mass"],
"charge": tr["charge"],
}
cond_cont, cond_cat = build_cond_features(cond_dict, pdg_map, mat_map, cond_norm)
cond_cont, cond_cat = build_cond_features(
cond_dict, pdg_map, mat_map, cond_norm, conditioning=conditioning
)
cc = torch.from_numpy(cond_cont).float().to(device)
ck = torch.from_numpy(cond_cat).long().to(device)
@@ -456,18 +475,24 @@ def _step_chunk(
n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64)
# --- Secondaries ---
sec_cont, sec_type_emb, _valid = sample_secondaries(
# No snapping: sec_mass/sec_charge are the model's raw predicted physical
# identity, used as-is for the spawned track's own future conditioning.
# sec_pdg_code below is a *separate*, reporting-only nearest-known-PDG
# label (never fed back into the model) — see giant/particles.py.
sec_cont, sec_phys, _valid = sample_secondaries(
sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
)
sec_pdg_idx = snap_type_to_pdg_idx(sec_type_emb, pdg_emb_weight)
sec_E, sec_dir_world, sec_pdg_code, sec_valid = decode_secondaries(
sec_cont.cpu().numpy(),
sec_pdg_idx.cpu().numpy(),
sec_full = torch.cat([sec_cont, sec_phys], dim=-1).cpu().numpy()
sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid = decode_secondaries(
sec_full,
n_sec_np,
e_sec,
tr["pre_dir"],
pdg_map_inv,
sec_phys_normalizer=sec_phys_norm,
)
sec_pdg_code = nearest_known_pdg(
sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()
).reshape(sec_mass.shape)
edep = edep.astype(np.float64)
post_E = post_E.astype(np.float64)
@@ -480,6 +505,8 @@ def _step_chunk(
sec_E,
sec_dir_world,
sec_pdg_code,
sec_mass,
sec_charge,
counts,
max_tracks_per_event,
)
@@ -537,6 +564,8 @@ def _step_chunk(
"pre_pos": post_pos[cont],
"pre_E": post_E[cont],
"pre_dir": post_dir_world[cont],
"mass": tr["mass"][cont],
"charge": tr["charge"][cont],
}
return _concat_frontiers([cont_frontier, new_tracks])
@@ -548,6 +577,8 @@ def _spawn_secondaries(
sec_E,
sec_dir_world,
sec_pdg_code,
sec_mass,
sec_charge,
counts,
max_tracks_per_event,
) -> tuple[dict[str, np.ndarray], np.ndarray]:
@@ -594,9 +625,14 @@ def _spawn_secondaries(
"parent_id": tr["track_id"][pr_k],
"generation": tr["generation"][pr_k] + 1,
"step_in_track": np.zeros(len(pr_k), dtype=np.int64),
# Reporting-only nominal PDG (nearest-known-PDG label, never fed back
# into the model) — the track's actual physical identity going
# forward is "mass"/"charge" below, the model's raw prediction.
"pdg": sec_pdg_code[pr_k, sl_k].astype(np.int64),
"pre_pos": post_pos[pr_k],
"pre_E": sec_E[pr_k, sl_k].astype(np.float64),
"pre_dir": sec_dir_world[pr_k, sl_k].astype(np.float64),
"mass": sec_mass[pr_k, sl_k].astype(np.float64),
"charge": sec_charge[pr_k, sl_k].astype(np.float64),
}
return frontier, dropped_edep
giant/sample.py
+11 -21
View File
@@ -43,10 +43,15 @@ def sample_secondaries(
n_sec_pred: (B,) int64 — number of valid secondaries per step
Returns (sec_cont, sec_type_emb, sec_valid):
sec_cont: (B, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z]
sec_type_emb: (B, K_MAX, emb_dim) — predicted type embedding per slot
sec_valid: (B, K_MAX) bool — True for slots i < n_sec_pred
Returns (sec_cont, sec_phys, sec_valid):
sec_cont: (B, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z]
sec_phys: (B, K_MAX, PARTICLE_PHYS_DIM) — predicted [log_mass, charge]
per slot (normalised iff the checkpoint's sec_phys
normalizer was applied at training time — denormalize
before treating as physical units; see
giant.data.transforms.decode_secondaries). Used as-is —
no snapping to a discrete PDG code.
sec_valid: (B, K_MAX) bool — True for slots i < n_sec_pred
"""
sec_decoder.eval()
B = cond_cont.size(0)
@@ -61,27 +66,12 @@ def sample_secondaries(
x_slots = x.view(B, K_MAX, SEC_SLOT_DIM)
sec_cont = x_slots[:, :, :4]
sec_type_emb = x_slots[:, :, 4:]
sec_phys = x_slots[:, :, 4:]
sec_valid = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
1
)
return sec_cont, sec_type_emb, sec_valid
def snap_type_to_pdg_idx(
sec_type_emb: torch.Tensor,
pdg_emb_weight: torch.Tensor,
) -> torch.Tensor:
"""Nearest-neighbour snap: predicted type embedding → PDG model-index.
sec_type_emb: (B, K_MAX, emb_dim)
Returns (B, K_MAX) int64 with model-indices.
"""
B, K, D = sec_type_emb.shape
flat = sec_type_emb.reshape(-1, D)
dists = torch.cdist(flat.float(), pdg_emb_weight.float())
return dists.argmin(dim=-1).reshape(B, K)
return sec_cont, sec_phys, sec_valid
@torch.no_grad()
giant/train.py
+7 -33
View File
@@ -80,31 +80,6 @@ class _GracefulShutdown:
)
def _build_sec_x1(
sec_cont: torch.Tensor,
sec_pdg_idx: torch.Tensor,
pdg_emb_weight: torch.Tensor,
) -> torch.Tensor:
"""Assemble the Stage-2 flow target by appending type embeddings.
sec_cont: (B, K_MAX, 4) — [stick_logit, dir_local]
sec_pdg_idx: (B, K_MAX) — integer PDG model-indices
pdg_emb_weight: (pdg_vocab, emb_dim) — live embedding table weights
Returns (B, SEC_DIM) = (B, K_MAX * (4 + emb_dim)).
Detaches the looked-up rows: this tensor becomes x1 in the flow-matching
loss (u_t = x1 - x0), so without detaching, the Stage-2 loss could pull
the embedding table itself toward whatever the decoder already predicts
(a moving, self-referential regression target) instead of only pulling
the decoder toward the table. The table is still trained normally via
its Stage-1 conditioning role and `predict_n_sec`.
"""
type_emb = pdg_emb_weight[sec_pdg_idx].detach() # (B, K_MAX, emb_dim)
x1_s2 = torch.cat([sec_cont, type_emb], dim=-1) # (B, K_MAX, 4+emb_dim)
return x1_s2.flatten(1) # (B, SEC_DIM)
@torch.no_grad()
def _update_ema(
ema_model: torch.nn.Module, model: torch.nn.Module, decay: float
@@ -128,13 +103,12 @@ def _compute_losses(
torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
]:
"""Compute (total_loss, L_s1, L_nsec, L_s2, L_balance, L_proc) for one batch."""
cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx = batch
cond_cont, cond_cat, x1_s1, n_sec, sec_cont, proc_idx = batch
cond_cont = cond_cont.to(device)
cond_cat = cond_cat.to(device)
x1_s1 = x1_s1.to(device)
n_sec = n_sec.to(device)
sec_cont = sec_cont.to(device)
sec_pdg_idx = sec_pdg_idx.to(device)
proc_idx = proc_idx.to(device)
# Stage-1 flow loss
@@ -150,14 +124,14 @@ def _compute_losses(
# Stage-2 secondary flow loss
# Use a noiseless Stage-1 target as context (detach to avoid back-prop
# coupling between the two flow paths through the same embedding table).
# The type-embedding lookup itself is also detached inside _build_sec_x1,
# so the shared PDG table is shaped only by its Stage-1 conditioning role
# and predict_n_sec, not by chasing the Stage-2 decoder's predictions.
# coupling between the two flow paths). sec_cont's log_mass/charge
# columns are already a fixed physics-derived regression target (see
# giant.data.transforms.encode_secondaries) rather than a learned/moving
# one, so — unlike the embedding-table target this replaced — no
# detaching is needed to keep the target from chasing the decoder.
from giant.constants import K_MAX
pdg_emb_weight = stage1_model.pdg_embedding_weight()
x1_s2 = _build_sec_x1(sec_cont, sec_pdg_idx, pdg_emb_weight)
x1_s2 = sec_cont.flatten(1) # (B, SEC_DIM)
sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
l_s2 = flow_matching_loss_secondary(
giant/validate.py
+37 -28
View File
@@ -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,
}
)
pyproject.toml
+1
View File
@@ -11,6 +11,7 @@ dependencies = [
"tqdm>=4.60,<5",
"typer>=0.12,<1",
"pyyaml>=6,<7",
"particle>=1.0,<2",
]
[project.optional-dependencies]
tests/test_materials.py
+105
View File
@@ -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
+118
View File
@@ -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
+136 -52
View File
@@ -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
+42 -8
View File
@@ -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
+2 -4
View File
@@ -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
+83 -4
View File
@@ -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])
uv.lock
Generated
+33
View File
@@ -54,6 +54,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/d2/39/e7eaf1799466a4aef85b6a4fe7bd175ad2b1c6345066aa33f1f58d4b18d0/asttokens-3.0.1-py3-none-any.whl", hash = "sha256:15a3ebc0f43c2d0a50eeafea25e19046c68398e487b9f1f5b517f7c0f40f976a", size = 27047, upload-time = "2025-11-15T16:43:16.109Z" },
]
[[package]]
name = "attrs"
version = "26.1.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/9a/8e/82a0fe20a541c03148528be8cac2408564a6c9a0cc7e9171802bc1d26985/attrs-26.1.0.tar.gz", hash = "sha256:d03ceb89cb322a8fd706d4fb91940737b6642aa36998fe130a9bc96c985eff32", size = 952055, upload-time = "2026-03-19T14:22:25.026Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" },
]
[[package]]
name = "awkward"
version = "2.9.1"
@@ -439,6 +448,7 @@ source = { editable = "." }
dependencies = [
{ name = "numpy" },
{ name = "pandas" },
{ name = "particle" },
{ name = "pyarrow" },
{ name = "pyyaml" },
{ name = "tqdm" },
@@ -486,6 +496,7 @@ requires-dist = [
{ name = "matplotlib", marker = "extra == 'analysis'", specifier = ">=3.8,<4" },
{ name = "numpy", specifier = ">=1.26,<3" },
{ name = "pandas", specifier = ">=2.2,<4" },
{ name = "particle", specifier = ">=1.0,<2" },
{ name = "polars", marker = "extra == 'analysis'", specifier = ">=1.0,<2" },
{ name = "polars", marker = "extra == 'convert'", specifier = ">=1.0,<2" },
{ name = "pyarrow", specifier = ">=16,<25" },
@@ -502,6 +513,15 @@ requires-dist = [
]
provides-extras = ["cpu", "cuda", "dev", "geometry", "convert", "analysis"]
[[package]]
name = "hepunits"
version = "2.4.6"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e7/1f/c5f7525bf7e5d14d61750f91845f9f9350a8f33f5ac95672cec92f317038/hepunits-2.4.6.tar.gz", hash = "sha256:bca6ada937147166d66e9fa152566f2378868d798211cad6990080399560fa34", size = 17985, upload-time = "2026-06-16T09:23:36.685Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/85/10/7f9c58d1ec6a0b7f7783fe552f3593f39cda30c2e1d7a9d148ae711e748d/hepunits-2.4.6-py3-none-any.whl", hash = "sha256:089c52c3b84ef67a159b5e9ee9bdd50e1a442e3fd0c101303cc409c1e9011c4d", size = 17090, upload-time = "2026-06-16T09:23:35.35Z" },
]
[[package]]
name = "iniconfig"
version = "2.3.0"
@@ -1184,6 +1204,19 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/99/5d/8268b644392ee874ee82a635cd0df1773de230bde356c38de28e298392cc/parso-0.8.7-py2.py3-none-any.whl", hash = "sha256:a8926eb2a1b915486941fdbd31e86a4baf88fe8c210f25f2f35ecec5b574ca1c", size = 107025, upload-time = "2026-05-01T23:12:58.867Z" },
]
[[package]]
name = "particle"
version = "1.0.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "attrs" },
{ name = "hepunits" },
]
sdist = { url = "https://files.pythonhosted.org/packages/c3/66/09911bbb658fdffe960903c12edecab95f6cced40fef4909d1cc04bd288b/particle-1.0.0.tar.gz", hash = "sha256:49145dec1cb5044b07f3e8e902280fa050950fa845b058003e69de519bb50492", size = 285766, upload-time = "2026-06-25T14:48:26.893Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/a8/92/05078b696cddbdd60963577c895d5b77e5bf829197b08677351b11761c2b/particle-1.0.0-py3-none-any.whl", hash = "sha256:fc2656f53e729be76e45430f56aa65dc20dea069565a393032016544425b64bf", size = 245760, upload-time = "2026-06-25T14:48:25.238Z" },
]
[[package]]
name = "pexpect"
version = "4.9.0"