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

Implement Phase 2: secondary particle prediction #9

Merged
lbogner merged 21 commits from phase2-secondary-prediction into master 2026-07-17 11:04:22 +02:00
Conversation 1 Commits 21 Files Changed 40
+5597 -363
40 changed files with 5597 additions and 363 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CLAUDE.md
+20 -8
View File
@@ -8,12 +8,16 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
uv sync --extra cpu # install dependencies with CPU-only torch (standard/default)
uv sync --extra cuda # install dependencies with CUDA 11.8 torch
uv sync --extra cpu --extra dev # add dev extras (pytest, etc.)
uv sync --extra cpu --extra geometry # add scikit-learn for the geometry oracle (giant rollout)
pytest # run tests
giant train path/to/steps.parquet --mode flow # train (flow matching)
giant train path/to/steps.parquet --mode ddpm # train (DDPM baseline)
giant predict path/to/steps.parquet --checkpoint ckpt/best.pt # per-step predictions
giant rollout path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl # full showers
dwarf --help # dataset/tooling CLI: convert, migrate, bump-gen,
# bump-schema, status, update-manifest, create-manifest,
# make-root, hparam-scan (see scripts/dwarf.py)
# make-root, build-geometry-oracle, hparam-scan
# (see scripts/dwarf.py)
```
`cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x; newer torch requires newer NVIDIA drivers). Plain `uv sync` with no extra will not install torch at all; uv has no concept of a "default extra", so `--extra cpu` should always be included unless you need GPU support.
@@ -30,22 +34,30 @@ Part of the `dev` extra. Run these periodically (not just at commit time) to cat
## Architecture
GIANT is a conditional generative surrogate for the Geant4 step function. It replaces the stochastic physics engine: given a pre-step particle state (conditioning), it samples a post-step outcome.
GIANT is a conditional generative surrogate for the Geant4 step function. It replaces the stochastic physics engine: given a pre-step particle state (conditioning), it samples a post-step outcome — now including the variable-length list of secondary particles the step produces (Phase 2, see Roadmap).
**Data pipeline** (`giant/data/`): parquet files from miniCaloSim are loaded into numpy arrays (`loader.py`), then log-transformed and rotated into a local coordinate frame where `pre_dir = ẑ` (`transforms.py`), before being wrapped in a PyTorch `Dataset` (`dataset.py`). Train/val split is by `event_id` to avoid leaking correlated steps from the same shower.
**Output space (9D):** `step_length` (log), `ΔE` (log), `edep` (log), `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). `post_pos` itself 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.
**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:** PDG code (embedding), pre-step position, log(pre-energy), pre-step direction, material (embedding), layer ID, number of secondaries (Phase 1 only — see Roadmap below).
**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.
**Model** (`giant/model/`): `DenoisingMLP` built from `ResBlock`s, with a `SinusoidalEmbedding` for the diffusion/flow time variable and a `ConditionEncoder` that fuses all conditioning inputs. `schedule.py` provides both a `CosineSchedule` for DDPM and the flow matching loss utilities (Lipman et al. 2022 conditional flow matching).
**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`).
`schedule.py` provides both a `CosineSchedule` for DDPM and the flow matching loss utilities (Lipman et al. 2022 conditional flow matching).
**Samplers** (`giant/sample.py`): DDPM, DDIM, and flow matching (ODE integration, ~10 steps). Flow matching is the primary mode.
**Validation** (`giant/validate.py`): step-level marginal comparisons; shower-level rollout validation is planned.
**Validation** (`giant/validate.py`): step-level marginal comparisons. Shower-level (rollout) observables live in `giant/analysis.py` (`compute_rollout_observables` + `plot_rollout_*`), fed by `giant rollout` output.
**Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower — each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on energy cutoff, per-track max steps, escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
## Roadmap
Phase 1 (current): number of secondaries is a conditioning input — model predicts only 9D post-step kinematics (including derived post_pos).
**Phase 1 (done):** number of secondaries and their total energy were conditioning inputs; the model predicted only the 9D primary post-step (energy-conservation PoC).
Phase 2 (target): model must jointly predict the number of secondaries and all their properties (energy, direction, species), requiring an extended output space and likely a set-based or autoregressive generation scheme for the variable-length secondary list.
**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`).
README.md
+32 -20
View File
@@ -2,35 +2,38 @@
**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate — a play on *Geant4* and the step function being the computationally heaviest part of the simulation.
Proof-of-concept surrogate model for the Geant4 step function. Given a pre-step particle state, the model samples a physically plausible post-step outcome — replacing the stochastic Geant4 physics engine with a trained conditional generative model.
Proof-of-concept surrogate model for the Geant4 step function. Given a pre-step particle state, the model samples a physically plausible post-step outcome — the primary's continuation plus the secondary particles it produces — replacing the stochastic Geant4 physics engine with a trained conditional generative model.
Training is driven entirely from parquet files of the miniCaloSim steps tree. No Geant4 runtime dependency.
## Architecture
Conditional **flow matching** model (Lipman et al. 2022): a small MLP learns a vector field mapping noise → step outcomes in ~10 ODE steps per sample. Falls back to DDPM for comparison.
A **two-stage conditional flow matching** model (Lipman et al. 2022): a small MLP learns a vector field mapping noise → step outcomes in ~10 ODE steps per sample. Falls back to DDPM for comparison.
**Output space (9D, diffused):**
**Stage 1 — primary (9D, diffused):**
| Index | Variable | Transform |
|-------|----------|-----------|
| Index | Variable | Encoding |
|-------|----------|----------|
| 0 | `step_length` [mm] | log |
| 1 | `ΔE = pre_E − post_E` [MeV] | log |
| 2 | `edep` [MeV] | log |
| 1–2 | `edep_logit`, `sec_logit` | ALR coords of the deposit/secondary/post-energy simplex |
| 3–5 | `post_dir` in local frame | unit vector |
| 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector |
The two energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — **energy conservation is built into the parametrization**, not left to the loss. Stage 1 also has a classifier head predicting the number of secondaries `n_sec ∈ {0..15}` from the conditioning alone.
Both `post_dir` and `travel_dir` are expressed in the coordinate frame where `pre_dir = ẑ`, making the scattering distribution nearly azimuthally symmetric. `post_pos` itself is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, so the two stay consistent by construction instead of being learned (and potentially diverging) independently.
**Conditioning:** PDG code (embedding), pre-step position, log(pre-energy), pre-step direction, material (embedding), layer ID, number of secondaries (Phase 1 only — see Roadmap below).
**Stage 2 — secondaries (`SecondaryDecoder`):** conditioned on the pre-step state *and* the Stage-1 outcome, a second flow net generates all `K_MAX = 15` secondary slots at once. Each slot carries a stick-breaking energy fraction, a local-frame direction, and a continuous particle-type embedding (snapped to the nearest PDG at inference), ordered by descending energy; slots beyond the predicted `n_sec` are masked. The secondary energies are a stick-breaking partition of the `e_sec` budget from Stage 1, so the full chain conserves energy. Each secondary's momentum is reconstructed afterward from `(energy, direction, species)` rather than predicted.
**Conditioning:** PDG code (embedding), pre-step position, log(pre-energy), pre-step direction, material (embedding), layer ID. (`n_sec` / `e_sec` are outputs now, not inputs.)
## Roadmap
The model is developed in two phases:
**Phase 1 (done):** `n_sec` and total secondary energy `e_sec` were conditioning inputs; the model predicted only the 9D primary post-step (energy-conservation PoC).
**Phase 1 (current):** The number of secondaries produced in each step is passed as a conditioning input. This makes training easier because the model has direct access to multiplicity information and can focus on learning the continuous post-step kinematics.
**Phase 2 (implemented — baseline):** the two-stage model above predicts `n_sec` and each secondary's energy, direction, and species jointly with the primary, so a shower rollout is fully self-contained.
**Phase 2 (target):** The number of secondaries is not given — the model must predict it jointly with all secondary properties (energy, direction, species) for each step. This requires extending the output space and likely an autoregressive or set-based generative approach for the variable-length secondary list.
**Next directions:** faster-eval architectures against a ~10× native-Geant4 budget (Wasserstein-GAN, mixture-of-experts routing tree), a multi-material sampling-calorimeter dataset, and physical-property conditioning over learned embeddings.
## Data
@@ -43,18 +46,21 @@ giant/
├── giant/
│ ├── data/
│ │ ├── loader.py # parquet → numpy arrays (incl. streaming/chunked reads)
│ │ ├── transforms.py # log transforms, local-frame rotation, normaliser
│ │ ├── transforms.py # log transforms, local-frame rotation, energy simplex, secondary encode/decode
│ │ └── dataset.py # StepsDataset / StreamingStepsDataset (PyTorch)
│ ├── model/
│ │ ├── network.py # SinusoidalEmbedding, ConditionEncoder, DenoisingMLP
│ │ ├── network.py # SinusoidalEmbedding, ConditionEncoder, DenoisingMLP, SecondaryDecoder
│ │ └── schedule.py # CosineSchedule (DDPM) and flow matching utilities
│ ├── constants.py # output/conditioning dims, K_MAX, secondary slot layout, schema keys
│ ├── config.py # default hyperparameters, TOML config merging, device autodetect
│ ├── pipeline.py # builds datasets/normalizers and kicks off a training run
│ ├── train.py # training loop, checkpointing, graceful shutdown
│ ├── sample.py # DDPM / DDIM / flow matching samplers
│ ├── train.py # two-stage training loop, checkpointing, graceful shutdown
│ ├── sample.py # DDPM / DDIM / flow matching samplers + secondary sampling
│ ├── geometry.py # GeometryOracle: position → (material, layer_id, escaped) for rollout
│ ├── rollout.py # autoregressive shower rollout driver
│ ├── validate.py # step-level marginal + KL-divergence validation
│ ├── analysis.py # notebook diagnostics: marginals, correlations, constraint checks
│ └── cli.py # `giant train` / `giant predict` Typer app
│ ├── analysis.py # step- and shower-level diagnostics: marginals, correlations, rollout observables
│ └── cli.py # `giant train` / `predict` / `rollout` Typer app
├── scripts/ # dataset/tooling logic, unified under the `dwarf` CLI (`uv run dwarf --help`)
│ ├── dwarf.py # Typer app: convert, migrate, bump-gen, bump-schema, status,
│ │ # update-manifest, create-manifest, make-root, hparam-scan
@@ -64,6 +70,7 @@ giant/
│ ├── bump_dataset_version.py # cut a new raw gen or parquet schema, with a logged reason —
│ │ # `dwarf bump-gen` / `bump-schema` / `status` / `update-manifest` / `create-manifest`
│ ├── create_root_files.py # generate new ROOT shards via a minicalosim executable — `dwarf make-root`
│ ├── geometry_oracle.py # fit a position → (material, layer_id) oracle — `dwarf build-geometry-oracle`
│ └── hparam_scan.py # hyperparameter grid scan over `giant train` runs — `dwarf hparam-scan`
└── tests/
```
@@ -77,18 +84,23 @@ uv sync --extra cpu --extra dev # add dev tools (pytest, ruff, ty)
`cpu` and `cuda` are mutually exclusive extras selecting the torch build; plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
## Training
## Training, prediction, and rollout
```bash
giant train path/to/steps.parquet --mode flow
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
# Full-shower rollout needs a geometry oracle (position → material/layer_id):
uv sync --extra cpu --extra geometry
dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl
giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
```
Both accept a TOML config file (`--config`) and CLI overrides for hyperparameters; see `--help` on either for the full option list.
`train`/`predict` accept a TOML config file (`--config`) and CLI overrides for hyperparameters; see `--help` on any command for the full option list. `giant rollout` seeds showers from the highest-energy entry step per event, then autoregressively steps the two-stage model to completion — pushing secondaries as new tracks and looking up `material`/`layer_id` from the oracle at each step. Tracks terminate on energy cutoff, per-track max steps, detector escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
## Validation
`giant.validate.validate_marginals` runs step-level marginal and KL-divergence checks during training (`--validate-every`). For deeper, notebook-driven diagnostics on a trained checkpoint — stratified marginals, correlation structure, physical constraint violations — see `giant.analysis`.
`giant.validate.validate_marginals` runs step-level marginal and KL-divergence checks during training (`--validate-every`). For deeper diagnostics on a trained checkpoint — stratified marginals, correlation structure, physical-constraint violations, and shower-level rollout observables (longitudinal/transverse profiles, PDG energy shares) — see `giant.analysis`.
## Development
analysis/compare_ode_steps_energy_conservation.py
+34 -14
View File
@@ -38,9 +38,7 @@ def per_event(file: str) -> dict:
E0 = float(primary_E[0])
real = pe["real_total_edep"].to_numpy()
gen = pe["gen_total_edep"].to_numpy()
n_steps = (
pl.scan_parquet(file).select(pl.len()).collect(engine="streaming").item()
)
n_steps = pl.scan_parquet(file).select(pl.len()).collect(engine="streaming").item()
return {
"E0": E0,
"n_events": pe.height,
@@ -78,9 +76,12 @@ fmt_row("gen mean/E0", lambda r: f"{r['gen'].mean() / r['E0']:.4f}")
fmt_row("gen max/E0", lambda r: f"{r['gen'].max() / r['E0']:.4f}")
fmt_row("gen p99/E0", lambda r: f"{np.quantile(r['gen'], 0.99) / r['E0']:.4f}")
fmt_row("frac events gen>E0", lambda r: f"{np.mean(r['gen'] > r['E0']):.4f}")
disp_ratio = lambda r: (r["gen"].std() / r["gen"].mean()) / (
r["real"].std() / r["real"].mean()
)
def disp_ratio(r):
return (r["gen"].std() / r["gen"].mean()) / (r["real"].std() / r["real"].mean())
fmt_row("dispersion ratio gen/real", lambda r: f"{disp_ratio(r):.2f}x")
print("=" * 80)
@@ -91,21 +92,31 @@ E0, real_tot, gen_tot = r["E0"], r["real"], r["gen"]
fig, ax = plt.subplots(figsize=(6, 4))
edges = _hist_edges(real_tot, gen_tot, bins=50).tolist()
ax.hist(
real_tot, bins=edges, density=True, histtype="step",
real_tot,
bins=edges,
density=True,
histtype="step",
label=f"real (σ/μ={real_tot.std() / real_tot.mean():.3f})",
)
ax.hist(
gen_tot, bins=edges, density=True, histtype="step",
gen_tot,
bins=edges,
density=True,
histtype="step",
label=f"generated (σ/μ={gen_tot.std() / gen_tot.mean():.3f}, "
f"{np.mean(gen_tot > E0):.1%} > E0)",
)
ax.axvline(E0, color="k", linestyle="--", linewidth=1, label=f"incident energy E0={E0:.0f} MeV")
ax.axvline(
E0, color="k", linestyle="--", linewidth=1, label=f"incident energy E0={E0:.0f} MeV"
)
ax.set_yscale("log")
ax.set_xlabel("total deposited energy per event [MeV]")
ax.set_title("20 ODE steps")
ax.legend(fontsize=8)
fig.tight_layout()
fig.savefig(OUT / f"{PREFIX20}-event-total-energy-vs-E0.png", dpi=150, bbox_inches="tight")
fig.savefig(
OUT / f"{PREFIX20}-event-total-energy-vs-E0.png", dpi=150, bbox_inches="tight"
)
fig, ax = plt.subplots(figsize=(6, 4))
ratio_real = real_tot / E0
@@ -128,12 +139,19 @@ all_arrays = [results["10-step (baseline)"]["real"]] + [
]
edges = _hist_edges(*all_arrays, bins=60).tolist()
ax.hist(
results["20-step"]["real"], bins=edges, density=True, histtype="step",
color="k", label="real",
results["20-step"]["real"],
bins=edges,
density=True,
histtype="step",
color="k",
label="real",
)
for name, r2 in results.items():
ax.hist(
r2["gen"], bins=edges, density=True, histtype="step",
r2["gen"],
bins=edges,
density=True,
histtype="step",
label=f"gen {name} ({np.mean(r2['gen'] > r2['E0']):.1%} > E0)",
)
ax.axvline(E0, color="gray", linestyle="--", linewidth=1, label=f"E0={E0:.0f} MeV")
@@ -142,6 +160,8 @@ ax.set_xlabel("total deposited energy per event [MeV]")
ax.set_title("Generated event energy: 10 vs 20 ODE steps")
ax.legend(fontsize=8)
fig.tight_layout()
fig.savefig(OUT / f"{PREFIX20}-compare-event-total-energy.png", dpi=150, bbox_inches="tight")
fig.savefig(
OUT / f"{PREFIX20}-compare-event-total-energy.png", dpi=150, bbox_inches="tight"
)
print("DONE")
analysis/compare_ode_steps_kl.py
+1 -3
View File
@@ -66,6 +66,4 @@ print(
f"{'SUM':<14}{tot['10-step']:>14.5f}{tot['20-step']:>14.5f}"
f"{tot['20-step'] / tot['10-step']:>14.2f}"
)
print(
f"{'MEAN':<14}{tot['10-step'] / 9:>14.5f}{tot['20-step'] / 9:>14.5f}"
)
print(f"{'MEAN':<14}{tot['10-step'] / 9:>14.5f}{tot['20-step'] / 9:>14.5f}")
analysis/export_energy_conservation_poc.py
+6 -2
View File
@@ -77,12 +77,16 @@ ax.hist(
label=f"generated (σ/μ={gen_tot.std() / gen_tot.mean():.3f}, "
f"{np.mean(gen_tot > E0):.1%} > E0)",
)
ax.axvline(E0, color="k", linestyle="--", linewidth=1, label=f"incident energy E0={E0:.0f} MeV")
ax.axvline(
E0, color="k", linestyle="--", linewidth=1, label=f"incident energy E0={E0:.0f} MeV"
)
ax.set_yscale("log")
ax.set_xlabel("total deposited energy per event [MeV]")
ax.legend(fontsize=8)
fig.tight_layout()
fig.savefig(OUT / f"{PREFIX}-event-total-energy-vs-E0.png", dpi=150, bbox_inches="tight")
fig.savefig(
OUT / f"{PREFIX}-event-total-energy-vs-E0.png", dpi=150, bbox_inches="tight"
)
print("=== plot: total edep / incident energy ratio ===")
fig, ax = plt.subplots(figsize=(6, 4))
analysis/export_rollout_observables.py
+48
View File
@@ -0,0 +1,48 @@
"""Export shower-level plots from a `giant rollout` steps parquet.
Not part of the package; run manually. Optionally overlays the real showers
seeded from the same events (a `giant predict --coord local` file) by passing a
reference path. Usage::
python analysis/export_rollout_observables.py ROLLOUT.parquet [REFERENCE_local.parquet]
"""
import sys
from pathlib import Path
import giant.analysis as a
rollout_file = sys.argv[1] if len(sys.argv) > 1 else "rollout.parquet"
reference_file = (
sys.argv[2] if len(sys.argv) > 2 and sys.argv[2] not in ("", "-") else None
)
OUT = Path(sys.argv[3]) if len(sys.argv) > 3 else Path(".")
print(f"=== computing rollout observables: {rollout_file} ===")
obs = a.compute_rollout_observables(rollout_file)
tbl = obs.event_table
print(f"n events: {len(tbl)}")
print(
f"total_edep/event: mean={tbl['total_edep'].mean():.4g} MeV "
f"leaked_E/event: mean={tbl['leaked_E'].mean():.4g} MeV "
f"n_tracks/event: mean={tbl['n_tracks'].mean():.1f} "
f"n_steps/event: mean={tbl['n_steps'].mean():.1f}"
)
reference = None
if reference_file is not None:
print(f"=== computing real reference: {reference_file} ===")
reference = a.compute_event_observables_pl(reference_file)
print("=== plots ===")
for name, fn in [
("longitudinal", a.plot_rollout_longitudinal),
("transverse", a.plot_rollout_transverse),
("total-energy", a.plot_rollout_total_energy),
]:
fig = fn(obs, reference=reference)
path = OUT / f"rollout-{name}.png"
fig.savefig(path, dpi=150, bbox_inches="tight")
print(f"wrote {path}")
print("DONE")
analysis/rollout_validation.ipynb
+454
View File
File diff suppressed because one or more lines are too long
docs/phase2_plan.md
+172
View File
@@ -0,0 +1,172 @@
# Phase 2: Secondary Particle Prediction
## Context
Phase 1 takes `n_sec` (secondary count) and `e_sec` (total secondary energy) as **conditioning inputs**. Phase 2 must instead **predict** them, making the surrogate self-contained for shower rollout. Per Jan's 2026-06-29 decision: hard discrete `n_sec` integer head; escalation to Gumbel-Softmax only if empirically needed.
Two-stage factorization:
- **Stage 1**: existing 9D flow model (reduced conditioning: drop `n_sec` + `log(e_sec)`) + a new discrete `n_sec` classification head
- **Stage 2**: non-AR flow matching over `K_MAX` secondary slots simultaneously, each slot predicting `(stick_break_logit, dir_local_3D, type_emb)` — conditioned on pre-step state + Stage 1 output; padded slots masked from loss
Training: joint, combined loss `L = L_flow_s1 + λ_nsec * L_nsec + λ_s2 * L_flow_s2`.
---
## Prerequisite: Determine K_MAX
Before implementing, run a quick analysis over existing parquet files to find `max(n_sec)` and the 99th percentile. Expected to be 5–20 for EM shower steps. Set `K_MAX` as a constant in `giant/constants.py` (suggest 15 as a starting point, revise from data).
---
## New Branch
```bash
git checkout -b phase2-secondary-prediction master
```
---
## Part A — Data Pipeline
### A1. `scripts/steps_to_parquet.py`
Extend `_add_secondary_energy` to also collect per-secondary attributes from the spawning tree join:
- For each `child_track_id`, look up the child's first step → get `pdg`, `pre_E`, `pre_dx/dy/dz`
- Emit list columns in the parquet: `sec_pdg_list`, `sec_E_list`, `sec_dx_list`, `sec_dy_list`, `sec_dz_list`
- Lists are sorted **descending by energy** at write time
- Truncate to `K_MAX` entries if needed (flag if any row truncated)
Re-run ROOT→parquet conversion after this change.
### A2. `giant/data/loader.py`
In `_df_to_dict`: read the five new list columns. Pad each to length `K_MAX` with zeros (energy) / sentinel values (pdg → 0, dir → (0,0,1)). Return as fixed-shape arrays `(N, K_MAX)` / `(N, K_MAX, 3)`.
Also return a boolean validity mask `sec_valid` of shape `(N, K_MAX)`: `True` for slots `i < n_sec`.
### A3. `giant/data/transforms.py`
Add `encode_secondaries(sec_pdg_list, sec_E_list, sec_dir_list, sec_valid, e_sec, pdg_emb_weight, pre_dir, K_MAX)`:
1. **Direction**: call existing `local_frame_rotation` per slot
2. **Energy (stick-breaking)**:
- Slot 0: `f_0 = E_0 / e_sec` → logit `log(f_0/(1-f_0))` (clamped)
- Slot i: `f_i = E_i / (e_sec - sum(E_0..E_{i-1}))` → logit
- Last valid slot: logit = large positive constant (takes all remaining budget)
- Padding slots (beyond `n_sec`): set logit = 0, masked out of loss anyway
3. **Type embedding**: index into `pdg_emb_weight` (the PDG embedding table weights) to get the target embedding vector for each secondary's `pdg`. Shape `(K_MAX, emb_dim)`.
Returns `sec_targets: (K_MAX, 1 + 3 + emb_dim)` and `sec_valid: (K_MAX,)`.
Inverse (`decode_secondaries`): sigmoid stick-breaking fractions → energies, inv local frame rotation → world dirs, nearest-neighbor lookup in PDG embedding table → pdg code.
### A4. `giant/data/dataset.py`
Update `build_features` and `StreamingStepsDataset.__iter__` to also yield `sec_targets` and `sec_valid` alongside the existing `(cond_cont, cond_cat, x1)` batch items.
---
## Part B — Constants (`giant/constants.py`)
- `COND_DIM`: 10 → **8** (remove `n_sec` and `log(e_sec)`)
- Add `K_MAX: int` (set after data analysis, e.g. 15)
- Add `SEC_SLOT_DIM: int` (= 4 + `emb_dim` = 20 for default emb_dim=16; 1 stick + 3 dir + 16 type)
- Add `SEC_DIM: int = K_MAX * SEC_SLOT_DIM` (flattened Stage 2 target dimension)
- Update `LOCAL_TARGET_NAMES` (Stage 1 only, still 9D)
---
## Part C — Model (`giant/model/network.py`)
### C1. `DenoisingMLP` — Stage 1 (minimal changes)
- `ConditionEncoder.cont_dim` drops from 10 to 8 (COND_DIM change propagates automatically)
- Add `n_sec_head = nn.Sequential(Linear(cond_out_dim, hidden_dim//2), SiLU(), Linear(hidden_dim//2, K_MAX + 1))` applied to `c_emb` (the condition encoding, not the diffused latent)
- Add method `predict_n_sec(cond_cont, cond_cat) -> Tensor[B, K_MAX+1]` — no diffusion, just encode conditioning and run the head
### C2. `SecondaryDecoder` — Stage 2 (new class)
Architecture mirrors `DenoisingMLP` but:
- **Input**: `x_t` of shape `(B, SEC_DIM)` (flattened K_MAX secondary slots)
- **Conditioning**: pre-step state (8D cont + 2 cat → same ConditionEncoder as Stage 1) concatenated with Stage 1 output (9D normalized target, detached from Stage 1 loss for stability initially). Total cond dim to the ResBlocks: `time_dim + cond_s1_out_dim + 9`
- **Output**: vector field of shape `(B, SEC_DIM)`
- Uses same `ResBlock` / `SinusoidalEmbedding` / `ConditionEncoder` building blocks
A `SecondaryConditionEncoder` wraps the base `ConditionEncoder` and concatenates the Stage 1 output:
```python
class SecondaryConditionEncoder(nn.Module):
# base: ConditionEncoder(pdg_vocab, mat_vocab, 8, emb_dim, cond_out_dim)
# stage1_proj: Linear(X_DIM, stage1_cond_dim)
# mlp: fuses both
```
---
## Part D — Loss / Training
### `giant/model/schedule.py`
Add `flow_matching_loss_masked(model, x1, cond_cont, cond_cat, mask)`:
- Same as `flow_matching_loss` but divides by `mask.sum()` instead of `B * SEC_DIM`, zeroing out padded slots before averaging. `mask` shape: `(B, K_MAX)`, broadcast over slot dims.
### `giant/train.py`
Batch now unpacks as `(cond_cont, cond_cat, x1_s1, n_sec_target, x1_s2, sec_mask)`.
Combined loss per batch:
```
L_s1 = flow_matching_loss(stage1_model, x1_s1, cond_cont, cond_cat)
L_nsec = cross_entropy(stage1_model.predict_n_sec(cond_cont, cond_cat), n_sec_target)
L_s2 = flow_matching_loss_masked(sec_decoder, x1_s2, cond_cont, cond_cat, stage1_detached, sec_mask)
L = L_s1 + lambda_nsec * L_nsec + lambda_s2 * L_s2
```
Config adds `lambda_nsec` (suggest 0.1) and `lambda_s2` (suggest 1.0) under `[train]`.
Both `stage1_model` and `sec_decoder` share a single `optimizer` (AdamW over all parameters).
Checkpoint saves both `stage1_model.state_dict()` and `sec_decoder.state_dict()`, plus `K_MAX` and `SEC_SLOT_DIM` in `model_config`.
### `giant/pipeline.py`
- Compute `K_MAX` from data (max `n_sec` over training events) before constructing models
- Build both `DenoisingMLP` and `SecondaryDecoder`, pass both to `run_training`
---
## Part E — Sampling (`giant/sample.py`)
```python
def sample_stage1(model, cond_cont, cond_cat, steps=10):
# Euler ODE → primary sample (9D), + argmax n_sec head
...
def sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec, steps=10):
# Euler ODE on SEC_DIM → decode stick-breaking → energies
# inv_local_frame_rotation → world-frame dirs
# nearest-neighbor in pdg_emb_weight → pdg codes
# mask slots >= n_sec
...
```
---
## Part F — Wiring
- **`giant/validate.py`**: add secondary-specific marginals (n_sec distribution, species distribution, energy fraction per slot)
- **`giant/cli.py`**: `predict` command loads both checkpoints, calls both samplers, appends secondary columns to output parquet
---
## Type embedding design note
The type embedding target at training is `pdg_emb.weight[sec_pdg_idx]` (the Stage 1 PDG embedding table rows). Gradients flow into the embedding table from both the conditioning path (input PDG) and the secondary type loss — this is intentional; the shared embedding space is the bridge. At inference, snap: `argmin_k ||pred_emb - pdg_emb.weight[k]||`.
---
## Verification
1. `uv run pytest` — existing tests pass (Stage 1 shape/interface unchanged beyond COND_DIM)
2. Unit tests for `encode_secondaries` / `decode_secondaries` (round-trip: energies sum to `e_sec`, directions are unit vectors)
3. Unit test for `flow_matching_loss_masked`: verify padded slots contribute zero gradient
4. Short training run (1–2 epochs): confirm all three loss components decrease
5. Sampling smoke test: verify `sum(sec_E) ≈ e_sec` per sample, all directions unit-normed
giant/analysis.py
+655 -19
View File
@@ -51,6 +51,21 @@ how much of the total energy/length, see `pdg_contribution_table_pl`::
plot_pdg_energy_share(table)
plot_pdg_length_share(table)
To run the tiers 1-3 checks above against a full `giant rollout` shower
instead of one-step-ahead predict output, use `load_rollout_vs_truth` in place
of `load_predicted_local` — it builds the same `SampleCollection` from a
`giant rollout` output file ("generated") and any held-out file sharing
`giant train`'s input schema ("real"); the two are independent, unpaired
files (a rollout doesn't replay real events row-for-row), unlike the other
loader's paired pred_*/true_* columns::
from giant.analysis import load_rollout_vs_truth
samples = load_rollout_vs_truth("path/to/rollout.parquet", "path/to/val.parquet")
plot_marginals(samples, group_by="energy")
# ... same plot_kl_bars / plot_correlation_matrices / plot_pairwise /
# plot_direction_alignment / plot_constraint_violations as above.
Four tiers of checks, building on the aggregate marginal/KL check in
`giant.validate.validate_marginals`:
@@ -91,11 +106,18 @@ from giant.constants import (
PREDICT_COORD_METADATA_KEY,
PREDICT_SCHEMA_VERSION,
PREDICT_SCHEMA_VERSION_KEY,
ROLLOUT_COORD_VALUE,
TERM_ENERGY_CUTOFF,
TERM_ESCAPED,
TERM_MAX_STEPS,
TERM_UNKNOWN_PDG,
)
from giant.data.transforms import (
energy_simplex_decode,
inv_log_transform,
local_frame_rotation,
reconstruct_post_pos,
travel_direction,
)
from giant.validate import _histogram_kl
@@ -177,7 +199,16 @@ class SampleCollection:
pdg: np.ndarray # (N,) raw PDG codes
material: np.ndarray # (N,) raw material names
real_raw: np.ndarray # (N, 9) denormalized + delogged real targets
gen_raw: np.ndarray # (N, 9) denormalized + delogged generated targets
gen_raw: np.ndarray # (M, 9) denormalized + delogged generated targets
# Set these three when real_raw/gen_raw are *unpaired* — independent files
# with their own row counts and conditioning (e.g. `load_rollout_vs_truth`,
# comparing a `giant rollout` shower against a held-out truth file) — rather
# than the row-for-row pred_*/true_* pairing `load_predicted_local` produces.
# None (the default) means "same as the real-side field above", which
# reproduces the original paired behavior exactly.
cond_cont_raw_gen: np.ndarray | None = None
pdg_gen: np.ndarray | None = None
material_gen: np.ndarray | None = None
def _check_predict_metadata(path: Path) -> None:
@@ -292,28 +323,66 @@ def load_predicted_local(
# ---------------------------------------------------------------------------
def _gen_pdg(collection: SampleCollection) -> np.ndarray:
return collection.pdg if collection.pdg_gen is None else collection.pdg_gen
def _gen_material(collection: SampleCollection) -> np.ndarray:
return (
collection.material
if collection.material_gen is None
else collection.material_gen
)
def _gen_cond(collection: SampleCollection) -> np.ndarray:
return (
collection.cond_cont_raw
if collection.cond_cont_raw_gen is None
else collection.cond_cont_raw_gen
)
def _group_labels(
collection: SampleCollection,
group_by: str | None,
n_energy_bins: int,
) -> list[tuple[str, np.ndarray]]:
n = len(collection.pdg)
) -> list[tuple[str, np.ndarray, np.ndarray]]:
"""Per-group `(label, mask_real, mask_gen)` triples.
`mask_real` indexes `collection.real_raw` (via `pdg`/`material`/
`cond_cont_raw`); `mask_gen` indexes `collection.gen_raw` via the `*_gen`
fields when set (unpaired real/gen — see `SampleCollection`), or the same
real-side arrays otherwise, which collapses to a single shared mask — the
original paired behavior (real/gen same length, row-for-row).
"""
gen_pdg, gen_material, gen_cond = (
_gen_pdg(collection),
_gen_material(collection),
_gen_cond(collection),
)
n_real, n_gen = len(collection.pdg), len(gen_pdg)
if group_by is None:
return [("all", np.ones(n, dtype=bool))]
return [("all", np.ones(n_real, dtype=bool), np.ones(n_gen, dtype=bool))]
if group_by == "pdg":
return [(f"pdg={v}", collection.pdg == v) for v in np.unique(collection.pdg)]
values = np.unique(np.concatenate([collection.pdg, gen_pdg]))
return [(f"pdg={v}", collection.pdg == v, gen_pdg == v) for v in values]
if group_by == "material":
values = np.unique(np.concatenate([collection.material, gen_material]))
return [
(f"material={v}", collection.material == v)
for v in np.unique(collection.material)
(f"material={v}", collection.material == v, gen_material == v)
for v in values
]
if group_by == "energy":
pre_E = collection.cond_cont_raw[:, 3]
edges = np.quantile(pre_E, np.linspace(0, 1, n_energy_bins + 1))
real_E, gen_E = collection.cond_cont_raw[:, 3], gen_cond[:, 3]
edges = np.quantile(
np.concatenate([real_E, gen_E]), np.linspace(0, 1, n_energy_bins + 1)
)
edges[-1] += 1e-6
bin_idx = np.digitize(pre_E, edges[1:-1])
real_bin = np.digitize(real_E, edges[1:-1])
gen_bin = np.digitize(gen_E, edges[1:-1])
return [
(f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})", bin_idx == i)
(f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})", real_bin == i, gen_bin == i)
for i in range(n_energy_bins)
]
raise ValueError(f"unknown group_by={group_by!r}")
@@ -332,16 +401,19 @@ def marginal_table(
failure modes hidden by the aggregate surface at the top.
"""
rows = []
for label, mask in _group_labels(collection, group_by, n_energy_bins):
if mask.sum() < 2:
for label, mask_real, mask_gen in _group_labels(
collection, group_by, n_energy_bins
):
if mask_real.sum() < 2 or mask_gen.sum() < 2:
continue
real, gen = collection.real_raw[mask], collection.gen_raw[mask]
real, gen = collection.real_raw[mask_real], collection.gen_raw[mask_gen]
for j, name in enumerate(RAW_TARGET_NAMES):
rows.append(
{
"group": label,
"dim": name,
"n": int(mask.sum()),
"n": int(mask_real.sum()),
"n_gen": int(mask_gen.sum()),
"real_mean": real[:, j].mean(),
"gen_mean": gen[:, j].mean(),
"real_std": real[:, j].std(),
@@ -647,8 +719,8 @@ def plot_marginals(
squeeze=False,
figsize=(figsize_per_axis[0] * n_cols, figsize_per_axis[1] * n_rows),
)
for row, (label, mask) in enumerate(groups):
real, gen = collection.real_raw[mask], collection.gen_raw[mask]
for row, (label, mask_real, mask_gen) in enumerate(groups):
real, gen = collection.real_raw[mask_real], collection.gen_raw[mask_gen]
for col, j in enumerate(dim_idx):
ax = axes[row][col]
edges = _hist_edges(real[:, j], gen[:, j], bins=bins)
@@ -816,8 +888,6 @@ def plot_pairwise(
"""
pairs = pairs or _DEFAULT_PAIRS
rng = np.random.default_rng(seed)
n = len(collection.pdg)
idx = rng.choice(n, size=min(n_sample, n), replace=False)
fig, axes = plt.subplots(2, len(pairs), squeeze=False, figsize=(4 * len(pairs), 7))
for col, (a, b) in enumerate(pairs):
@@ -825,6 +895,8 @@ def plot_pairwise(
for row, (data, title) in enumerate(
[(collection.real_raw, "real"), (collection.gen_raw, "generated")]
):
n = len(data)
idx = rng.choice(n, size=min(n_sample, n), replace=False)
ax = axes[row][col]
ax.scatter(data[idx, ia], data[idx, ib], s=3, alpha=0.3)
ax.set_xlabel(a)
@@ -1693,3 +1765,567 @@ def plot_pdg_length_share(table: pl.DataFrame, max_slices: int = 6):
"length traveled share by particle type",
max_slices,
)
# ── Shower-rollout observables (single-sided; `giant rollout` output) ──────────
#
# The rollout writes a world-frame steps parquet (post_x/y/z + edep already in
# physical units), unlike the local predict schema these other diagnostics read.
# So this is a standalone, generated-only path: no paired `true_*` columns, no
# energy-simplex decode. Because rollout showers are seeded from real events, the
# same event_ids can be overlaid against a real reference computed elsewhere.
_ROLLOUT_COLS = [
"event_id",
"track_id",
"termination_reason",
"pre_x",
"pre_y",
"pre_z",
"pre_dx",
"pre_dy",
"pre_dz",
"pre_E",
"post_x",
"post_y",
"post_z",
"edep",
]
def _check_rollout_metadata(path: Path) -> None:
"""Raise if `path` carries coord metadata that isn't `ROLLOUT_COORD_VALUE`.
A missing tag (older rollout output, predating tagging) is let through
silently, matching `giant predict`/`giant rollout`'s own leniency; a tag
that's present but wrong is a real mismatch.
"""
metadata = pq.read_schema(path).metadata or {}
coord = metadata.get(PREDICT_COORD_METADATA_KEY.encode())
if coord is not None and coord.decode() != ROLLOUT_COORD_VALUE:
raise ValueError(
f"{path} is not a rollout file (coord={coord.decode()!r}); "
"expected a `giant rollout` output"
)
@dataclass
class RolloutObservables:
event_table: pd.DataFrame # one row per event_id (mm/MeV)
depth_edges: np.ndarray # (depth_bins+1,) mm, along shower axis
transverse_edges: np.ndarray # (transverse_bins+1,) mm, perpendicular
depth_profile: np.ndarray # (depth_bins,) mean edep/event/bin, MeV
depth_profile_std: np.ndarray
transverse_profile: np.ndarray
transverse_profile_std: np.ndarray
def _event_axis_depth_transverse(
df: pd.DataFrame, depth_bins: int, transverse_bins: int
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Shared core of `compute_rollout_observables`/`compute_truth_observables`.
`df` needs `event_id`/`pre_x,y,z`/`pre_dx,dy,dz`/`pre_E`/`post_x,y,z`/`edep`
— both a `giant rollout` file and a raw truth-schema steps file carry these.
Per event, the highest-`pre_E` row fixes the shower axis/entry point; every
row's `edep` at `post_pos` is projected onto depth-along-axis and
transverse-distance-from-axis, then binned.
Returns `(ev_ids, row_ev, depth_edges, transverse_edges, depth_ev, trans_ev)`
— `row_ev` maps each row of `df` to its event's index in `ev_ids`;
`depth_ev`/`trans_ev` are `(n_events, bins)` per-event-per-bin edep sums.
"""
entry_idx = df.groupby("event_id")["pre_E"].idxmax()
entry = df.loc[entry_idx].set_index("event_id")
ax = entry[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float64).copy()
ax /= np.clip(np.linalg.norm(ax, axis=1, keepdims=True), 1e-12, None)
ent = entry[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float64)
ev_ids = entry.index.to_numpy()
ev_row = {int(e): i for i, e in enumerate(ev_ids)}
row_ev = df["event_id"].map(ev_row).to_numpy()
post = df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float64)
rel = post - ent[row_ev]
depth = np.einsum("ij,ij->i", rel, ax[row_ev])
transverse = np.linalg.norm(rel - depth[:, None] * ax[row_ev], axis=1)
edep = df["edep"].to_numpy(dtype=np.float64)
n_events = len(ev_ids)
depth_lo, depth_hi = np.quantile(depth, [0.001, 0.999])
if not (depth_hi - depth_lo > 1e-6 * max(abs(depth_hi), 1.0)):
depth_lo, depth_hi = depth_lo - 0.5, depth_hi + 0.5
depth_edges = np.linspace(depth_lo, depth_hi, depth_bins + 1)
transverse_edges = np.linspace(
0.0, max(np.quantile(transverse, 0.999), 1e-6), transverse_bins + 1
)
d_bin = np.clip(np.digitize(depth, depth_edges) - 1, 0, depth_bins - 1)
t_bin = np.clip(
np.digitize(transverse, transverse_edges) - 1, 0, transverse_bins - 1
)
# Per-(event, bin) edep sums, then mean/std across events.
depth_ev = np.zeros((n_events, depth_bins))
trans_ev = np.zeros((n_events, transverse_bins))
np.add.at(depth_ev, (row_ev, d_bin), edep)
np.add.at(trans_ev, (row_ev, t_bin), edep)
return ev_ids, row_ev, depth_edges, transverse_edges, depth_ev, trans_ev
def _centroid_depth(depth_ev: np.ndarray, depth_edges: np.ndarray) -> np.ndarray:
"""Energy-weighted centroid depth per event, from the binned edep sums."""
bin_centers = 0.5 * (depth_edges[:-1] + depth_edges[1:])
tot = depth_ev.sum(axis=1)
centroid = (depth_ev * bin_centers).sum(axis=1) / np.where(tot > 0, tot, 1.0)
return np.where(tot > 0, centroid, 0.0)
def compute_rollout_observables(
path: str | Path,
depth_bins: int = 20,
transverse_bins: int = 20,
) -> RolloutObservables:
"""Compute shower observables from a `giant rollout` steps parquet.
Per event the primary entry step (highest-`pre_E` row) fixes the shower axis
and entry point; every step's deposit (`edep` at `post_pos`) is projected onto
depth-along-axis and transverse-distance-from-axis, then binned. Returns
per-event totals plus dataset-mean longitudinal/transverse profiles.
See `compute_truth_observables` for the truth-schema counterpart, shaped to
plug straight into this function's own output as a `plot_rollout_*`
`reference=` overlay.
"""
_check_rollout_metadata(Path(path))
df = pd.read_parquet(path, columns=_ROLLOUT_COLS)
ev_ids, row_ev, depth_edges, transverse_edges, depth_ev, trans_ev = (
_event_axis_depth_transverse(df, depth_bins, transverse_bins)
)
# Per-event scalar table.
is_leak = df["termination_reason"].to_numpy() == TERM_ESCAPED
per_ev = df.groupby("event_id").agg(
n_steps=("edep", "size"),
total_edep=("edep", "sum"),
)
per_ev["n_tracks"] = df.groupby("event_id")["track_id"].nunique()
per_ev["leaked_E"] = (
df.assign(_leak=np.where(is_leak, df["pre_E"], 0.0))
.groupby("event_id")["_leak"]
.sum()
)
per_ev = per_ev.reindex(ev_ids)
per_ev["centroid_depth"] = _centroid_depth(depth_ev, depth_edges)
per_ev.index.name = "event_id"
return RolloutObservables(
event_table=per_ev.reset_index(),
depth_edges=depth_edges,
transverse_edges=transverse_edges,
depth_profile=depth_ev.mean(0),
depth_profile_std=depth_ev.std(0),
transverse_profile=trans_ev.mean(0),
transverse_profile_std=trans_ev.std(0),
)
_TRUTH_EVENT_COLS = [
"event_id",
"pre_x",
"pre_y",
"pre_z",
"pre_dx",
"pre_dy",
"pre_dz",
"pre_E",
"post_x",
"post_y",
"post_z",
"edep",
"step_length",
]
@dataclass
class TruthObservables:
"""Truth-schema event-level observables, shaped for `plot_rollout_*`'s `reference=`.
The truth-side counterpart to `RolloutObservables`, computed directly from
a raw truth-schema steps file (the same file `load_rollout_vs_truth` takes
as `truth_path`) rather than needing a separate paired `giant predict
--coord local` file covering the same events. Field names mirror
`EventObservables`'s `real_*` convention — `plot_rollout_longitudinal`/
`plot_rollout_transverse`/`plot_rollout_total_energy` already read exactly
these names off their `reference` argument — since this object is only
ever used as a reference overlay, never plotted standalone.
"""
event_table: pd.DataFrame # one row per event_id (mm/MeV)
depth_edges: np.ndarray
transverse_edges: np.ndarray
real_depth_profile: np.ndarray
real_depth_profile_std: np.ndarray
real_transverse_profile: np.ndarray
real_transverse_profile_std: np.ndarray
def compute_truth_observables(
path: str | Path,
depth_bins: int = 20,
transverse_bins: int = 20,
) -> TruthObservables:
"""Compute shower observables from a raw truth-schema steps parquet.
Lets the same held-out truth file used by `load_rollout_vs_truth` (Tier
1-3) also supply the Tier 4 real-shower overlay — pass the result as
`reference=` to `plot_rollout_longitudinal`/`plot_rollout_transverse`/
`plot_rollout_total_energy` — without needing a separate paired `giant
predict --coord local` file for the same events. Same per-event
axis/depth/transverse construction as `compute_rollout_observables` (see
`_event_axis_depth_transverse`), just without the `track_id`/
`termination_reason` columns a rollout file (but not a truth file) has, so
there's no `n_tracks`/`leaked_E` in `event_table`.
"""
df = pd.read_parquet(path, columns=_TRUTH_EVENT_COLS)
ev_ids, row_ev, depth_edges, transverse_edges, depth_ev, trans_ev = (
_event_axis_depth_transverse(df, depth_bins, transverse_bins)
)
per_ev = df.groupby("event_id").agg(
n_steps=("edep", "size"),
real_total_edep=("edep", "sum"),
real_total_length=("step_length", "sum"),
)
per_ev = per_ev.reindex(ev_ids)
per_ev["real_centroid_depth"] = _centroid_depth(depth_ev, depth_edges)
per_ev.index.name = "event_id"
return TruthObservables(
event_table=per_ev.reset_index(),
depth_edges=depth_edges,
transverse_edges=transverse_edges,
real_depth_profile=depth_ev.mean(0),
real_depth_profile_std=depth_ev.std(0),
real_transverse_profile=trans_ev.mean(0),
real_transverse_profile_std=trans_ev.std(0),
)
def plot_rollout_longitudinal(obs: RolloutObservables, reference=None):
"""Mean edep vs depth along the shower axis; optional real-reference overlay.
`reference` may be an `EventObservables` or a `TruthObservables` (either
exposes `real_depth_profile`) to overlay the real showers seeded from the
same events — `compute_truth_observables` builds the latter directly from
a raw truth-schema file, with no paired predict-schema file needed.
"""
centers = 0.5 * (obs.depth_edges[:-1] + obs.depth_edges[1:])
fig, ax = plt.subplots(figsize=(7, 4))
ax.plot(centers, obs.depth_profile, label="rollout", color="C0")
ax.fill_between(
centers,
obs.depth_profile - obs.depth_profile_std,
obs.depth_profile + obs.depth_profile_std,
alpha=0.2,
color="C0",
)
if reference is not None:
rc = 0.5 * (reference.depth_edges[:-1] + reference.depth_edges[1:])
ax.plot(rc, reference.real_depth_profile, label="real", color="k", ls="--")
ax.set_xlabel("depth along axis [mm]")
ax.set_ylabel("mean E_dep / event [MeV]")
ax.set_title("Longitudinal shower profile")
ax.legend()
fig.tight_layout()
return fig
def plot_rollout_transverse(obs: RolloutObservables, reference=None):
"""Mean edep vs transverse distance from the shower axis (Molière-style)."""
centers = 0.5 * (obs.transverse_edges[:-1] + obs.transverse_edges[1:])
fig, ax = plt.subplots(figsize=(7, 4))
ax.plot(centers, obs.transverse_profile, label="rollout", color="C0")
ax.fill_between(
centers,
obs.transverse_profile - obs.transverse_profile_std,
obs.transverse_profile + obs.transverse_profile_std,
alpha=0.2,
color="C0",
)
if reference is not None:
rc = 0.5 * (reference.transverse_edges[:-1] + reference.transverse_edges[1:])
ax.plot(rc, reference.real_transverse_profile, label="real", color="k", ls="--")
ax.set_xlabel("transverse distance [mm]")
ax.set_ylabel("mean E_dep / event [MeV]")
ax.set_title("Transverse shower profile")
ax.set_yscale("log")
ax.legend()
fig.tight_layout()
return fig
def plot_rollout_total_energy(obs: RolloutObservables, bins: int = 50, reference=None):
"""Distribution of total deposited energy per shower."""
fig, ax = plt.subplots(figsize=(7, 4))
ax.hist(
obs.event_table["total_edep"],
bins=bins,
histtype="step",
label="rollout",
color="C0",
)
if reference is not None:
ax.hist(
reference.event_table["real_total_edep"].to_numpy(),
bins=bins,
histtype="step",
label="real",
color="k",
ls="--",
)
ax.set_xlabel("total E_dep / event [MeV]")
ax.set_ylabel("events")
ax.set_title("Total deposited energy")
ax.legend()
fig.tight_layout()
return fig
# ── Rollout vs. held-out truth (Tier 1-3, unpaired) ────────────────────────────
#
# `compute_rollout_observables` above covers Tier 4 (event-level) comparisons.
# This builds a `SampleCollection` instead, so the Tier 1-3 diagnostics
# (plot_marginals, plot_kl_bars, correlation_matrices, plot_pairwise,
# direction_alignment, constraint_report) also work on a rollout: `giant
# rollout` output and a training-input-schema truth file (see
# `giant.data.loader.load_steps`) both carry pre_*/post_*/edep/step_length in
# the same physical, world-frame units, so both sides decode into
# RAW_TARGET_NAMES space via the same local_frame_rotation/travel_direction
# construction `giant.data.transforms.build_features` uses for `target_s1` —
# just skipping the log/ALR encode step, since neither file needs it decoded.
#
# Unlike `load_predicted_local`, the two files are independent rather than
# row-for-row paired (a rollout doesn't replay real events step-by-step), so
# real_raw/gen_raw may have different lengths; `SampleCollection`'s `*_gen`
# fields carry the rollout side's own pdg/material/conditioning for grouping.
_WORLD_FRAME_STEP_COLS = [
"pdg",
"material",
"layer_id",
"pre_x",
"pre_y",
"pre_z",
"pre_E",
"pre_dx",
"pre_dy",
"pre_dz",
"post_x",
"post_y",
"post_z",
"post_E",
"post_dx",
"post_dy",
"post_dz",
"edep",
"step_length",
]
def _world_frame_raw_targets(batch: pl.DataFrame) -> np.ndarray:
"""(N, 9) RAW_TARGET_NAMES array from a world-frame steps batch.
See the module note above `_WORLD_FRAME_STEP_COLS` — both `load_rollout_vs_truth`
inputs share this schema.
"""
pre_pos = batch.select(["pre_x", "pre_y", "pre_z"]).to_numpy().astype(np.float32)
pre_dir = batch.select(["pre_dx", "pre_dy", "pre_dz"]).to_numpy().astype(np.float32)
post_pos = (
batch.select(["post_x", "post_y", "post_z"]).to_numpy().astype(np.float32)
)
post_dir = (
batch.select(["post_dx", "post_dy", "post_dz"]).to_numpy().astype(np.float32)
)
pre_E = batch["pre_E"].to_numpy().astype(np.float32)
post_E = batch["post_E"].to_numpy().astype(np.float32)
post_dir_local = local_frame_rotation(pre_dir, post_dir)
travel_dir_local = local_frame_rotation(
pre_dir, travel_direction(pre_pos, post_pos)
)
return np.column_stack(
[
batch["step_length"].to_numpy().astype(np.float32),
pre_E - post_E,
batch["edep"].to_numpy().astype(np.float32),
post_dir_local,
travel_dir_local,
]
).astype(np.float32)
def _world_frame_cond(batch: pl.DataFrame) -> np.ndarray:
"""(N, 9) `_COND_CONT_COLS`-layout array from a world-frame steps batch.
`n_sec` comes from `child_track_ids` (truth schema) or `n_sec_pred`
(rollout schema) — whichever the batch has; kept only for shape parity
with `load_predicted_local`'s `cond_cont_raw` — nothing downstream in this
module groups by it, only `pre_E` (index 3, for `group_by="energy"`).
"""
if "n_sec_pred" in batch.columns:
n_sec = batch["n_sec_pred"].to_numpy().astype(np.float32)
elif "child_track_ids" in batch.columns:
n_sec = batch["child_track_ids"].list.len().to_numpy().astype(np.float32)
else:
n_sec = np.zeros(batch.height, dtype=np.float32)
return np.column_stack(
[
batch.select(["pre_x", "pre_y", "pre_z"]).to_numpy(),
batch["pre_E"].to_numpy(),
batch.select(["pre_dx", "pre_dy", "pre_dz"]).to_numpy(),
batch["layer_id"].to_numpy().astype(np.float32),
n_sec,
]
).astype(np.float32)
def _existing_columns(
source: str | Path | pl.LazyFrame, wanted: list[str]
) -> list[str]:
if isinstance(source, pl.LazyFrame):
available = set(source.collect_schema().names())
else:
available = set(pq.ParquetFile(Path(source)).schema_arrow.names)
return [c for c in wanted if c in available]
# `rollout.py`'s `_terminal_rows` writes one synthetic bookkeeping row per track
# for these termination reasons (escaped/unknown_pdg/energy_cutoff/max_steps):
# step_length=0, post_pos=pre_pos, and — for every reason but escaped — the
# track's *entire remaining pre_E* dumped into `edep` in one row, so the shower's
# total energy still conserves. These aren't steps in any physical sense (truth
# data has no equivalent), so mixing them into a per-step real-vs-generated
# comparison would inject a spurious step_length=0 spike and roughly double the
# apparent mean edep purely from bookkeeping, not model behavior. Real generated
# steps carry "" (continuing) or `TERM_NATURAL_END` (the track's last real step,
# which does have genuine step_length/edep) and are kept.
_SYNTHETIC_ROLLOUT_TERMINATION_REASONS = frozenset(
{TERM_ESCAPED, TERM_UNKNOWN_PDG, TERM_ENERGY_CUTOFF, TERM_MAX_STEPS}
)
def _load_world_frame_side(
source: str | Path | pl.LazyFrame,
sample_frac: float,
seed: int,
batch_size: int,
extra_cols: list[str],
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Stream one world-frame steps file/LazyFrame into (raw9, cond9, pdg, material).
`extra_cols` (e.g. `["n_sec_pred", "termination_reason"]` or
`["child_track_ids"]`) are included only when present, so the same helper
serves both the rollout and truth schemas without either needing the
other's columns. When `termination_reason` is present (the rollout side),
rows carrying one of `_SYNTHETIC_ROLLOUT_TERMINATION_REASONS` are dropped
before decoding — see that constant's docstring for why.
"""
columns = _WORLD_FRAME_STEP_COLS + _existing_columns(source, extra_cols)
threshold = int(sample_frac * 2**32) if sample_frac < 1.0 else None
raw_parts, cond_parts, pdg_parts, mat_parts = [], [], [], []
offset = 0
for batch in _iter_predicted_local_batches(source, columns, batch_size):
n = batch.height
if threshold is not None:
row_idx = pl.arange(offset, offset + n, eager=True).cast(pl.UInt32)
batch = batch.filter((row_idx.hash(seed=seed) % 2**32) < threshold)
offset += n
if "termination_reason" in batch.columns:
batch = batch.filter(
~pl.col("termination_reason").is_in(
list(_SYNTHETIC_ROLLOUT_TERMINATION_REASONS)
)
)
if batch.height == 0:
continue
raw_parts.append(_world_frame_raw_targets(batch))
cond_parts.append(_world_frame_cond(batch))
pdg_parts.append(batch["pdg"].to_numpy())
mat_parts.append(batch["material"].to_numpy())
if not raw_parts:
raise ValueError(f"{source}: no rows survived (sample_frac={sample_frac})")
return (
np.concatenate(raw_parts, axis=0),
np.concatenate(cond_parts, axis=0),
np.concatenate(pdg_parts, axis=0),
np.concatenate(mat_parts, axis=0),
)
def load_rollout_vs_truth(
rollout_path: str | Path | pl.LazyFrame,
truth_path: str | Path | pl.LazyFrame,
sample_frac: float = 1.0,
seed: int = 0,
batch_size: int = 1_000_000,
) -> SampleCollection:
"""Build a `SampleCollection` comparing a `giant rollout` shower to a truth file.
`truth_path` — any file sharing `giant train`'s input schema (real
miniCaloSim steps, e.g. a held-out/val parquet) — is treated as "real";
`rollout_path` (`giant rollout` output) is treated as "generated". The two
are independent files (not row-for-row paired, since a rollout doesn't
replay real events step-by-step): `real_raw`/`gen_raw` may have different
lengths, and `pdg`/`material`/`cond_cont_raw` are computed separately per
side (`SampleCollection`'s `*_gen` fields) — `_group_labels` (used by
`marginal_table`/`plot_marginals`/`plot_kl_bars`) builds independent masks
for each. `correlation_matrices`, `direction_alignment`, and
`constraint_report` never paired real/gen row-for-row to begin with, so
they need no special handling here.
Every raw target dim is decoded from the world-frame `pre_*`/`post_*`/
`edep`/`step_length` columns both files share — see the module note above
`_WORLD_FRAME_STEP_COLS`. `sample_frac`/`seed`/`batch_size` behave as in
`load_predicted_local`, applied independently to each file.
The rollout side drops synthetic termination-bookkeeping rows (see
`_SYNTHETIC_ROLLOUT_TERMINATION_REASONS`) before decoding, since those
aren't real generated steps. One remaining case where `edep` still isn't
perfectly analogous between the two files: `decode_secondaries` rescales
the valid secondary slots to sum to exactly `e_sec` whenever `n_sec > 0`
(see that function's docstring), but when Stage 1 predicts a nonzero
`e_sec` while Stage 2's `n_sec` head predicts 0 secondaries, there's no
slot to carry that budget at all — `rollout.py` deposits it into that
step's `edep` instead, which truth's Geant4-recorded `edep` never does.
That specific disagreement between the two Stage-1/Stage-2 heads is rare
but not otherwise fixable at decode time.
"""
if not (0 < sample_frac <= 1):
raise ValueError(f"sample_frac must be in (0, 1], got {sample_frac}")
if not isinstance(rollout_path, pl.LazyFrame):
_check_rollout_metadata(Path(rollout_path))
gen_raw, gen_cond, gen_pdg, gen_material = _load_world_frame_side(
rollout_path,
sample_frac,
seed,
batch_size,
extra_cols=["n_sec_pred", "termination_reason"],
)
real_raw, real_cond, real_pdg, real_material = _load_world_frame_side(
truth_path, sample_frac, seed, batch_size, extra_cols=["child_track_ids"]
)
return SampleCollection(
cond_cont_raw=real_cond,
pdg=real_pdg,
material=real_material,
real_raw=real_raw,
gen_raw=gen_raw,
cond_cont_raw_gen=gen_cond,
pdg_gen=gen_pdg,
material_gen=gen_material,
)
giant/cli.py
+287 -10
View File
@@ -21,6 +21,7 @@ from giant.constants import (
PREDICT_COORD_METADATA_KEY,
PREDICT_SCHEMA_VERSION,
PREDICT_SCHEMA_VERSION_KEY,
ROLLOUT_COORD_VALUE,
)
from giant.data.loader import (
find_parquet_files,
@@ -30,15 +31,36 @@ from giant.data.loader import (
from giant.data.transforms import (
build_features,
build_cond_features,
decode_secondaries,
energy_simplex_decode,
inv_local_frame_rotation,
inv_log_transform,
reconstruct_post_pos,
Normalizer,
)
from giant.model.network import DenoisingMLP
from giant.geometry import GeometryOracle
from giant.model.network import DenoisingMLP, SecondaryDecoder
from giant.pipeline import run_train_job
from giant.sample import sample_flow
from giant.rollout import rollout as run_rollout
from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx
_STAGE1_MODEL_KEYS = {
"pdg_vocab",
"mat_vocab",
"hidden_dim",
"n_blocks",
"emb_dim",
"dropout",
"k_max",
}
_SEC_DECODER_MODEL_KEYS = {
"pdg_vocab",
"mat_vocab",
"hidden_dim",
"n_blocks",
"emb_dim",
"dropout",
}
app = typer.Typer(no_args_is_help=True)
@@ -354,6 +376,13 @@ def predict(
)
raise typer.Exit(1)
if "sec_decoder" not in ckpt:
typer.echo(
"error: checkpoint has no sec_decoder — retrain with the current code",
err=True,
)
raise typer.Exit(1)
model_cfg = ckpt["model_config"]
if batch_size_auto:
@@ -370,16 +399,27 @@ def predict(
typer.echo(
f"batch_size: {batch_size_value} (auto-estimated from free GPU memory)"
)
assert batch_size_value is not None
bs = batch_size_value
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"])
model = DenoisingMLP(**model_cfg)
model = DenoisingMLP(
**{k: v for k, v in model_cfg.items() if k in _STAGE1_MODEL_KEYS}
)
model.load_state_dict(ckpt["model"])
model.to(_device).eval()
sec_decoder = SecondaryDecoder(
**{k: v for k, v in model_cfg.items() if k in _SEC_DECODER_MODEL_KEYS}
)
sec_decoder.load_state_dict(ckpt["sec_decoder"])
sec_decoder.to(_device).eval()
typer.echo(f"loaded checkpoint: {checkpoint}")
gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
@@ -408,7 +448,7 @@ def predict(
nonlocal writer, total
if coord == Coord.local:
cond_cont, cond_cat, target_raw, _, _ = build_features(
cond_cont, cond_cat, target_raw, _, _, _, _, _ = build_features(
piece, pdg_map, mat_map
)
cond_cont = cond_norm.transform(cond_cont)
@@ -419,7 +459,20 @@ def predict(
cc = torch.from_numpy(cond_cont).float().to(_device)
ck = torch.from_numpy(cond_cat).long().to(_device)
pred = sample_flow(model, cc, ck, steps=steps).cpu().numpy() # normalised
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_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()
n_sec_pred_np = n_sec_pred.cpu().numpy()
pred = stage1_norm.cpu().numpy() # normalised
# Inverse-normalise → local frame, log-scaled scalars
raw = tgt_norm.inverse_transform(pred)
@@ -453,8 +506,10 @@ def predict(
step_length = inv_log_transform(raw[:, 0])
# Columns 1:3 are ALR coords of the deposit/secondary/post energy
# simplex; decode them against pre_E so edep + e_sec + post_E == pre_E
# (hence delta_e == edep + e_sec) holds by construction.
edep, _e_sec, _post_E, delta_e = energy_simplex_decode(
# (hence delta_e == edep + e_sec) holds by construction. e_sec_pred
# doubles as the stick-breaking energy budget for the Stage-2 decode
# below, since the model has no other source for it at inference.
edep, e_sec_pred, _post_E, delta_e = energy_simplex_decode(
raw[:, 1:3], piece["pre_E"]
)
@@ -473,6 +528,28 @@ 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,
n_sec_pred_np,
e_sec_pred,
piece["pre_dir"],
pdg_map_inv,
)
sec_pdg_list = [
sec_pdg_code[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)
]
sec_E_list = [sec_E[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)]
sec_dx_list = [
sec_dir_world[i, :n, 0].tolist() for i, n in enumerate(n_sec_pred_np)
]
sec_dy_list = [
sec_dir_world[i, :n, 1].tolist() for i, n in enumerate(n_sec_pred_np)
]
sec_dz_list = [
sec_dir_world[i, :n, 2].tolist() for i, n in enumerate(n_sec_pred_np)
]
table = pa.table(
{
"event_id": piece["event_id"],
@@ -487,6 +564,7 @@ def predict(
"material": piece["material"],
"layer_id": piece["layer_id"],
"n_sec": piece["n_sec"],
"n_sec_pred": n_sec_pred_np,
"step_length": step_length,
"delta_e": delta_e,
"edep": edep,
@@ -496,6 +574,11 @@ def predict(
"post_x": post_pos_world[:, 0],
"post_y": post_pos_world[:, 1],
"post_z": post_pos_world[:, 2],
"sec_pdg_list": sec_pdg_list,
"sec_E_list": sec_E_list,
"sec_dx_list": sec_dx_list,
"sec_dy_list": sec_dy_list,
"sec_dz_list": sec_dz_list,
}
)
@@ -543,9 +626,7 @@ def predict(
if writer is not None:
writer.close()
ref_path = _write_prediction_ref(
checkpoint, pred_uuid, out, dataset_path, comment
)
ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path, comment)
typer.echo(f"reference: {ref_path}")
if skipped:
@@ -559,5 +640,201 @@ def predict(
typer.echo(f"wrote {total:,} rows → {out}")
def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.ndarray]:
"""Pick each event's primary entry state (argmax-pre_E row) as a shower seed.
Streams conditioning columns and keeps the highest-pre_E step per event_id —
the codebase's convention for the primary (a secondary always carries less
energy than its parent). See giant/analysis.py:_entry_axis_and_bin_edges.
"""
best_E: dict[int, float] = {}
best: dict[int, tuple] = {}
for path in files:
for chunk in iter_cond_chunks(path):
ev = chunk["event_id"]
pe = chunk["pre_E"]
for i in range(len(ev)):
e = int(ev[i])
if pe[i] > best_E.get(e, -np.inf):
best_E[e] = float(pe[i])
best[e] = (
int(chunk["pdg"][i]),
chunk["pre_pos"][i].astype(np.float64),
float(pe[i]),
chunk["pre_dir"][i].astype(np.float64),
)
event_ids = sorted(best)
if n_events is not None:
event_ids = event_ids[:n_events]
if not event_ids:
raise ValueError("no events found to seed from")
return {
"event_id": np.array(event_ids, dtype=np.int64),
"pdg": np.array([best[e][0] for e in event_ids], dtype=np.int64),
"pre_pos": np.stack([best[e][1] for e in event_ids]),
"pre_E": np.array([best[e][2] for e in event_ids], dtype=np.float64),
"pre_dir": np.stack([best[e][3] for e in event_ids]),
}
@app.command()
def rollout(
data: Annotated[
Path, typer.Argument(help="Parquet file/dir to seed showers from (real events)")
],
checkpoint: Annotated[
Path,
typer.Option("--checkpoint", "-c", help="Checkpoint .pt (best.pt/last.pt)"),
],
geometry: Annotated[
Path,
typer.Option(
"--geometry",
"-g",
help="Geometry oracle .pkl (dwarf build-geometry-oracle)",
),
],
energy_cutoff: Annotated[
float,
typer.Option(
"--energy-cutoff",
help="Stop a track when its energy drops below this [MeV]",
),
] = 0.1,
max_steps: Annotated[
int, typer.Option("--max-steps", help="Max steps per individual track")
] = 1000,
steps: Annotated[
int,
typer.Option("--steps", "-s", help="Flow matching ODE steps per model call"),
] = 10,
batch_size: Annotated[
int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward")
] = 4096,
max_tracks_per_event: Annotated[
Optional[int],
typer.Option(
"--max-tracks-per-event",
help="Safety cap on tracks per shower (sub-cap secondaries deposit in place)",
),
] = None,
escape_threshold: Annotated[
Optional[float],
typer.Option(
"--escape-threshold",
help="Override the oracle's NN-distance escape threshold [mm]",
),
] = None,
n_events: Annotated[
Optional[int], typer.Option("--n-events", help="Cap number of seed events")
] = None,
device: Annotated[
Optional[str], typer.Option("--device", "-d", help="cpu | cuda | mps (auto)")
] = None,
out: Annotated[
Optional[Path], typer.Option("--out", "-o", help="Output steps parquet")
] = None,
seed: Annotated[
Optional[int],
typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
] = None,
) -> None:
"""Roll the surrogate forward into full showers (autoregressive)."""
if seed is not None:
torch.manual_seed(seed)
np.random.seed(seed)
_device = torch.device(device) if device else gconfig.auto_device()
typer.echo(f"device: {_device}")
ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
for key in ("model_config", "sec_decoder"):
if key not in ckpt:
typer.echo(
f"error: checkpoint has no {key} — retrain with the current code",
err=True,
)
raise typer.Exit(1)
model_cfg = ckpt["model_config"]
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"])
model = DenoisingMLP(
**{k: v for k, v in model_cfg.items() if k in _STAGE1_MODEL_KEYS}
)
model.load_state_dict(ckpt["model"])
model.to(_device).eval()
sec_decoder = SecondaryDecoder(
**{k: v for k, v in model_cfg.items() if k in _SEC_DECODER_MODEL_KEYS}
)
sec_decoder.load_state_dict(ckpt["sec_decoder"])
sec_decoder.to(_device).eval()
typer.echo(f"loaded checkpoint: {checkpoint}")
oracle = GeometryOracle.load(geometry)
typer.echo(
f"loaded geometry oracle: {geometry} "
f"(escape_threshold={oracle.escape_threshold:.3f})"
)
files = find_parquet_files(data)
seeds = _seed_from_data(files, n_events)
typer.echo(f"seeded {len(seeds['event_id']):,} shower(s)")
records = run_rollout(
model,
sec_decoder,
oracle,
seeds,
cond_norm,
tgt_norm,
pdg_map,
mat_map,
energy_cutoff=energy_cutoff,
max_steps=max_steps,
steps=steps,
batch_size=batch_size,
device=_device,
max_tracks_per_event=max_tracks_per_event,
escape_threshold=escape_threshold,
)
out, dataset_path, pred_uuid = _resolve_prediction_output(data, out)
out.parent.mkdir(parents=True, exist_ok=True)
table = pa.table(records).replace_schema_metadata(
{
PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE,
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
}
)
pq.write_table(table, out)
ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path)
ref = yaml.safe_load(ref_path.read_text())
ref.update(
{
"kind": "rollout",
"geometry_oracle": str(geometry.resolve()),
"energy_cutoff": energy_cutoff,
"max_steps": max_steps,
"steps": steps,
"max_tracks_per_event": max_tracks_per_event,
"n_seed_events": int(len(seeds["event_id"])),
}
)
ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False))
n_rows = len(records["event_id"])
reasons = Counter(r for r in records["termination_reason"].tolist() if r)
typer.echo(f"wrote {n_rows:,} step rows → {out}")
typer.echo(f"terminations: {dict(reasons)}")
typer.echo(f"reference: {ref_path}")
if __name__ == "__main__":
app()
giant/config.py
+11 -8
View File
@@ -20,6 +20,8 @@ DEFAULT_CONFIG: dict = {
"validate_every": 10,
"validate_steps": 10,
"warmup_epochs": 5,
"lambda_nsec": 0.1,
"lambda_s2": 1.0,
},
"model": {
"hidden_dim": 256,
@@ -51,14 +53,15 @@ def auto_device() -> torch.device:
return torch.device("cpu")
# Calibration point for estimate_batch_size(training=True): hidden_dim=512,
# n_blocks=6, batch_size=131072 measured at ~8 GiB VRAM. Activation memory is
# assumed to scale linearly with batch_size * hidden_dim * n_blocks (the
# ResBlock stack dominates), so this is a rough estimate rather than a
# guaranteed bound.
_REF_BYTES = 8 * 1024**3
_REF_BATCH_SIZE = 131072
_REF_HIDDEN_DIM = 512
# Calibration point for estimate_batch_size(training=True): hidden_dim=1024,
# n_blocks=6, batch_size=29696 measured at ~7683 MiB VRAM (post-Phase-2
# architecture, including the Stage-2 secondary decoder and n_sec head).
# Activation memory is assumed to scale linearly with
# batch_size * hidden_dim * n_blocks (the ResBlock stack dominates), so this
# is a rough estimate rather than a guaranteed bound.
_REF_BYTES = 7683 * 1024**2
_REF_BATCH_SIZE = 29696
_REF_HIDDEN_DIM = 1024
_REF_N_BLOCKS = 6
# Calibration point for estimate_batch_size(training=False): inference has no
giant/constants.py
+35 -10
View File
@@ -1,16 +1,27 @@
X_DIM = 9
# Conditioning continuous-feature width: pre_pos(3), log(pre_E)(1), pre_dir(3),
# layer_id(1), n_sec(1), log(e_sec)(1). One wider than X_DIM because e_sec
# (secondary energy) is a conditioning input in the energy-conservation PoC.
COND_DIM = 10
# 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
# The two energy columns are additive-log-ratio (ALR) coordinates of the
# deposit/secondary/post energy simplex (fractions of pre_E that sum to 1),
# referenced to the post-energy fraction — see giant.data.transforms
# .energy_simplex_encode/.energy_simplex_decode. They replace the former
# independent log_delta_e / log_edep targets so energy conservation holds by
# construction after decoding.
# 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 (non-embedding) width: stick-breaking logit + local dir.
CONT_SLOT_DIM = SEC_SLOT_DIM - EMB_DIM # 4
# Flattened Stage-2 target dimension
SEC_DIM = K_MAX * SEC_SLOT_DIM # 15 * 20 = 300
# Stage-1 9D target names (unchanged from energy-conservation PoC)
LOCAL_TARGET_NAMES = [
"log_step_length",
"edep_logit",
@@ -29,3 +40,17 @@ LOCAL_TARGET_NAMES = [
PREDICT_COORD_METADATA_KEY = "giant.predict.coord"
PREDICT_SCHEMA_VERSION_KEY = "giant.predict.schema_version"
PREDICT_SCHEMA_VERSION = "2"
# Coord-metadata value tagging a `giant rollout` steps parquet (world frame,
# autoregressive shower output). Distinct from predict's "global"/"local".
ROLLOUT_COORD_VALUE = "rollout"
# Per-track termination reasons recorded by the rollout driver. "escaped"
# energy is treated as leakage (not deposited); every other stop deposits the
# track's remaining energy locally so total energy is conserved.
TERM_ESCAPED = "escaped"
TERM_ENERGY_CUTOFF = "energy_cutoff"
TERM_MAX_STEPS = "max_steps"
TERM_NATURAL_END = "natural_end"
TERM_UNKNOWN_PDG = "unknown_pdg"
TERM_MAX_TRACKS = "max_tracks"
giant/data/dataset.py
+68 -60
View File
@@ -4,53 +4,12 @@ from pathlib import Path
import numpy as np
import torch
from torch.utils.data import Dataset, IterableDataset
from torch.utils.data import IterableDataset
from giant.data.loader import iter_file_chunks
from giant.data.transforms import Normalizer, build_features
class StepsDataset(Dataset):
def __init__(
self,
cond_cont: np.ndarray,
cond_cat: np.ndarray,
target: np.ndarray,
) -> None:
self.cond_cont = torch.from_numpy(cond_cont).float()
self.cond_cat = torch.from_numpy(cond_cat).long()
self.target = torch.from_numpy(target).float()
def __len__(self) -> int:
return len(self.target)
def __getitem__(self, index):
return self.cond_cont[index], self.cond_cat[index], self.target[index]
def train_val_split(
data: dict,
cond_cont: np.ndarray,
cond_cat: np.ndarray,
target: np.ndarray,
val_fraction: float = 0.1,
seed: int = 42,
) -> tuple[StepsDataset, StepsDataset]:
rng = np.random.default_rng(seed)
unique_events = np.unique(data["event_id"])
rng.shuffle(unique_events)
n_val = max(1, int(len(unique_events) * val_fraction))
val_events = set(unique_events[:n_val].tolist())
val_mask = np.array([e in val_events for e in data["event_id"]])
train_mask = ~val_mask
return (
StepsDataset(cond_cont[train_mask], cond_cat[train_mask], target[train_mask]),
StepsDataset(cond_cont[val_mask], cond_cat[val_mask], target[val_mask]),
)
def make_event_split(
all_event_ids: np.ndarray,
val_fraction: float = 0.1,
@@ -75,6 +34,16 @@ class StreamingStepsDataset(IterableDataset):
Yields whole batches (use with `DataLoader(..., batch_size=None)`)
rather than single rows, so the batch is assembled with vectorized
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)
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
"""
def __init__(
@@ -91,13 +60,12 @@ class StreamingStepsDataset(IterableDataset):
) -> None:
self.files = list(files)
self.split_events = split_events
self._events_arr = np.array(sorted(split_events)) # for np.isin
self._events_arr = np.array(sorted(split_events))
self.pdg_map = pdg_map
self.mat_map = mat_map
self.cond_normalizer = cond_normalizer
self.target_normalizer = target_normalizer
self.batch_size = batch_size
# Buffer must hold at least one batch or we could never emit one.
self.shuffle_buffer = max(shuffle_buffer, batch_size)
self.shuffle = shuffle
@@ -114,6 +82,9 @@ class StreamingStepsDataset(IterableDataset):
buf_cont: list[np.ndarray] = []
buf_cat: list[np.ndarray] = []
buf_tgt: list[np.ndarray] = []
buf_nsec: list[np.ndarray] = []
buf_sec: list[np.ndarray] = []
buf_spdg: list[np.ndarray] = []
buf_n = 0
for path in files:
@@ -123,43 +94,69 @@ class StreamingStepsDataset(IterableDataset):
continue
chunk = {k: v[mask] for k, v in chunk.items()}
cond_cont, cond_cat, target, _, _ = build_features(
chunk,
self.pdg_map,
self.mat_map,
cond_normalizer=self.cond_normalizer,
target_normalizer=self.target_normalizer,
cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, _, _ = (
build_features(
chunk,
self.pdg_map,
self.mat_map,
cond_normalizer=self.cond_normalizer,
target_normalizer=self.target_normalizer,
require_secondaries=True,
)
)
buf_cont.append(cond_cont)
buf_cat.append(cond_cat)
buf_tgt.append(target)
buf_tgt.append(target_s1)
buf_nsec.append(n_sec)
buf_sec.append(sec_cont)
buf_spdg.append(sec_pdg_idx)
buf_n += len(cond_cont)
if buf_n >= self.shuffle_buffer:
buf_cont, buf_cat, buf_tgt, buf_n = yield from self._flush(
buf_cont, buf_cat, buf_tgt, final=False
(
buf_cont,
buf_cat,
buf_tgt,
buf_nsec,
buf_sec,
buf_spdg,
buf_n,
) = yield from self._flush(
buf_cont,
buf_cat,
buf_tgt,
buf_nsec,
buf_sec,
buf_spdg,
final=False,
)
if buf_n > 0:
yield from self._flush(buf_cont, buf_cat, buf_tgt, final=True)
yield from self._flush(
buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, final=True
)
def _flush(
self,
buf_cont: list[np.ndarray],
buf_cat: list[np.ndarray],
buf_tgt: list[np.ndarray],
buf_nsec: list[np.ndarray],
buf_sec: list[np.ndarray],
buf_spdg: list[np.ndarray],
final: bool,
):
"""Yield full batches of `batch_size`; carry any remainder back to the caller.
All batching is done via vectorized numpy slicing (no per-row Python loop).
"""
cont = np.concatenate(buf_cont)
cat = np.concatenate(buf_cat)
tgt = np.concatenate(buf_tgt)
nsec = np.concatenate(buf_nsec)
sec = np.concatenate(buf_sec)
spdg = np.concatenate(buf_spdg)
if self.shuffle:
idx = np.random.permutation(len(cont))
cont, cat, tgt = cont[idx], cat[idx], tgt[idx]
nsec, sec, spdg = nsec[idx], sec[idx], spdg[idx]
bs = self.batch_size
n = len(cont)
@@ -170,9 +167,20 @@ class StreamingStepsDataset(IterableDataset):
torch.from_numpy(cont[start:end]).float(),
torch.from_numpy(cat[start:end]).long(),
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(),
)
if final:
return [], [], [], 0
return [], [], [], [], [], [], 0
rem = n_full * bs
return [cont[rem:]], [cat[rem:]], [tgt[rem:]], n - rem
return (
[cont[rem:]],
[cat[rem:]],
[tgt[rem:]],
[nsec[rem:]],
[sec[rem:]],
[spdg[rem:]],
n - rem,
)
giant/data/loader.py
+52 -1
View File
@@ -40,8 +40,50 @@ def find_parquet_files(path: str | Path) -> list[Path]:
return [p]
def _pad_list_col(series: pd.Series, K: int, fill: float = 0.0) -> np.ndarray:
"""Pad / truncate a list-valued Series to fixed width K → (N, K) float32."""
out = np.full((len(series), K), fill, dtype=np.float32)
for i, lst in enumerate(series):
if lst is not None and len(lst) > 0:
n = min(len(lst), K)
out[i, :n] = lst[:n]
return out
def _pad_list_col_int(series: pd.Series, K: int, fill: int = 0) -> np.ndarray:
"""Pad / truncate a list-valued integer Series to fixed width K → (N, K) int64."""
out = np.full((len(series), K), fill, dtype=np.int64)
for i, lst in enumerate(series):
if lst is not None and len(lst) > 0:
n = min(len(lst), K)
out[i, :n] = lst[:n]
return out
def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndarray:
"""Pad three list-valued direction columns → (N, K, 3) float32.
Padding direction defaults to (0,0,1) (forward) so it is a valid unit vector.
"""
N = len(dx)
out = np.zeros((N, K, 3), dtype=np.float32)
out[:, :, 2] = 1.0
for i in range(N):
lx, ly, lz = dx.iloc[i], dy.iloc[i], dz.iloc[i]
if lx is not None and len(lx) > 0:
n = min(len(lx), K)
out[i, :n, 0] = lx[:n]
out[i, :n, 1] = ly[:n]
out[i, :n, 2] = lz[:n]
return out
def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
return {
from giant.constants import K_MAX
has_sec_lists = "sec_E_list" in df.columns
d: dict[str, np.ndarray] = {
"event_id": df["event_id"].to_numpy(),
"pdg": df["pdg"].to_numpy(dtype=np.int32),
"pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
@@ -59,6 +101,15 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
"post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32),
}
if has_sec_lists:
d["sec_E_list"] = _pad_list_col(df["sec_E_list"], K_MAX)
d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], K_MAX)
d["sec_dir_list"] = _pad_dir_col(
df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], K_MAX
)
return d
def load_steps(path: str | Path) -> dict[str, np.ndarray]:
return _df_to_dict(pd.read_parquet(path))
giant/data/transforms.py
+230 -13
View File
@@ -280,6 +280,146 @@ def inv_local_frame_rotation(
)
_STICK_LOGIT_CLIP = 10.0 # logit value used for the last valid secondary slot
def encode_secondaries(
sec_E_list: np.ndarray,
sec_dir_list: np.ndarray,
sec_valid: np.ndarray,
e_sec: np.ndarray,
pre_dir: np.ndarray,
) -> 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]
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.
"""
N, K = sec_E_list.shape
e_sec = np.asarray(e_sec, dtype=np.float64)
stick_logits = np.zeros((N, K), dtype=np.float32)
for i in range(K):
if i == 0:
remaining = np.maximum(e_sec, _EPS)
else:
remaining = np.maximum(e_sec - sec_E_list[:, :i].sum(axis=1), _EPS)
f = np.clip(sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS)
logit = np.log(f / (1.0 - f)).astype(np.float32)
# Last valid slot: give it the full remaining budget
is_last = sec_valid[:, i] & ~(
sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool)
)
logit = np.where(is_last, _STICK_LOGIT_CLIP, logit)
logit = np.where(
sec_valid[:, i], np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP), 0.0
)
stick_logits[:, i] = logit.astype(np.float32)
# Rotate each slot's direction into the local frame of the primary.
# pre_dir is broadcast across all K slots.
dir_local = np.zeros((N, K, 3), dtype=np.float32)
for i in range(K):
# Only rotate valid slots; leave padded slots as (0,0,1) or whatever.
valid_mask = sec_valid[:, i]
if valid_mask.any():
dir_local[valid_mask, i] = local_frame_rotation(
pre_dir[valid_mask], sec_dir_list[valid_mask, i]
)
sec_cont = np.concatenate(
[stick_logits[:, :, None], dir_local], axis=-1
) # (N, K, 4)
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]:
"""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)
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.
"""
N, K, _ = sec_cont.shape
stick_logits = sec_cont[:, :, 0] # (N, K)
dir_local = sec_cont[:, :, 1:].copy() # (N, K, 3)
# Flow-matching output isn't guaranteed unit norm; normalise before the
# rotation below, which preserves magnitude rather than fixing it up.
norms = np.linalg.norm(dir_local, axis=-1, keepdims=True)
dir_local /= np.where(norms < 1e-8, 1.0, norms)
fractions = 1.0 / (1.0 + np.exp(-stick_logits.astype(np.float64)))
sec_E = np.zeros((N, K), dtype=np.float64)
e_sec = np.asarray(e_sec, dtype=np.float64)
remaining = e_sec.copy()
for i in range(K):
sec_E[:, i] = fractions[:, i] * remaining
remaining = np.maximum(remaining - sec_E[:, i], 0.0)
sec_valid = np.arange(K)[None, :] < n_sec[:, None] # (N, K)
# Stick-breaking guarantees sum(sec_E[valid]) <= e_sec (each fraction is in
# [0,1] of an already-shrinking remainder) but rarely hits it exactly, so
# rescale the valid slots by one common per-row factor to close that gap —
# rather than dumping the shortfall into whichever slot happens to be last
# by energy rank, which would let one low-energy secondary balloon and
# distort the shower's topology. This preserves each row's relative split
# across its secondaries and only ever scales up (valid_sum <= e_sec).
# Rows where every valid slot decoded to ~zero (scale undefined) fall back
# to an even split of e_sec across the n_sec valid slots.
sec_E = sec_E * sec_valid
valid_sum = sec_E.sum(axis=1)
degenerate = (valid_sum <= _EPS) & (n_sec > 0)
scale = np.where(valid_sum > _EPS, e_sec / np.maximum(valid_sum, _EPS), 0.0)
sec_E = sec_E * scale[:, None]
even_share = e_sec / np.maximum(n_sec, 1).astype(np.float64)
sec_E = np.where(degenerate[:, None] & sec_valid, even_share[:, None], sec_E)
sec_E = sec_E.astype(np.float32)
sec_dir_world = np.zeros((N, K, 3), dtype=np.float32)
for i in range(K):
valid = sec_valid[:, i]
if valid.any():
sec_dir_world[valid, i] = inv_local_frame_rotation(
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,
)
return sec_E, sec_dir_world, sec_pdg_code, sec_valid
def build_cond_features(
data: dict[str, np.ndarray],
pdg_map: dict[int, int],
@@ -293,8 +433,6 @@ def build_cond_features(
log_transform(data["pre_E"]),
data["pre_dir"],
data["layer_id"].astype(np.float32),
data["n_sec"].astype(np.float32),
log_transform(data["e_sec"]),
]
).astype(np.float32)
@@ -315,11 +453,32 @@ def build_features(
cond_normalizer: Normalizer | None = None,
target_normalizer: Normalizer | None = None,
fit: bool = False,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, Normalizer | None, Normalizer | None]:
"""Assemble (cond_cont, cond_cat, target) arrays ready for StepsDataset.
require_secondaries: bool = False,
) -> tuple[
np.ndarray,
np.ndarray,
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) arrays.
When fit=True, new Normalizers are fitted on the supplied 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
require_secondaries: when True, raise if any step has n_sec > 0 but the
per-secondary list columns are absent (a mis-converted file that would
otherwise silently zero all Stage-2 targets). Training paths set this;
Stage-1-only callers (e.g. `giant predict`) leave it False.
"""
from giant.constants import K_MAX
post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])
travel_dir_local = local_frame_rotation(
data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"])
@@ -327,9 +486,9 @@ def build_features(
energy_z = energy_simplex_encode(
data["edep"], data["e_sec"], data["post_E"], data["pre_E"]
) # (N, 2): ALR coords of the deposit/secondary/post energy simplex
) # (N, 2)
target = np.column_stack(
target_s1 = np.column_stack(
[
log_transform(data["step_length"]),
energy_z,
@@ -338,28 +497,86 @@ def build_features(
]
).astype(np.float32) # (N, 9)
# Phase 2: conditioning drops n_sec and log(e_sec)
cond_cont = np.column_stack(
[
data["pre_pos"],
log_transform(data["pre_E"]),
data["pre_dir"],
data["layer_id"].astype(np.float32),
data["n_sec"].astype(np.float32),
log_transform(data["e_sec"]),
]
).astype(np.float32) # (N, COND_DIM)
).astype(np.float32) # (N, COND_DIM=8)
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)
cond_cat = np.column_stack([pdg_idx, mat_idx]) # (N, 2)
n_sec_raw = data["n_sec"].astype(
np.int64
) # (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.
n_sec = np.minimum(n_sec_raw, K_MAX).astype(np.int64) # (N,)
# Secondary continuous targets
sec_E_list = data.get("sec_E_list")
sec_dir_list = data.get("sec_dir_list")
sec_pdg_list = data.get("sec_pdg_list")
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)
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
# per-secondary columns are absent, the file was never run through the
# parent->child join (steps_to_parquet._add_secondary_attributes /
# `dwarf convert`). Zero-filling here would collapse every secondary to
# PDG index 0 and a constant energy fraction — a broken Stage 2 with no
# error. Callers that only need Stage-1 (e.g. `giant predict`) keep the
# default require_secondaries=False.
if require_secondaries and n_sec_raw.max(initial=0) > 0:
n_with_sec = int((n_sec_raw > 0).sum())
raise ValueError(
f"{n_with_sec} step(s) have secondaries (n_sec > 0) but the "
"per-secondary columns (sec_E_list / sec_pdg_list / sec_dx_list "
"…) are missing. This parquet was not run through the "
"parent->child join (steps_to_parquet._add_secondary_attributes "
"/ `dwarf convert`); training on it would silently zero all "
"Stage-2 targets. Re-convert the file, or pass "
"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)
if fit:
cond_normalizer = Normalizer().fit(cond_cont)
target_normalizer = Normalizer().fit(target)
target_normalizer = Normalizer().fit(target_s1)
if cond_normalizer is not None:
cond_cont = cond_normalizer.transform(cond_cont)
if target_normalizer is not None:
target = target_normalizer.transform(target)
target_s1 = target_normalizer.transform(target_s1)
return cond_cont, cond_cat, target, cond_normalizer, target_normalizer
return (
cond_cont,
cond_cat,
target_s1,
n_sec,
sec_cont,
sec_pdg_idx,
cond_normalizer,
target_normalizer,
)
giant/geometry.py
+486
View File
@@ -0,0 +1,486 @@
"""Geometry oracle: learn a position -> (material, layer_id) map from step data.
The surrogate conditions on `material` and `layer_id`, but does not predict
them — during a shower rollout they must be looked up from the new position.
There is no in-repo detector geometry (it lives in external miniCaloSim), so we
learn it from positions sampled from a real steps dataset.
miniCaloSim's detector is a stack of planar layer slabs along one axis (see
`physics/detector-design/minicalosim-geometry.md`), so `material`/`layer_id`
are a pure function of depth. The default ("slab") method exploits this: fit a
1D lookup table of depth-axis segment boundaries and do an exact O(log
#segments) binary search per query, with escape decided by depth/transverse
bounds — far cheaper per call than a nearest-neighbour search over hundreds of
thousands of reference points, which matters because this oracle is queried on
every autoregressive step of a shower rollout. "knn"/"svm" remain as generic
fallbacks (a classifier over 3D positions, escape decided by distance to the
nearest reference point) for geometries that aren't simple slab stacks.
scikit-learn / joblib are an optional dependency (the `geometry` extra) needed
by "knn"/"svm" and by `save`/`load`; they are imported lazily so the core
install stays lean.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Iterable
import numpy as np
import pandas as pd
import pyarrow.parquet as pq
_INSTALL_HINT = (
"the geometry oracle needs scikit-learn — install it with "
"`uv sync --extra cpu --extra geometry`"
)
def _require_sklearn():
try:
import joblib # noqa: F401
from sklearn.neighbors import KNeighborsClassifier # noqa: F401
from sklearn.svm import SVC # noqa: F401
except ImportError as exc: # pragma: no cover - exercised only without extra
raise ImportError(_INSTALL_HINT) from exc
@dataclass
class _SlabLookup:
"""Fast path for a detector that is a stack of planar layer slabs along one
axis (miniCaloSim's actual geometry — see `giant/geometry.py` module docstring
and `physics/detector-design/minicalosim-geometry.md`). `material`/`layer_id`
are then a pure function of depth, found by binary search over `z_edges`
instead of a nearest-neighbour search over the whole reference point cloud —
O(log(#segments)) instead of O(log(#reference points)), with a far smaller
constant factor, and exact rather than approximate.
"""
axis: int # which of the 3 position components is the depth axis
z_edges: np.ndarray # (n_segments + 1,) sorted boundaries between segments
materials: np.ndarray # (n_segments,) object, material of each segment
layer_ids: np.ndarray # (n_segments,) int64, layer_id of each segment
radius_max: float # largest transverse radius seen in training data
def query(
self, pos: np.ndarray, margin: float
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
other = [i for i in range(3) if i != self.axis]
z = pos[:, self.axis]
radius = np.sqrt(pos[:, other[0]] ** 2 + pos[:, other[1]] ** 2)
idx = np.searchsorted(self.z_edges, z, side="right") - 1
idx = np.clip(idx, 0, len(self.materials) - 1)
material = self.materials[idx]
layer_id = self.layer_ids[idx]
escaped = (
(z < self.z_edges[0] - margin)
| (z > self.z_edges[-1] + margin)
| (radius > self.radius_max + margin)
)
return material, layer_id, escaped
@dataclass
class GeometryOracle:
"""Maps world-frame position -> (material, layer_id, escaped).
Two lookup strategies are supported (`metadata["method"]`):
- `"slab"`: exact O(log #segments) binary search exploiting the known
layered-slab detector geometry (see `_SlabLookup`). Fast and preferred.
- `"knn"` / `"svm"`: a generic sklearn classifier over 3D positions,
predicting a class index into `classes` (a list of (material, layer_id)
pairs). Kept as a fallback for geometries that aren't simple slab stacks.
`escape_threshold` is a distance in position units (mm). For knn/svm it's
compared against the nearest training reference point. For slab it's the
slack allowed beyond the observed depth range / transverse radius before a
point is flagged `escaped`.
"""
estimator: Any
classes: list[tuple[str, int]]
escape_threshold: float
metadata: dict
# Only populated for non-neighbour estimators (SVM) to answer the escape
# distance query; KNeighborsClassifier answers it directly.
_ref_tree: Any = field(default=None)
# Only populated when metadata["method"] == "slab".
_slab: _SlabLookup | None = field(default=None)
def query(self, pos: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Return (material (N,) str, layer_id (N,) int, escaped (N,) bool)."""
pos = np.ascontiguousarray(np.asarray(pos, dtype=np.float64))
if pos.ndim != 2 or pos.shape[1] != 3:
raise ValueError(f"pos must be (N, 3), got {pos.shape}")
if len(pos) == 0:
return (
np.empty(0, dtype=object),
np.empty(0, dtype=np.int64),
np.empty(0, dtype=bool),
)
if self._slab is not None:
return self._slab.query(pos, self.escape_threshold)
# kneighbors gives the distance to the nearest reference point, which is
# what the escape test needs; it exists on both KNeighborsClassifier and
# (via a stored reference tree) our SVM wrapper below.
dist = self._nearest_distance(pos)
escaped = dist > self.escape_threshold
cls_idx = self.estimator.predict(pos).astype(np.int64)
material = np.array([self.classes[i][0] for i in cls_idx], dtype=object)
layer_id = np.array([self.classes[i][1] for i in cls_idx], dtype=np.int64)
return material, layer_id, escaped
def _nearest_distance(self, pos: np.ndarray) -> np.ndarray:
from sklearn.neighbors import KNeighborsClassifier
if isinstance(self.estimator, KNeighborsClassifier):
dist, _ = self.estimator.kneighbors(pos, n_neighbors=1)
return dist[:, 0]
# SVM (or any non-neighbour estimator): use the separately stored
# NearestNeighbors index purely for the escape distance.
dist, _ = self._ref_tree.kneighbors(pos, n_neighbors=1)
return dist[:, 0]
def save(self, path: str | Path) -> None:
_require_sklearn()
import joblib
joblib.dump(
{
"estimator": self.estimator,
"classes": self.classes,
"escape_threshold": self.escape_threshold,
"metadata": self.metadata,
"ref_tree": getattr(self, "_ref_tree", None),
"slab": getattr(self, "_slab", None),
},
path,
)
@classmethod
def load(cls, path: str | Path) -> "GeometryOracle":
_require_sklearn()
import joblib
d = joblib.load(path)
obj = cls(
estimator=d["estimator"],
classes=[tuple(c) for c in d["classes"]],
escape_threshold=float(d["escape_threshold"]),
metadata=d.get("metadata", {}),
)
obj._ref_tree = d.get("ref_tree")
obj._slab = d.get("slab")
return obj
_PRE_COLS = ["pre_x", "pre_y", "pre_z"]
_POST_COLS = ["post_x", "post_y", "post_z"]
_LABEL_COLS = ["material", "layer_id"]
def _iter_point_batches(path: Path, batch_size: int = 1_000_000):
"""Yield (pos (M,3), material (M,), layer_id (M,)) from any parquet with a
position + material + layer_id schema (raw steps *or* predict output).
Only the needed columns are read. post_pos points are included when present
(they share their step's label) so boundary regions are densely sampled.
"""
pf = pq.ParquetFile(path)
have = set(pf.schema_arrow.names)
missing = [c for c in (*_PRE_COLS, *_LABEL_COLS) if c not in have]
if missing:
raise ValueError(
f"{path} is missing columns {missing} needed to build a geometry "
"oracle (expected pre_x/y/z, material, layer_id)"
)
has_post = all(c in have for c in _POST_COLS)
cols = [*_PRE_COLS, *_LABEL_COLS] + (_POST_COLS if has_post else [])
for batch in pf.iter_batches(batch_size=batch_size, columns=cols):
d = batch.to_pydict()
pre = np.array([d[c] for c in _PRE_COLS], dtype=np.float32).T
mat = np.array([str(m) for m in d["material"]], dtype=object)
lay = np.asarray(d["layer_id"], dtype=np.int64)
if has_post:
post = np.array([d[c] for c in _POST_COLS], dtype=np.float32).T
yield (
np.concatenate([pre, post], axis=0),
np.concatenate([mat, mat], axis=0),
np.concatenate([lay, lay], axis=0),
)
else:
yield pre, mat, lay
def _collect_points(
files: Iterable[Path],
subsample: int,
seed: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Stream files, reservoir-sample (pos, material, layer_id) points.
Reservoir sampling keeps memory bounded regardless of total file size.
"""
rng = np.random.default_rng(seed)
res_pos = np.empty((subsample, 3), dtype=np.float32)
res_mat = np.empty(subsample, dtype=object)
res_lay = np.empty(subsample, dtype=np.int64)
seen = 0
for path in files:
for pos, mat, lay in _iter_point_batches(path):
m = len(pos)
if seen < subsample:
take = min(subsample - seen, m)
res_pos[seen : seen + take] = pos[:take]
res_mat[seen : seen + take] = mat[:take]
res_lay[seen : seen + take] = lay[:take]
seen += take
if take < m:
# Reservoir is now full; run the standard replacement rule
# on the remainder of this chunk.
_reservoir_replace(
res_pos,
res_mat,
res_lay,
pos[take:],
mat[take:],
lay[take:],
seen,
rng,
)
seen += m - take
else:
_reservoir_replace(res_pos, res_mat, res_lay, pos, mat, lay, seen, rng)
seen += m
n = min(seen, subsample)
return res_pos[:n], res_mat[:n], res_lay[:n]
def _reservoir_replace(res_pos, res_mat, res_lay, pos, mat, lay, seen, rng) -> None:
"""Vectorized reservoir replacement for a batch of incoming points."""
m = len(pos)
k = res_pos.shape[0]
# For incoming global index j (seen..seen+m-1), keep with prob k/(j+1),
# replacing a uniformly-chosen reservoir slot.
idx = seen + np.arange(m)
j = rng.integers(0, idx + 1) # j in [0, global_index]
keep = j < k
slots = j[keep]
res_pos[slots] = pos[keep]
res_mat[slots] = mat[keep]
res_lay[slots] = lay[keep]
def _fit_slab_lookup(
pos: np.ndarray,
mat: np.ndarray,
lay: np.ndarray,
axis: int,
n_bins: int,
) -> tuple[_SlabLookup, dict]:
"""Fit a `_SlabLookup` assuming material/layer_id are a function of depth
(`pos[:, axis]`) alone — true for a stack of planar layer slabs.
Bins the depth axis into `n_bins` equal-width bins, takes the majority
(material, layer_id) label per bin (robust to the handful of points near a
boundary whose true label is ambiguous at bin resolution), fills any empty
bins from the nearest populated bin, then run-length-encodes consecutive
bins sharing a label into segments. Binary search over the segment
boundaries then answers a query in O(log #segments).
"""
other = [i for i in range(3) if i != axis]
z = pos[:, axis].astype(np.float64)
radius = np.sqrt(
pos[:, other[0]].astype(np.float64) ** 2
+ pos[:, other[1]].astype(np.float64) ** 2
)
z_min, z_max = float(z.min()), float(z.max())
if z_min == z_max:
raise ValueError(
"all points share the same depth-axis coordinate — pick a "
"different `depth_axis` or use method='knn'/'svm'"
)
edges = np.linspace(z_min, z_max, n_bins + 1)
bin_idx = np.clip(np.searchsorted(edges, z, side="right") - 1, 0, n_bins - 1)
counts = (
pd.DataFrame({"bin": bin_idx, "material": mat, "layer_id": lay})
.groupby(["bin", "material", "layer_id"])
.size()
.to_frame("n")
.reset_index()
.sort_values("n", ascending=False)
.drop_duplicates("bin")
)
bin_material = np.full(n_bins, "", dtype=object)
bin_layer = np.full(n_bins, -1, dtype=np.int64)
has_data = np.zeros(n_bins, dtype=bool)
idx = counts["bin"].to_numpy()
bin_material[idx] = counts["material"].to_numpy()
bin_layer[idx] = counts["layer_id"].to_numpy()
has_data[idx] = True
# Forward/backward-fill bins with no samples from the nearest populated one.
fill_from = np.where(has_data, np.arange(n_bins), -1)
for b in range(1, n_bins):
if fill_from[b] == -1:
fill_from[b] = fill_from[b - 1]
for b in range(n_bins - 2, -1, -1):
if fill_from[b] == -1:
fill_from[b] = fill_from[b + 1]
bin_material = bin_material[fill_from]
bin_layer = bin_layer[fill_from]
# Run-length-encode consecutive bins sharing a label into segments.
changed = (
np.flatnonzero(
(bin_material[1:] != bin_material[:-1]) | (bin_layer[1:] != bin_layer[:-1])
)
+ 1
)
seg_starts = np.concatenate([[0], changed])
z_edges = np.concatenate([edges[seg_starts], edges[-1:]])
materials = bin_material[seg_starts]
layer_ids = bin_layer[seg_starts]
unique_z = np.unique(z)
median_spacing = float(np.median(np.diff(unique_z))) if len(unique_z) > 1 else 1.0
radius_max = float(radius.max())
slab = _SlabLookup(
axis=axis,
z_edges=z_edges,
materials=materials,
layer_ids=layer_ids,
radius_max=radius_max,
)
info = {
"n_segments": int(len(materials)),
"z_range": (z_min, z_max),
"median_z_spacing": median_spacing,
"radius_max": radius_max,
}
return slab, info
def build_geometry_oracle(
files: list[Path],
method: str = "slab",
k: int = 1,
subsample: int = 500_000,
escape_factor: float = 5.0,
seed: int = 0,
depth_axis: int = 2,
n_bins: int = 2000,
) -> GeometryOracle:
"""Fit a position -> (material, layer_id) classifier from steps files.
method: "slab" (default-recommended fast path exploiting the known
layered-slab detector geometry — see `_SlabLookup`), "knn"
(KNeighborsClassifier), or "svm" (SVC). "slab" is O(log #segments) per
query and exact; "knn"/"svm" are generic fallbacks for geometries that
aren't simple slab stacks, at the cost of a much slower query (a
nearest-neighbour or kernel evaluation against up to `subsample`
reference points) and, for "svm", occasional misclassification.
k: neighbours for the knn classifier (ignored otherwise).
subsample: max reference points held in memory / used for the fit.
escape_factor: escape_threshold = escape_factor * median spacing of the
reference points along the relevant axis/axes, so it scales with the
sampling density of the data.
depth_axis: index (0/1/2 -> x/y/z) of the position component that layers
stack along. Only used by method="slab"; default 2 (z) matches
miniCaloSim's beam-axis-aligned layer stack.
n_bins: depth-axis resolution for method="slab" — should be finer than the
thinnest layer.
"""
pos, mat, lay = _collect_points(files, subsample, seed)
if len(pos) == 0:
raise ValueError("no points collected — are these steps parquet files?")
# Combined (material, layer_id) class label, used for `classes` regardless
# of method (informational for slab; the actual classifier index for
# knn/svm).
pairs = list(zip((str(m) for m in mat), (int(v) for v in lay)))
classes = sorted(set(pairs))
if method == "slab":
slab, info = _fit_slab_lookup(pos, mat, lay, axis=depth_axis, n_bins=n_bins)
escape_threshold = escape_factor * info["median_z_spacing"]
oracle = GeometryOracle(
estimator=None,
classes=classes,
escape_threshold=escape_threshold,
metadata={
"method": "slab",
"depth_axis": depth_axis,
"n_bins": n_bins,
"n_reference_points": int(len(pos)),
"escape_factor": escape_factor,
"n_files": len(files),
**info,
},
)
oracle._slab = slab
return oracle
_require_sklearn()
from sklearn.neighbors import KNeighborsClassifier, NearestNeighbors
from sklearn.svm import SVC
class_to_idx = {c: i for i, c in enumerate(classes)}
y = np.array([class_to_idx[p] for p in pairs], dtype=np.int64)
X = pos.astype(np.float64)
if method == "knn":
estimator = KNeighborsClassifier(n_neighbors=k)
estimator.fit(X, y)
ref_tree = None
elif method == "svm":
estimator = SVC(kernel="rbf")
estimator.fit(X, y)
# SVM cannot answer nearest-neighbour distance queries, so keep a light
# reference tree alongside it purely for the escape test.
ref_tree = NearestNeighbors(n_neighbors=1).fit(X)
else:
raise ValueError(f"unknown method {method!r}; use 'slab', 'knn', or 'svm'")
# Escape threshold from the reference point spacing. Sample a subset for the
# median 2-NN distance (the 1st neighbour of a training point is itself).
nn = NearestNeighbors(n_neighbors=2).fit(X)
probe = (
X
if len(X) <= 20_000
else X[np.random.default_rng(seed).choice(len(X), 20_000, replace=False)]
)
d2, _ = nn.kneighbors(probe, n_neighbors=2)
median_nn = float(np.median(d2[:, 1]))
escape_threshold = escape_factor * median_nn
oracle = GeometryOracle(
estimator=estimator,
classes=classes,
escape_threshold=escape_threshold,
metadata={
"method": method,
"k": k,
"n_reference_points": int(len(X)),
"median_nn_dist": median_nn,
"escape_factor": escape_factor,
"n_files": len(files),
},
)
oracle._ref_tree = ref_tree
return oracle
giant/model/network.py
+124 -2
View File
@@ -3,7 +3,7 @@ import math
import torch
import torch.nn as nn
from giant.constants import COND_DIM, X_DIM
from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
class SinusoidalEmbedding(nn.Module):
@@ -70,6 +70,12 @@ class ResBlock(nn.Module):
class DenoisingMLP(nn.Module):
"""Stage-1 model: predicts the 9D primary post-step vector field + n_sec logits.
The n_sec head runs on the condition encoding only (no diffusion noise),
so it can be called at inference time independently via `predict_n_sec`.
"""
def __init__(
self,
pdg_vocab: int,
@@ -81,6 +87,7 @@ class DenoisingMLP(nn.Module):
cond_out_dim: int = 128,
x_dim: int = X_DIM,
dropout: float = 0.1,
k_max: int = K_MAX,
) -> None:
super().__init__()
self.time_emb = SinusoidalEmbedding(time_dim)
@@ -99,6 +106,13 @@ class DenoisingMLP(nn.Module):
]
)
self.out_proj = nn.Linear(hidden_dim, x_dim)
# Predicts n_sec as classification over {0, 1, ..., k_max}.
# Applied to the condition encoding (not the diffused latent).
self.n_sec_head = nn.Sequential(
nn.Linear(cond_out_dim, hidden_dim // 2),
nn.SiLU(),
nn.Linear(hidden_dim // 2, k_max + 1),
)
def forward(
self,
@@ -109,7 +123,115 @@ class DenoisingMLP(nn.Module):
) -> torch.Tensor:
t_emb = self.time_emb(t) # (B, time_dim)
c_emb = self.cond_enc(cond_cont, cond_cat) # (B, cond_out_dim)
cond = torch.cat([t_emb, c_emb], dim=-1) # (B, time_dim+cond_out_dim)
cond = torch.cat([t_emb, c_emb], dim=-1)
x = self.input_proj(x_t)
for block in self.blocks:
x = block(x, cond)
return self.out_proj(x)
def predict_n_sec(
self,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
) -> torch.Tensor:
"""Return n_sec logits (B, K_MAX+1) from conditioning alone."""
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."""
def __init__(
self,
pdg_vocab: int,
mat_vocab: int,
emb_dim: int = 16,
cond_out_dim: int = 128,
stage1_dim: int = X_DIM,
stage1_proj_dim: int = 64,
out_dim: int = 128,
) -> None:
super().__init__()
self.base = ConditionEncoder(
pdg_vocab=pdg_vocab,
mat_vocab=mat_vocab,
emb_dim=emb_dim,
out_dim=cond_out_dim,
)
self.stage1_proj = nn.Linear(stage1_dim, stage1_proj_dim)
fused_dim = cond_out_dim + stage1_proj_dim
self.fuse = nn.Sequential(
nn.Linear(fused_dim, out_dim),
nn.SiLU(),
)
def forward(
self,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
stage1_out: torch.Tensor,
) -> torch.Tensor:
base = self.base(cond_cont, cond_cat) # (B, cond_out_dim)
s1 = self.stage1_proj(stage1_out).tanh() # (B, stage1_proj_dim)
return self.fuse(torch.cat([base, s1], dim=-1)) # (B, out_dim)
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.
"""
def __init__(
self,
pdg_vocab: int,
mat_vocab: int,
hidden_dim: int = 256,
n_blocks: int = 6,
emb_dim: int = 16,
time_dim: int = 64,
cond_out_dim: int = 128,
stage1_proj_dim: int = 64,
sec_dim: int = SEC_DIM,
dropout: float = 0.1,
) -> None:
super().__init__()
self.time_emb = SinusoidalEmbedding(time_dim)
self.cond_enc = SecondaryConditionEncoder(
pdg_vocab=pdg_vocab,
mat_vocab=mat_vocab,
emb_dim=emb_dim,
cond_out_dim=cond_out_dim,
stage1_proj_dim=stage1_proj_dim,
out_dim=cond_out_dim,
)
merged_cond_dim = time_dim + cond_out_dim
self.input_proj = nn.Linear(sec_dim, hidden_dim)
self.blocks = nn.ModuleList(
[
ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
for _ in range(n_blocks)
]
)
self.out_proj = nn.Linear(hidden_dim, sec_dim)
def forward(
self,
x_t: torch.Tensor,
t: torch.Tensor,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
stage1_out: torch.Tensor,
) -> torch.Tensor:
t_emb = self.time_emb(t)
c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
cond = torch.cat([t_emb, c_emb], dim=-1)
x = self.input_proj(x_t)
for block in self.blocks:
x = block(x, cond)
giant/model/schedule.py
+42
View File
@@ -69,3 +69,45 @@ def flow_matching_loss(
u_t = x1 - x0
v_t = model(x_t, t, cond_cont, cond_cat)
return F.mse_loss(v_t, u_t)
def flow_matching_loss_secondary(
model: torch.nn.Module,
x1: torch.Tensor,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
stage1_out: torch.Tensor,
sec_mask: torch.Tensor,
) -> 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)
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.
"""
from giant.constants import CONT_SLOT_DIM, EMB_DIM, K_MAX, SEC_SLOT_DIM
B = x1.size(0)
t = torch.rand(B, device=x1.device)
x0 = torch.randn_like(x1)
x_t = (1.0 - t.view(-1, 1)) * x0 + t.view(-1, 1) * x1
u_t = x1 - x0
v_t = model(x_t, t, cond_cont, cond_cat, stage1_out)
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)
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
giant/pipeline.py
+35 -11
View File
@@ -5,7 +5,7 @@ import torch
from torch.utils.data import DataLoader
from giant import config
from giant.constants import COND_DIM, X_DIM
from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_SLOT_DIM, X_DIM
from giant.data.loader import (
find_parquet_files,
load_event_ids,
@@ -14,7 +14,7 @@ from giant.data.loader import (
)
from giant.data.transforms import build_features, _WelfordAccumulator
from giant.data.dataset import make_event_split, StreamingStepsDataset
from giant.model.network import DenoisingMLP
from giant.model.network import DenoisingMLP, SecondaryDecoder
from giant.train import train as run_training
@@ -63,9 +63,11 @@ def run_train_job(
if not mask.any():
continue
chunk_tr = {k: v[mask] for k, v in chunk.items()}
cond_cont, _, target, _, _ = build_features(chunk_tr, pdg_map, mat_map)
cond_cont, _, target_s1, _n_sec, _sec_cont, _sec_pdg, _, _ = build_features(
chunk_tr, pdg_map, mat_map, require_secondaries=True
)
cond_acc.update(cond_cont)
tgt_acc.update(target)
tgt_acc.update(target_s1)
cond_norm = cond_acc.to_normalizer()
tgt_norm = tgt_acc.to_normalizer()
@@ -91,8 +93,6 @@ def run_train_job(
shuffle=False,
)
# Dataset yields whole batches already, so batch_size=None tells DataLoader
# to pass them through instead of re-collating row-by-row in Python.
pin = device.type == "cuda"
train_loader = DataLoader(
train_ds,
@@ -107,15 +107,34 @@ def run_train_job(
pin_memory=pin,
)
model = DenoisingMLP(
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"
)
stage1_model = DenoisingMLP(
pdg_vocab=len(pdg_map),
mat_vocab=len(mat_map),
hidden_dim=m["hidden_dim"],
n_blocks=m["n_blocks"],
emb_dim=m["emb_dim"],
emb_dim=emb_dim,
dropout=m["dropout"],
k_max=K_MAX,
)
sec_decoder = SecondaryDecoder(
pdg_vocab=len(pdg_map),
mat_vocab=len(mat_map),
hidden_dim=m["hidden_dim"],
n_blocks=m["n_blocks"],
emb_dim=emb_dim,
dropout=m["dropout"],
)
echo(f"model: {sum(p.numel() for p in model.parameters()):,} parameters")
echo(
f"stage1: {sum(p.numel() for p in stage1_model.parameters()):,} parameters | "
f"sec_decoder: {sum(p.numel() for p in sec_decoder.parameters()):,} parameters"
)
out_dir.mkdir(parents=True, exist_ok=True)
meta = config.build_run_meta(
@@ -134,12 +153,15 @@ def run_train_job(
"mat_vocab": len(mat_map),
"hidden_dim": m["hidden_dim"],
"n_blocks": m["n_blocks"],
"emb_dim": m["emb_dim"],
"emb_dim": emb_dim,
"dropout": m["dropout"],
"k_max": K_MAX,
"sec_slot_dim": SEC_SLOT_DIM,
}
run_training(
model=model,
stage1_model=stage1_model,
sec_decoder=sec_decoder,
train_loader=train_loader,
val_loader=val_loader,
mode=t["mode"],
@@ -148,6 +170,8 @@ def run_train_job(
warmup_epochs=t["warmup_epochs"],
device=device,
out_dir=out_dir,
lambda_nsec=t.get("lambda_nsec", 0.1),
lambda_s2=t.get("lambda_s2", 1.0),
normalizer_dict={"cond": cond_norm.to_dict(), "target": tgt_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()},
giant/rollout.py
+512
View File
@@ -0,0 +1,512 @@
"""Autoregressive shower rollout driver.
Steps the two-stage GIANT surrogate forward into a full particle shower: each
primary post-step becomes the next pre-step, secondaries are pushed as new tracks,
and the material/layer_id conditioning at every step comes from a `GeometryOracle`
(the surrogate does not predict them).
Tracks are advanced breadth-first: every sweep steps all currently-active tracks
once (in `batch_size` chunks), so many tracks share each model forward pass. A
track terminates on one of the recorded `termination_reason`s in constants.py.
Energy accounting: on every terminal stop except escape, the track's remaining
energy is deposited locally so the shower conserves energy; escaped energy is
treated as detector leakage and not deposited.
"""
from __future__ import annotations
import numpy as np
import torch
from giant.constants import (
TERM_ENERGY_CUTOFF,
TERM_ESCAPED,
TERM_MAX_STEPS,
TERM_NATURAL_END,
TERM_UNKNOWN_PDG,
)
from giant.data.transforms import (
Normalizer,
build_cond_features,
decode_secondaries,
energy_simplex_decode,
inv_local_frame_rotation,
inv_log_transform,
reconstruct_post_pos,
)
from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx
# Record columns produced per step / per terminal marker.
_RECORD_KEYS = [
"event_id",
"track_id",
"parent_id",
"generation",
"step_no",
"pdg",
"pre_x",
"pre_y",
"pre_z",
"pre_E",
"pre_dx",
"pre_dy",
"pre_dz",
"post_x",
"post_y",
"post_z",
"post_E",
"post_dx",
"post_dy",
"post_dz",
"edep",
"step_length",
"material",
"layer_id",
"n_sec_pred",
"termination_reason",
]
def _empty_frontier() -> dict[str, np.ndarray]:
return {
"event_id": np.empty(0, dtype=np.int64),
"track_id": np.empty(0, dtype=np.int64),
"parent_id": np.empty(0, dtype=np.int64),
"generation": np.empty(0, dtype=np.int64),
"step_in_track": np.empty(0, dtype=np.int64),
"pdg": np.empty(0, dtype=np.int64),
"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),
}
def _concat_frontiers(parts: list[dict[str, np.ndarray]]) -> dict[str, np.ndarray]:
parts = [p for p in parts if len(p["event_id"]) > 0]
if not parts:
return _empty_frontier()
return {k: np.concatenate([p[k] for p in parts], axis=0) for k in parts[0]}
class _Recorder:
"""Accumulates per-step rows into column lists, materialised at the end."""
def __init__(self) -> None:
self._cols: dict[str, list] = {k: [] for k in _RECORD_KEYS}
def add(self, **cols) -> None:
n = len(cols["event_id"])
if n == 0:
return
for k in _RECORD_KEYS:
v = cols[k]
self._cols[k].append(np.asarray(v).reshape(n))
def to_dict(self) -> dict[str, np.ndarray]:
out = {}
for k, chunks in self._cols.items():
if chunks:
out[k] = np.concatenate(chunks, axis=0)
else:
out[k] = np.empty(
0,
dtype=object
if k in ("material", "termination_reason")
else np.float64,
)
return out
def make_seed_frontier(
event_id: np.ndarray,
pdg: np.ndarray,
pre_pos: np.ndarray,
pre_E: np.ndarray,
pre_dir: np.ndarray,
) -> tuple[dict[str, np.ndarray], dict[int, int]]:
"""Build the initial frontier from primary entry states.
Returns (frontier, event_track_count) where the latter tracks the next
unused track_id per event (each primary gets a fresh id starting from 0).
"""
event_id = np.asarray(event_id, dtype=np.int64)
n = len(event_id)
track_id = np.empty(n, dtype=np.int64)
counts: dict[int, int] = {}
for i, ev in enumerate(event_id.tolist()):
c = counts.get(ev, 0)
track_id[i] = c
counts[ev] = c + 1
dir_ = np.asarray(pre_dir, dtype=np.float64)
dir_ = dir_ / np.clip(np.linalg.norm(dir_, axis=1, keepdims=True), 1e-12, None)
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),
"pre_pos": np.asarray(pre_pos, dtype=np.float64),
"pre_E": np.asarray(pre_E, dtype=np.float64),
"pre_dir": dir_,
}
return frontier, counts
def _terminal_rows(tr: dict[str, np.ndarray], sel: np.ndarray, reason: str, edep):
"""Assemble terminal-marker record columns for the selected tracks."""
pos = tr["pre_pos"][sel]
dir_ = tr["pre_dir"][sel]
n = int(sel.sum())
return dict(
event_id=tr["event_id"][sel],
track_id=tr["track_id"][sel],
parent_id=tr["parent_id"][sel],
generation=tr["generation"][sel],
step_no=tr["step_in_track"][sel],
pdg=tr["pdg"][sel],
pre_x=pos[:, 0],
pre_y=pos[:, 1],
pre_z=pos[:, 2],
pre_E=tr["pre_E"][sel],
pre_dx=dir_[:, 0],
pre_dy=dir_[:, 1],
pre_dz=dir_[:, 2],
post_x=pos[:, 0],
post_y=pos[:, 1],
post_z=pos[:, 2],
post_E=np.zeros(n),
post_dx=dir_[:, 0],
post_dy=dir_[:, 1],
post_dz=dir_[:, 2],
edep=np.asarray(edep, dtype=np.float64).reshape(n),
step_length=np.zeros(n),
material=tr.get("_material", np.full(len(sel), "", dtype=object))[sel],
layer_id=tr.get("_layer_id", np.zeros(len(sel), dtype=np.int64))[sel],
n_sec_pred=np.zeros(n, dtype=np.int64),
termination_reason=np.full(n, reason, dtype=object),
)
@torch.no_grad()
def rollout(
stage1_model: torch.nn.Module,
sec_decoder: torch.nn.Module,
oracle,
seeds: dict[str, np.ndarray],
cond_norm: Normalizer,
tgt_norm: Normalizer,
pdg_map: dict[int, int],
mat_map: dict[str, int],
*,
energy_cutoff: float,
max_steps: int,
steps: int = 10,
batch_size: int = 4096,
device: torch.device | None = None,
max_tracks_per_event: int | None = None,
escape_threshold: float | None = None,
) -> dict[str, np.ndarray]:
"""Run showers to completion; return a step-record dict (see _RECORD_KEYS)."""
device = device or torch.device("cpu")
stage1_model.eval()
sec_decoder.eval()
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"],
seeds["pre_pos"],
seeds["pre_E"],
seeds["pre_dir"],
)
rec = _Recorder()
while len(frontier["event_id"]) > 0:
next_parts: list[dict[str, np.ndarray]] = []
n_total = len(frontier["event_id"])
for start in range(0, n_total, batch_size):
chunk = {k: v[start : start + batch_size] for k, v in frontier.items()}
next_parts.append(
_step_chunk(
chunk,
stage1_model,
sec_decoder,
oracle,
cond_norm,
tgt_norm,
pdg_map,
mat_map,
pdg_map_inv,
pdg_emb_weight,
rec,
counts,
energy_cutoff,
max_steps,
steps,
device,
max_tracks_per_event,
)
)
frontier = _concat_frontiers(next_parts)
return rec.to_dict()
def _step_chunk(
tr,
stage1_model,
sec_decoder,
oracle,
cond_norm,
tgt_norm,
pdg_map,
mat_map,
pdg_map_inv,
pdg_emb_weight,
rec,
counts,
energy_cutoff,
max_steps,
steps,
device,
max_tracks_per_event,
) -> dict[str, np.ndarray]:
"""Advance one chunk of tracks by a single step; return the next frontier."""
n = len(tr["event_id"])
# --- Geometry lookup + material/layer conditioning ---
material, layer_id, escaped = oracle.query(tr["pre_pos"])
tr = dict(tr)
tr["_material"] = material
tr["_layer_id"] = layer_id
known_pdg = np.array([int(p) in pdg_map for p in tr["pdg"]], dtype=bool)
# --- Pre-step termination gates (in priority order; each track picks one) ---
stop = np.zeros(n, dtype=bool)
escaped_sel = escaped & ~stop
rec.add(
**_terminal_rows(
tr, escaped_sel, TERM_ESCAPED, edep=np.zeros(int(escaped_sel.sum()))
)
)
stop |= escaped_sel
unknown_sel = ~known_pdg & ~stop
rec.add(
**_terminal_rows(
tr, unknown_sel, TERM_UNKNOWN_PDG, edep=tr["pre_E"][unknown_sel]
)
)
stop |= unknown_sel
cutoff_sel = (tr["pre_E"] < energy_cutoff) & ~stop
rec.add(
**_terminal_rows(
tr, cutoff_sel, TERM_ENERGY_CUTOFF, edep=tr["pre_E"][cutoff_sel]
)
)
stop |= cutoff_sel
maxstep_sel = (tr["step_in_track"] >= max_steps) & ~stop
rec.add(
**_terminal_rows(tr, maxstep_sel, TERM_MAX_STEPS, edep=tr["pre_E"][maxstep_sel])
)
stop |= maxstep_sel
active = ~stop
if not active.any():
return _empty_frontier()
tr = {k: v[active] for k, v in tr.items()}
material = tr["_material"]
layer_id = tr["_layer_id"]
# --- Build conditioning and run the two stages ---
cond_dict = {
"pre_pos": tr["pre_pos"],
"pre_E": tr["pre_E"],
"pre_dir": tr["pre_dir"],
"layer_id": layer_id,
"material": material,
"pdg": tr["pdg"],
}
cond_cont, cond_cat = build_cond_features(cond_dict, pdg_map, mat_map, cond_norm)
cc = torch.from_numpy(cond_cont).float().to(device)
ck = torch.from_numpy(cond_cat).long().to(device)
stage1_norm, n_sec_pred = sample_flow(stage1_model, cc, ck, steps=steps)
raw = tgt_norm.inverse_transform(stage1_norm.cpu().numpy())
step_length = inv_log_transform(raw[:, 0])
edep, e_sec, post_E, _delta = energy_simplex_decode(raw[:, 1:3], tr["pre_E"])
post_dir_local = raw[:, 3:6].copy()
post_dir_local /= np.clip(
np.linalg.norm(post_dir_local, axis=1, keepdims=True), 1e-8, None
)
post_dir_world = inv_local_frame_rotation(tr["pre_dir"], post_dir_local)
travel_dir_local = raw[:, 6:9].copy()
travel_dir_local /= np.clip(
np.linalg.norm(travel_dir_local, axis=1, keepdims=True), 1e-8, None
)
post_pos = reconstruct_post_pos(
tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local
)
n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64)
# --- Secondaries ---
sec_cont, sec_type_emb, _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(),
n_sec_np,
e_sec,
tr["pre_dir"],
pdg_map_inv,
)
edep = edep.astype(np.float64)
post_E = post_E.astype(np.float64)
# --- Spawn secondaries (with per-event track cap) ---
new_tracks, dropped_edep = _spawn_secondaries(
tr,
post_pos,
sec_valid,
sec_E,
sec_dir_world,
sec_pdg_code,
counts,
max_tracks_per_event,
)
# Energy bookkeeping so each step conserves exactly (edep + carried + post_E
# == pre_E): `decode_secondaries` already rescales valid slots to sum to
# exactly `e_sec` whenever n_sec > 0, so `residual` here is ~0 except when
# n_sec == 0 (no secondary to carry the budget at all — the whole `e_sec`
# becomes residual). Also deposit the energy of any sub-cap secondaries
# we dropped for hitting `max_tracks_per_event`.
sec_E_valid_sum = (sec_E * sec_valid).sum(axis=1)
residual = np.maximum(e_sec - sec_E_valid_sum, 0.0)
edep = edep + residual + dropped_edep
# --- Record the stepped rows; mark natural_end where the primary died ---
natural = post_E <= 0.0
reason = np.where(natural, TERM_NATURAL_END, "").astype(object)
rec.add(
event_id=tr["event_id"],
track_id=tr["track_id"],
parent_id=tr["parent_id"],
generation=tr["generation"],
step_no=tr["step_in_track"],
pdg=tr["pdg"],
pre_x=tr["pre_pos"][:, 0],
pre_y=tr["pre_pos"][:, 1],
pre_z=tr["pre_pos"][:, 2],
pre_E=tr["pre_E"],
pre_dx=tr["pre_dir"][:, 0],
pre_dy=tr["pre_dir"][:, 1],
pre_dz=tr["pre_dir"][:, 2],
post_x=post_pos[:, 0],
post_y=post_pos[:, 1],
post_z=post_pos[:, 2],
post_E=post_E,
post_dx=post_dir_world[:, 0],
post_dy=post_dir_world[:, 1],
post_dz=post_dir_world[:, 2],
edep=edep,
step_length=step_length,
material=material,
layer_id=layer_id,
n_sec_pred=n_sec_np,
termination_reason=reason,
)
# --- Continue surviving primaries ---
cont = ~natural
cont_frontier = {
"event_id": tr["event_id"][cont],
"track_id": tr["track_id"][cont],
"parent_id": tr["parent_id"][cont],
"generation": tr["generation"][cont],
"step_in_track": tr["step_in_track"][cont] + 1,
"pdg": tr["pdg"][cont],
"pre_pos": post_pos[cont],
"pre_E": post_E[cont],
"pre_dir": post_dir_world[cont],
}
return _concat_frontiers([cont_frontier, new_tracks])
def _spawn_secondaries(
tr,
post_pos,
sec_valid,
sec_E,
sec_dir_world,
sec_pdg_code,
counts,
max_tracks_per_event,
) -> tuple[dict[str, np.ndarray], np.ndarray]:
"""Turn valid secondaries into new tracks; return (frontier, per-parent dropped edep).
Secondaries are born at their parent's post_pos. When `max_tracks_per_event`
is set and an event is at its cap, further secondaries are not spawned; their
energy is returned as `dropped_edep` (indexed by parent row) so it is
deposited into the parent step instead of vanishing.
"""
B = len(tr["event_id"])
dropped_edep = np.zeros(B, dtype=np.float64)
pr, sl = np.nonzero(sec_valid) # parent-row idx, slot idx
if len(pr) == 0:
return _empty_frontier(), dropped_edep
# Assign a fresh per-event track_id to each candidate in stable parent order,
# applying the per-event cap. The candidate count per chunk is small
# (<= batch_size * K_MAX), so a plain loop is clear and fast enough.
order = np.lexsort((sl, pr)) # group by parent row, slot ascending
kept_pr, kept_sl, kept_tid = [], [], []
for j in order:
ev = int(tr["event_id"][pr[j]])
cur = counts.get(ev, 0)
if max_tracks_per_event is not None and cur >= max_tracks_per_event:
dropped_edep[pr[j]] += float(sec_E[pr[j], sl[j]])
continue
kept_pr.append(pr[j])
kept_sl.append(sl[j])
kept_tid.append(cur)
counts[ev] = cur + 1
if not kept_pr:
return _empty_frontier(), dropped_edep
pr_k = np.array(kept_pr, dtype=np.int64)
sl_k = np.array(kept_sl, dtype=np.int64)
tid_k = np.array(kept_tid, dtype=np.int64)
frontier = {
"event_id": tr["event_id"][pr_k],
"track_id": tid_k,
"parent_id": tr["track_id"][pr_k],
"generation": tr["generation"][pr_k] + 1,
"step_in_track": np.zeros(len(pr_k), dtype=np.int64),
"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),
}
return frontier, dropped_edep
giant/sample.py
+75 -10
View File
@@ -1,6 +1,6 @@
import torch
from giant.constants import X_DIM
from giant.constants import K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM
@torch.no_grad()
@@ -9,8 +9,13 @@ def sample_flow(
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
steps: int = 10,
) -> torch.Tensor:
"""Euler integration of the learned vector field from t=0 to t=1."""
) -> tuple[torch.Tensor, torch.Tensor]:
"""Euler integration of the Stage-1 vector field from t=0 to t=1.
Returns (primary_sample, n_sec_pred):
primary_sample: (B, X_DIM) — normalised 9D primary post-step output
n_sec_pred: (B,) int64 — predicted secondary count
"""
model.eval()
B = cond_cont.size(0)
device = cond_cont.device
@@ -20,7 +25,63 @@ def sample_flow(
t = torch.full((B,), i * dt, device=device)
v = model(x, t, cond_cont, cond_cat)
x = x + v * dt
return x
n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
n_sec_pred = n_sec_logits.argmax(dim=-1)
return x, n_sec_pred
@torch.no_grad()
def sample_secondaries(
sec_decoder: torch.nn.Module,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
stage1_out: torch.Tensor,
n_sec_pred: torch.Tensor,
steps: int = 10,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Euler integration of the Stage-2 vector field; return raw slot outputs.
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
"""
sec_decoder.eval()
B = cond_cont.size(0)
device = cond_cont.device
x = torch.randn(B, SEC_DIM, device=device)
dt = 1.0 / steps
for i in range(steps):
t = torch.full((B,), i * dt, device=device)
v = sec_decoder(x, t, cond_cont, cond_cat, stage1_out)
x = x + v * dt
x_slots = x.view(B, K_MAX, SEC_SLOT_DIM)
sec_cont = x_slots[:, :, :4]
sec_type_emb = 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)
@torch.no_grad()
@@ -29,8 +90,8 @@ def sample_ddpm(
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
schedule,
) -> torch.Tensor:
"""Full DDPM ancestral sampling (T reverse steps)."""
) -> tuple[torch.Tensor, torch.Tensor]:
"""Full DDPM ancestral sampling (T reverse steps). Returns (sample, n_sec_pred)."""
model.eval()
B = cond_cont.size(0)
device = cond_cont.device
@@ -46,7 +107,9 @@ def sample_ddpm(
x = (1.0 / alpha.sqrt()) * (
x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred
) + beta.sqrt() * z
return x
n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
n_sec_pred = n_sec_logits.argmax(dim=-1)
return x, n_sec_pred
@torch.no_grad()
@@ -56,8 +119,8 @@ def sample_ddim(
cond_cat: torch.Tensor,
schedule,
steps: int = 50,
) -> torch.Tensor:
"""DDIM deterministic sampling (Song et al. 2020) with `steps` substeps."""
) -> tuple[torch.Tensor, torch.Tensor]:
"""DDIM deterministic sampling (Song et al. 2020). Returns (sample, n_sec_pred)."""
model.eval()
B = cond_cont.size(0)
device = cond_cont.device
@@ -75,4 +138,6 @@ def sample_ddim(
ab_prev = torch.ones(1, device=device)
x0_pred = (x - (1.0 - ab_t).sqrt() * eps_pred) / ab_t.sqrt()
x = ab_prev.sqrt() * x0_pred + (1.0 - ab_prev).sqrt() * eps_pred
return x
n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
n_sec_pred = n_sec_logits.argmax(dim=-1)
return x, n_sec_pred
giant/train.py
+172 -41
View File
@@ -8,14 +8,31 @@ from types import FrameType
from typing import Callable
import torch
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from tqdm import tqdm
from giant.model.schedule import CosineSchedule, flow_matching_loss
from giant.model.schedule import (
CosineSchedule,
flow_matching_loss,
flow_matching_loss_secondary,
)
from giant.validate import validate_marginals
_METRICS_FIELDS = ["epoch", "train_loss", "val_loss", "lr", "epoch_time_s"]
_METRICS_FIELDS = [
"epoch",
"train_loss",
"train_loss_s1",
"train_loss_nsec",
"train_loss_s2",
"val_loss",
"val_loss_s1",
"val_loss_nsec",
"val_loss_s2",
"lr",
"epoch_time_s",
]
_CATCHABLE_SIGNALS = (signal.SIGINT, signal.SIGTERM)
@@ -57,8 +74,89 @@ 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)
def _compute_losses(
stage1_model: torch.nn.Module,
sec_decoder: torch.nn.Module,
batch: tuple,
mode: str,
ddpm_schedule,
device: torch.device,
lambda_nsec: float,
lambda_s2: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Compute (total_loss, L_s1, L_nsec, L_s2) for one batch."""
cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_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)
# Stage-1 flow loss
if mode == "flow":
l_s1 = flow_matching_loss(stage1_model, x1_s1, cond_cont, cond_cat)
else:
assert ddpm_schedule is not None
l_s1 = ddpm_schedule.loss(stage1_model, x1_s1, cond_cont, cond_cat)
# n_sec classification loss
n_sec_logits = stage1_model.predict_n_sec(cond_cont, cond_cat)
l_nsec = F.cross_entropy(n_sec_logits, n_sec)
# 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.
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)
sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
l_s2 = flow_matching_loss_secondary(
sec_decoder,
x1_s2,
cond_cont,
cond_cat,
x1_s1.detach(),
sec_mask,
)
total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2
return total, l_s1, l_nsec, l_s2
def train(
model: torch.nn.Module,
stage1_model: torch.nn.Module,
sec_decoder: torch.nn.Module,
train_loader: DataLoader,
val_loader: DataLoader,
mode: str,
@@ -67,6 +165,8 @@ def train(
warmup_epochs: int,
device: torch.device,
out_dir: str | Path,
lambda_nsec: float = 0.1,
lambda_s2: float = 1.0,
normalizer_dict: dict | None = None,
pdg_map: dict | None = None,
mat_map: dict | None = None,
@@ -79,8 +179,11 @@ def train(
out_dir = Path(out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
model = model.to(device)
optimizer = optim.AdamW(model.parameters(), lr=lr)
stage1_model = stage1_model.to(device)
sec_decoder = sec_decoder.to(device)
all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters())
optimizer = optim.AdamW(all_params, lr=lr)
def _lr_lambda(epoch: int) -> float:
if warmup_epochs > 0 and epoch < warmup_epochs:
@@ -97,7 +200,8 @@ def train(
best_val_loss = float("inf")
if resume_path is not None:
ckpt = torch.load(resume_path, map_location=device, weights_only=False)
model.load_state_dict(ckpt["model"])
stage1_model.load_state_dict(ckpt["model"])
sec_decoder.load_state_dict(ckpt["sec_decoder"])
optimizer.load_state_dict(ckpt["optimizer"])
lr_sched.load_state_dict(ckpt["lr_sched"])
start_epoch = ckpt.get("epoch", 0) + 1
@@ -135,8 +239,12 @@ def train(
for epoch in range(start_epoch, epochs + 1):
epoch_start = time.monotonic()
current_lr = optimizer.param_groups[0]["lr"]
model.train()
stage1_model.train()
sec_decoder.train()
train_loss_sum = 0.0
train_s1_sum = 0.0
train_nsec_sum = 0.0
train_s2_sum = 0.0
train_n = 0
ema_loss = 0.0
bar = tqdm(
@@ -147,28 +255,31 @@ def train(
unit="batch",
dynamic_ncols=True,
)
for cond_cont, cond_cat, x1 in bar:
cond_cont = cond_cont.to(device)
cond_cat = cond_cat.to(device)
x1 = x1.to(device)
if mode == "flow":
loss = flow_matching_loss(model, x1, cond_cont, cond_cat)
else:
assert ddpm_schedule is not None
loss = ddpm_schedule.loss(model, x1, cond_cont, cond_cat)
for batch in bar:
loss, l_s1, l_nsec, l_s2 = _compute_losses(
stage1_model,
sec_decoder,
batch,
mode,
ddpm_schedule,
device,
lambda_nsec,
lambda_s2,
)
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
torch.nn.utils.clip_grad_norm_(all_params, 1.0)
optimizer.step()
B = batch[0].size(0)
batch_loss = loss.item()
train_loss_sum += batch_loss * x1.size(0)
train_n += x1.size(0)
train_loss_sum += batch_loss * B
train_s1_sum += l_s1.item() * B
train_nsec_sum += l_nsec.item() * B
train_s2_sum += l_s2.item() * B
train_n += B
ema_loss = (
batch_loss
if train_n == x1.size(0)
else 0.95 * ema_loss + 0.05 * batch_loss
batch_loss if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss
)
bar.set_postfix_str(f"loss={ema_loss:.4f}", refresh=False)
@@ -177,28 +288,36 @@ def train(
bar.close()
if shutdown.requested:
# Mid-epoch: discard the partial epoch rather than persist an
# inconsistent (lr_sched not stepped, no validation) checkpoint.
break
train_loss = train_loss_sum / max(train_n, 1)
lr_sched.step()
model.eval()
stage1_model.eval()
sec_decoder.eval()
val_loss_sum = 0.0
val_s1_sum = 0.0
val_nsec_sum = 0.0
val_s2_sum = 0.0
val_n = 0
with torch.no_grad():
for cond_cont, cond_cat, x1 in val_loader:
cond_cont = cond_cont.to(device)
cond_cat = cond_cat.to(device)
x1 = x1.to(device)
if mode == "flow":
loss = flow_matching_loss(model, x1, cond_cont, cond_cat)
else:
assert ddpm_schedule is not None
loss = ddpm_schedule.loss(model, x1, cond_cont, cond_cat)
val_loss_sum += loss.item() * x1.size(0)
val_n += x1.size(0)
for batch in val_loader:
loss, l_s1, l_nsec, l_s2 = _compute_losses(
stage1_model,
sec_decoder,
batch,
mode,
ddpm_schedule,
device,
lambda_nsec,
lambda_s2,
)
B = batch[0].size(0)
val_loss_sum += loss.item() * B
val_s1_sum += l_s1.item() * B
val_nsec_sum += l_nsec.item() * B
val_s2_sum += l_s2.item() * B
val_n += B
val_loss = val_loss_sum / max(val_n, 1)
epoch_time = time.monotonic() - epoch_start
@@ -206,14 +325,24 @@ def train(
marker = " [best]" if is_best else ""
print(
f"epoch {epoch:{epoch_w}d}/{epochs}"
f" train {train_loss:.4f} val {val_loss:.4f}"
f" train {train_loss:.4f}"
f" (s1={train_s1_sum / max(train_n, 1):.3f}"
f" nsec={train_nsec_sum / max(train_n, 1):.3f}"
f" s2={train_s2_sum / max(train_n, 1):.3f})"
f" val {val_loss:.4f}"
f" lr {current_lr:.2e} {epoch_time:.1f}s{marker}"
)
metrics_writer.writerow(
{
"epoch": epoch,
"train_loss": train_loss,
"train_loss_s1": train_s1_sum / max(train_n, 1),
"train_loss_nsec": train_nsec_sum / max(train_n, 1),
"train_loss_s2": train_s2_sum / max(train_n, 1),
"val_loss": val_loss,
"val_loss_s1": val_s1_sum / max(val_n, 1),
"val_loss_nsec": val_nsec_sum / max(val_n, 1),
"val_loss_s2": val_s2_sum / max(val_n, 1),
"lr": current_lr,
"epoch_time_s": epoch_time,
}
@@ -223,16 +352,18 @@ def train(
if validate_every > 0 and epoch % validate_every == 0:
print(f"[epoch {epoch}] marginal validation:")
validate_marginals(
model,
stage1_model,
val_loader,
mode=mode,
schedule=ddpm_schedule,
device=device,
steps=validate_steps,
sec_decoder=sec_decoder,
)
ckpt: dict = {
"model": model.state_dict(),
"model": stage1_model.state_dict(),
"sec_decoder": sec_decoder.state_dict(),
"optimizer": optimizer.state_dict(),
"lr_sched": lr_sched.state_dict(),
"epoch": epoch,
giant/validate.py
+140 -8
View File
@@ -2,8 +2,14 @@ import numpy as np
import torch
from torch.utils.data import DataLoader
from giant.constants import LOCAL_TARGET_NAMES
from giant.sample import sample_flow, sample_ddpm, sample_ddim
from giant.constants import K_MAX, LOCAL_TARGET_NAMES
from giant.sample import (
sample_flow,
sample_ddpm,
sample_ddim,
sample_secondaries,
snap_type_to_pdg_idx,
)
def _kw(steps: int | None) -> dict[str, int]:
@@ -28,6 +34,13 @@ def _histogram_kl(
return float(np.sum(p * np.log(p / q)))
def _bincount_frac(x: np.ndarray, minlength: int) -> np.ndarray:
"""Fraction of samples per integer value in [0, minlength), as a distribution."""
counts = np.bincount(x, minlength=minlength)[:minlength].astype(np.float64)
total = counts.sum()
return counts / total if total > 0 else counts
def validate_marginals(
model: torch.nn.Module,
val_loader: DataLoader,
@@ -37,7 +50,8 @@ def validate_marginals(
n_batches: int | None = None,
kl_bins: int = 50,
steps: int | None = None,
) -> dict[str, np.ndarray]:
sec_decoder: torch.nn.Module | None = None,
) -> dict[str, np.ndarray | float]:
"""Compare per-dimension marginals of generated vs. real steps.
Returns {"real": (N,9), "generated": (N,9), "kl_divergence": (9,)} in
@@ -47,28 +61,80 @@ def validate_marginals(
`steps` overrides the number of sampler steps (flow ODE steps or DDIM
substeps); `None` keeps each sampler's own default. Unused in "ddpm"
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.
"""
if device is None:
device = next(model.parameters()).device
model.eval()
if sec_decoder is not None:
sec_decoder.eval()
all_real, all_gen = [], []
for i, (cond_cont, cond_cat, x1) in enumerate(val_loader):
all_n_sec_real, all_n_sec_pred = [], []
all_species_real, all_species_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).
cond_cont, cond_cat, x1, n_sec, sec_cont, sec_pdg_idx = batch
cond_cont = cond_cont.to(device)
cond_cat = cond_cat.to(device)
if mode == "flow":
gen = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
gen, n_sec_pred = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
elif mode == "ddpm":
gen = sample_ddpm(model, cond_cont, cond_cat, schedule)
gen, n_sec_pred = sample_ddpm(model, cond_cont, cond_cat, schedule)
else:
gen = sample_ddim(model, cond_cont, cond_cat, schedule, **_kw(steps))
gen, n_sec_pred = sample_ddim(
model, cond_cont, cond_cat, schedule, **_kw(steps)
)
all_real.append(x1.numpy())
all_gen.append(gen.cpu().numpy())
if sec_decoder is None:
continue
n_sec_pred_np = n_sec_pred.cpu().numpy()
n_sec_np = n_sec.numpy()
all_n_sec_real.append(n_sec_np)
all_n_sec_pred.append(n_sec_pred_np)
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()
sec_cont_pred, sec_type_emb, sec_valid_pred = sample_secondaries(
sec_decoder,
cond_cont,
cond_cat,
gen,
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_valid = sec_valid_pred.cpu().numpy()
all_species_real.append(real_species[real_valid])
all_species_gen.append(gen_species[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])
real = np.concatenate(all_real, axis=0)
generated = np.concatenate(all_gen, axis=0)
@@ -92,4 +158,70 @@ def validate_marginals(
f"{r.std():>10.4f} {g.std():>10.4f} {kl_divergence[j]:>14.4f}"
)
return {"real": real, "generated": generated, "kl_divergence": kl_divergence}
result: dict[str, np.ndarray | float] = {
"real": real,
"generated": generated,
"kl_divergence": kl_divergence,
}
if sec_decoder is None:
return result
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)
n_classes = (
max(int(species_real.max(initial=0)), int(species_gen.max(initial=0))) + 1
)
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(
f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f} "
f"mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}"
)
n_sec_dist_header = f"{'n_sec value':<20} {'real_frac':>10} {'gen_frac':>10}"
print(n_sec_dist_header)
print("-" * len(n_sec_dist_header))
max_n_sec = max(int(n_sec_real.max()), int(n_sec_pred_all.max())) + 1
real_n_sec_dist = _bincount_frac(n_sec_real, max_n_sec)
gen_n_sec_dist = _bincount_frac(n_sec_pred_all, max_n_sec)
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 slot (energy frac.)':<24} {'real_mean':>10} {'gen_mean':>10} "
f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
)
print("-" * 90)
for j in range(K_MAX):
r = np.concatenate(all_frac_real[j]) if all_frac_real[j] else np.array([])
g = np.concatenate(all_frac_gen[j]) if all_frac_gen[j] else np.array([])
if len(r) == 0 or len(g) == 0:
continue
kl = _histogram_kl(r, g, bins=kl_bins)
energy_fraction_kl[j] = kl
print(
f"{j:<24} {r.mean():>10.4f} {g.mean():>10.4f} "
f"{r.std():>10.4f} {g.std():>10.4f} {kl:>14.4f}"
)
result.update(
{
"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,
"energy_fraction_kl": energy_fraction_kl,
}
)
return result
pyproject.toml
+4 -1
View File
@@ -24,7 +24,10 @@ dev = [
"pytest>=8,<10",
"ruff>=0.15,<1",
"ty>=0.0.50,<0.1",
"giant[convert,analysis]",
"giant[convert,analysis,geometry]",
]
geometry = [
"scikit-learn>=1.4,<2",
]
convert = [
"uproot>=5.3,<6",
scripts/create_root_files.py
+55 -20
View File
@@ -3,15 +3,17 @@ run_pbwo4, run_sampling) and filing the output into the dataset's raw/ tree:
raw/<kind>/<gen>/<detector>/shard-NNN.root
These executables take `[configName] nEvents` and always write a fixed-name
*.root file into the current directory — so running several in parallel
needs separate working directories, and the output filename has to be
discovered rather than assumed (it differs per executable: run_pbwo4 writes
pbwo4_<n>events_hits.root, run_sampling writes sampling_<config>_<n>events_hits.root,
others may differ again). This script gives each run its own scratch
directory under <dataset-root>/.sim-tmp/, requires exactly one *.root to
appear there, and moves it to the next free shard index for that detector
(existing shards are never overwritten).
These executables take `[configName] nEvents [energy_GeV]` (configName is
only accepted by executables with a config selector, e.g. run_sampling;
energy_GeV defaults to 1.0 in the executable itself if omitted here) and
always write a fixed-name *.root file into the current directory — so
running several in parallel needs separate working directories, and the
output filename has to be discovered rather than assumed (it differs per
executable: run_pbwo4 writes pbwo4_<n>events_hits.root, run_sampling writes
sampling_<config>_<n>events_hits.root, others may differ again). This script
gives each run its own scratch directory under <dataset-root>/.sim-tmp/,
requires exactly one *.root to appear there, and moves it to the next free
shard index for that detector (existing shards are never overwritten).
--gen must already exist under raw/<kind>/ — create one first with
`dwarf bump-gen`.
@@ -25,6 +27,7 @@ import shutil
import subprocess
import sys
import uuid
import zlib
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
@@ -104,10 +107,41 @@ def plan_jobs(
return jobs
def job_seed(kind: str, gen: str, job: SimJob, energy_gev: float | None) -> int:
"""Deterministic RNG seed for one sim job, unique per (kind, gen, detector, config, shard, energy).
Jobs run concurrently (ThreadPoolExecutor below) and can start within the
same wall-clock second; minicalosim's default seed falls back to
time(NULL) in that case, so two concurrently-launched jobs can silently
get identical RNG state and produce byte-identical physics despite
landing in separate shard files. Deriving the seed from the full job
identity instead keeps it both unique and reproducible.
"""
key = (
f"{kind}|{gen}|{job.detector}|{job.config or ''}|{job.shard_index}"
f"|{energy_gev if energy_gev is not None else ''}"
)
return zlib.crc32(key.encode()) & 0x7FFFFFFF
def build_cmd(
executable: Path, job: SimJob, events_per_file: int, energy_gev: float | None
) -> list[str]:
"""minicalosim executables take positional `[configName] nEvents [energy_GeV]`."""
cmd = [str(executable)]
if job.config:
cmd.append(job.config)
cmd.append(str(events_per_file))
if energy_gev is not None:
cmd.append(str(energy_gev))
return cmd
def run_job(
job: SimJob,
executable: Path,
events_per_file: int,
energy_gev: float | None,
dataset_root: Path,
kind: str,
gen: str,
@@ -119,12 +153,10 @@ def run_job(
)
workdir.mkdir(parents=True)
cmd = [str(executable)]
if job.config:
cmd.append(job.config)
cmd.append(str(events_per_file))
cmd = build_cmd(executable, job, events_per_file, energy_gev)
result = subprocess.run(cmd, cwd=workdir, capture_output=True, text=True)
env = dict(os.environ, MINICALOSIM_SEED=str(job_seed(kind, gen, job, energy_gev)))
result = subprocess.run(cmd, cwd=workdir, capture_output=True, text=True, env=env)
if result.returncode != 0:
return JobResult(
@@ -178,6 +210,7 @@ def run_all(
jobs: list[SimJob],
executable: Path,
events_per_file: int,
energy_gev: float | None,
dataset_root: Path,
kind: str,
gen: str,
@@ -192,6 +225,7 @@ def run_all(
job,
executable,
events_per_file,
energy_gev,
dataset_root,
kind,
gen,
@@ -222,6 +256,7 @@ def run_make_root(
dataset_root: str,
jobs: int,
execute: bool,
energy_gev: float | None = None,
) -> None:
if jobs < 1:
raise SystemExit("error: --jobs must be >= 1")
@@ -229,6 +264,8 @@ def run_make_root(
raise SystemExit("error: --num-files must be >= 1")
if events_per_file < 1:
raise SystemExit("error: --events-per-file must be >= 1")
if energy_gev is not None and energy_gev <= 0:
raise SystemExit("error: --energy-gev must be > 0")
if not executable.is_file() or not os.access(executable, os.X_OK):
raise SystemExit(f"error: {executable} is not an executable file")
@@ -241,11 +278,7 @@ def run_make_root(
print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
print(f"executable: {executable}")
for job in planned_jobs:
cmd = (
[str(executable)]
+ ([job.config] if job.config else [])
+ [str(events_per_file)]
)
cmd = build_cmd(executable, job, events_per_file, energy_gev)
dest = (
dataset_root_path
/ "raw"
@@ -254,7 +287,8 @@ def run_make_root(
/ job.detector
/ f"shard-{job.shard_index:03d}.root"
)
print(f" {' '.join(cmd)} -> {dest}")
seed = job_seed(kind, gen, job, energy_gev)
print(f" MINICALOSIM_SEED={seed} {' '.join(cmd)} -> {dest}")
if not execute:
print("\nDry run only — pass --execute to apply.")
@@ -266,6 +300,7 @@ def run_make_root(
planned_jobs,
executable,
events_per_file,
energy_gev,
dataset_root_path,
kind,
gen,
scripts/dwarf.py
+87 -1
View File
@@ -20,6 +20,7 @@ from scripts.bump_dataset_version import (
run_update_manifest,
)
from scripts.create_root_files import run_make_root
from scripts.geometry_oracle import run_build_geometry_oracle
from scripts.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
from scripts.migrate_geant_steps import run_migration
from scripts.steps_to_parquet import convert_steps_to_parquet
@@ -115,14 +116,21 @@ def convert(
"error: --output can only be used with a single input file", err=True
)
raise typer.Exit(1)
total_orphaned = 0
for root_file in root_files:
convert_steps_to_parquet(
_, n_orphaned = convert_steps_to_parquet(
root_file,
output_path=output,
batch_size=batch_size,
tree_name=tree,
compression=compression_value,
)
total_orphaned += n_orphaned
if total_orphaned:
typer.echo(
f"\n{total_orphaned} orphaned child track(s) dropped across "
f"{len(root_files)} file(s)."
)
return
if output is not None:
@@ -355,6 +363,16 @@ def make_root(
gen: Annotated[
str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")
],
energy_gev: Annotated[
float | None,
typer.Option(
"--energy-gev",
help="energy_GeV passed to the executable (default: executable's own "
"default, currently 1.0). Note the dataset detector label is not "
"derived from this — e.g. use '--detector pbwo4_10gev --energy-gev 10' "
"to name the dataset accordingly.",
),
] = None,
kind: Annotated[
str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
] = "steps",
@@ -382,6 +400,74 @@ def make_root(
dataset_root=str(dataset_root),
jobs=jobs,
execute=execute,
energy_gev=energy_gev,
)
class OracleMethod(str, Enum):
slab = "slab"
knn = "knn"
svm = "svm"
@app.command("build-geometry-oracle")
def build_geometry_oracle(
data: Annotated[
Path, typer.Argument(help="Steps parquet file or directory of steps files")
],
out: Annotated[Path, typer.Option("--out", "-o", help="Output oracle .pkl path")],
method: Annotated[
OracleMethod,
typer.Option(
"--method",
help=(
"Lookup strategy: slab (default; exact O(log #segments) fast "
"path for the layered-slab detector geometry), knn, or svm "
"(generic fallbacks for non-slab geometries)"
),
),
] = OracleMethod.slab,
k: Annotated[
int, typer.Option("--k", help="Neighbours for the knn classifier")
] = 1,
subsample: Annotated[
int,
typer.Option("--subsample", help="Max reference points sampled from the data"),
] = 500_000,
escape_factor: Annotated[
float,
typer.Option(
"--escape-factor",
help="escape_threshold = this x median spacing of reference points",
),
] = 5.0,
seed: Annotated[int, typer.Option("--seed", help="Sampling seed")] = 0,
depth_axis: Annotated[
int,
typer.Option(
"--depth-axis",
help="0/1/2 -> x/y/z axis the layers stack along (method=slab only)",
),
] = 2,
n_bins: Annotated[
int,
typer.Option(
"--n-bins",
help="Depth-axis resolution, finer than the thinnest layer (method=slab only)",
),
] = 2000,
) -> None:
"""Fit a position -> (material, layer_id) oracle for `giant rollout`."""
run_build_geometry_oracle(
data=data,
out=out,
method=method.value,
k=k,
subsample=subsample,
escape_factor=escape_factor,
seed=seed,
depth_axis=depth_axis,
n_bins=n_bins,
)
scripts/geometry_oracle.py
+74
View File
@@ -0,0 +1,74 @@
"""Build a position -> (material, layer_id) geometry oracle from steps parquet.
Backs the `dwarf build-geometry-oracle` subcommand. The oracle is consumed by
`giant rollout` to supply the material/layer conditioning at each step, since the
surrogate does not predict them.
"""
from __future__ import annotations
from pathlib import Path
from giant.data.loader import find_parquet_files
from giant.geometry import build_geometry_oracle
def run_build_geometry_oracle(
data: Path,
out: Path,
method: str = "slab",
k: int = 1,
subsample: int = 500_000,
escape_factor: float = 5.0,
seed: int = 0,
depth_axis: int = 2,
n_bins: int = 2000,
) -> None:
files = find_parquet_files(data)
print(f"found {len(files)} parquet file(s); sampling up to {subsample:,} points")
oracle = build_geometry_oracle(
files,
method=method,
k=k,
subsample=subsample,
escape_factor=escape_factor,
seed=seed,
depth_axis=depth_axis,
n_bins=n_bins,
)
print(
f"method: {method} reference points: {oracle.metadata['n_reference_points']:,}"
)
print("classes (material, layer_id):")
for material, layer_id in oracle.classes:
print(f" {material:<12} layer_id={layer_id}")
if method == "slab":
z_lo, z_hi = oracle.metadata["z_range"]
print(
f"depth axis: {'xyz'[depth_axis]} segments: {oracle.metadata['n_segments']} "
f"z range: [{z_lo:.3f}, {z_hi:.3f}] radius_max: {oracle.metadata['radius_max']:.3f}"
)
print(
f"median depth spacing: {oracle.metadata['median_z_spacing']:.3f} "
f"escape_threshold: {oracle.escape_threshold:.3f} (= {escape_factor}x spacing)"
)
else:
print(
f"median NN spacing: {oracle.metadata['median_nn_dist']:.3f} "
f"escape_threshold: {oracle.escape_threshold:.3f} "
f"(= {escape_factor}x spacing)"
)
if oracle.escape_threshold <= 0.0:
print(
"warning: escape_threshold is 0 (reference points are coincident) — "
"every rollout query would be flagged as escaped. Pass "
"`giant rollout --escape-threshold <mm>` to override, or use data "
"with distinct step positions."
)
out.parent.mkdir(parents=True, exist_ok=True)
oracle.save(out)
print(f"wrote oracle -> {out}")
scripts/steps_to_parquet.py
+89 -23
View File
@@ -4,7 +4,7 @@ See `uv run dwarf convert --help` for the CLI.
"""
from pathlib import Path
from typing import Literal
from typing import Literal, cast
import awkward as ak
import polars as pl
@@ -13,41 +13,101 @@ import uproot
ParquetCompression = Literal["lz4", "uncompressed", "snappy", "gzip", "brotli", "zstd"]
def _add_secondary_energy(df: pl.DataFrame) -> pl.DataFrame:
"""Add per-step `e_sec`: total initial kinetic energy of the secondaries born in it.
def _add_secondary_attributes(df: pl.DataFrame) -> tuple[pl.DataFrame, int]:
"""Add per-step secondary attributes via the parent→child track join.
Each secondary's creation energy is the `pre_E` of that child track's first step
(min `step_no`) in the same event, so for a parent step
`e_sec = Σ over child_track_ids of the child track's first-step pre_E`. Steps that
spawn nothing get 0.0. The full event must be present in `df` (it is — the writer
concatenates every batch before this runs), since a child track's first step can
live in a different read batch than its parent step.
For each step that spawns secondaries, collects each child track's birth
state (from the child track's first step in the same event) and emits:
e_sec float64 — total secondary energy (sum of child first-step pre_E)
sec_E_list list[f64] — per-secondary energy, sorted descending
sec_pdg_list list[i32] — per-secondary PDG code, same order
sec_dx_list list[f64] — per-secondary birth direction x, same order
sec_dy_list list[f64] — per-secondary birth direction y, same order
sec_dz_list list[f64] — per-secondary birth direction z, same order
Steps with no children get 0.0 / empty lists. The full event must be
present in `df` (it is — the writer concatenates before calling this).
A listed child_track_id can fail to match any row in `first_step` — the
child track never took a recorded step (e.g. absorbed below the tracking
threshold at birth). Such orphans carry no physical secondary data, so
they're dropped from child_track_ids/sec_*_list rather than left as nulls:
a null in a float32 list silently becomes NaN once the parquet round-trips
through the loader (`giant/data/loader.py:_pad_list_col`), and that NaN
poisons every later secondary slot in the same step via the cumulative-sum
"remaining budget" in `encode_secondaries`.
Returns (df, n_orphaned) — the caller uses the count to report/aggregate
across files rather than relying solely on the printed message here.
"""
first_E = (
first_step = (
df.sort("step_no")
.group_by(["event_id", "track_id"])
.agg(pl.col("pre_E").first().alias("child_E"))
.agg(
pl.col("pre_E").first().alias("child_E"),
pl.col("pdg").first().alias("child_pdg"),
pl.col("pre_dx").first().alias("child_dx"),
pl.col("pre_dy").first().alias("child_dy"),
pl.col("pre_dz").first().alias("child_dz"),
)
.rename({"track_id": "child_track_id"})
)
child_track_id_dtype = cast(pl.List, df.schema["child_track_ids"]).inner
exploded = (
df.select(["event_id", "child_track_ids"])
.with_row_index("_step_row")
.explode("child_track_ids")
.rename({"child_track_ids": "child_track_id"})
.drop_nulls("child_track_id") # steps with no children explode to a null row
.drop_nulls("child_track_id")
)
summed = (
exploded.join(first_E, on=["event_id", "child_track_id"], how="left")
joined = exploded.join(first_step, on=["event_id", "child_track_id"], how="left")
n_orphaned = joined["child_E"].null_count()
if n_orphaned:
print(
f" dropping {n_orphaned} orphaned child_track_id(s) with no "
"recorded first step (absorbed below tracking threshold?)"
)
joined = joined.drop_nulls("child_E")
# Sort each step's secondaries by descending energy, then aggregate into lists
per_step = (
joined.sort("child_E", descending=True)
.group_by("_step_row")
.agg(pl.col("child_E").sum().alias("e_sec"))
.agg(
pl.col("child_track_id").alias("child_track_ids"),
pl.col("child_E").sum().alias("e_sec"),
pl.col("child_E").alias("sec_E_list"),
pl.col("child_pdg").alias("sec_pdg_list"),
pl.col("child_dx").alias("sec_dx_list"),
pl.col("child_dy").alias("sec_dy_list"),
pl.col("child_dz").alias("sec_dz_list"),
)
)
return (
df.with_row_index("_step_row")
.join(summed, on="_step_row", how="left")
.with_columns(pl.col("e_sec").fill_null(0.0).cast(pl.Float64))
empty_list_f64 = pl.Series("x", [[]], dtype=pl.List(pl.Float64))
empty_list_i32 = pl.Series("x", [[]], dtype=pl.List(pl.Int32))
empty_list_child_id = pl.Series("x", [[]], dtype=pl.List(child_track_id_dtype))
out = (
df.drop("child_track_ids")
.with_row_index("_step_row")
.join(per_step, on="_step_row", how="left")
.with_columns(
pl.col("child_track_ids").fill_null(empty_list_child_id),
pl.col("e_sec").fill_null(0.0).cast(pl.Float64),
pl.col("sec_E_list").fill_null(empty_list_f64),
pl.col("sec_pdg_list").fill_null(empty_list_i32),
pl.col("sec_dx_list").fill_null(empty_list_f64),
pl.col("sec_dy_list").fill_null(empty_list_f64),
pl.col("sec_dz_list").fill_null(empty_list_f64),
)
.drop("_step_row")
)
return out, n_orphaned
def _batch_to_polars(batch: ak.Array) -> pl.DataFrame:
@@ -73,7 +133,7 @@ def convert_steps_to_parquet(
batch_size: str = "100 MB",
tree_name: str = "Steps",
compression: ParquetCompression = "snappy",
) -> Path:
) -> tuple[Path, int]:
"""Read *tree_name* from *root_path* and write it to a Parquet file.
Reads in batches of *batch_size* so that peak ROOT-deserialization memory
@@ -89,6 +149,11 @@ def convert_steps_to_parquet(
integer row count (500_000).
tree_name: Name of the TTree inside the ROOT file.
compression: Parquet compression codec (snappy | lz4 | zstd | gzip | none).
Returns (output_path, n_orphaned) — n_orphaned is the count of dropped
orphaned child_track_ids (see `_add_secondary_attributes`), 0 if the tree
has no child_track_ids column at all. Callers converting many files use
it to aggregate a total instead of grepping the printed per-file message.
"""
root_path = Path(root_path)
if output_path is None:
@@ -111,11 +176,12 @@ def convert_steps_to_parquet(
df = pl.concat(batches)
# Steps tree carries the parent→child links needed to derive secondary energy;
# other trees (e.g. Hits) don't, so only augment when the column is present.
n_orphaned = 0
if "child_track_ids" in df.columns:
print("\nComputing per-step secondary energy (e_sec) …", end=" ", flush=True)
df = _add_secondary_energy(df)
print("\nComputing per-step secondary attributes …", end=" ", flush=True)
df, n_orphaned = _add_secondary_attributes(df)
print(f"\nWriting {output_path} …", end=" ", flush=True)
df.write_parquet(output_path, compression=compression)
print(f"done ({output_path.stat().st_size / 1e6:.1f} MB)")
return output_path
return output_path, n_orphaned
scripts/steps_to_parquet_parallel.py
+11
View File
@@ -26,6 +26,12 @@ from pathlib import Path
GEN_RE = re.compile(r"^gen\d+$")
SCHEMA_RE = re.compile(r"^schema(\d+)$")
# Matches the per-file orphan-drop message printed by
# steps_to_parquet._add_secondary_attributes — each subprocess's count is
# parsed back out of its captured stdout since there's no in-process return
# value across the subprocess boundary.
_ORPHAN_RE = re.compile(r"dropping (\d+) orphaned child_track_id")
class DestinationError(ValueError):
pass
@@ -210,4 +216,9 @@ def run_parallel_job(
print(f" {root_file}", file=sys.stderr)
raise SystemExit(1)
total_orphaned = sum(
int(m.group(1)) for _, _, stdout, _ in results for m in _ORPHAN_RE.finditer(stdout)
)
if total_orphaned:
print(f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s).")
print(f"\nAll {len(results)} conversion(s) completed.")
tests/test_analysis.py
+526
View File
@@ -12,11 +12,14 @@ from giant.analysis import (
RAW_TARGET_NAMES,
SampleCollection,
compute_event_observables_pl,
compute_rollout_observables,
compute_truth_observables,
constraint_report,
constraint_report_pl,
correlation_matrices,
direction_alignment,
load_predicted_local,
load_rollout_vs_truth,
marginal_table,
marginal_table_pl,
pdg_contribution_table_pl,
@@ -30,6 +33,9 @@ from giant.analysis import (
plot_pairwise,
plot_pdg_energy_share,
plot_pdg_length_share,
plot_rollout_longitudinal,
plot_rollout_total_energy,
plot_rollout_transverse,
plot_shower_max_depth,
plot_total_energy,
plot_total_length,
@@ -40,12 +46,15 @@ from giant.constants import (
PREDICT_COORD_METADATA_KEY,
PREDICT_SCHEMA_VERSION,
PREDICT_SCHEMA_VERSION_KEY,
ROLLOUT_COORD_VALUE,
)
from giant.data.transforms import (
energy_simplex_decode,
inv_log_transform,
local_frame_rotation,
log_transform,
reconstruct_post_pos,
travel_direction,
)
@@ -660,3 +669,520 @@ def test_plot_pdg_energy_share_caps_slices():
fig = plot_pdg_energy_share(table, max_slices=4)
for ax in fig.axes:
assert len(ax.patches) == 4
# ---------------------------------------------------------------------------
# load_rollout_vs_truth: unpaired rollout-vs-truth SampleCollection
# ---------------------------------------------------------------------------
def _make_world_frame_physical(rng, n):
"""Random-but-physical pre/post step fields shared by the truth/rollout schemas."""
pre_pos = rng.uniform(-5.0, 5.0, (n, 3)).astype(np.float32)
pre_dir = _unit_vectors(rng, n)
pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)
step_length = rng.uniform(0.1, 5.0, n).astype(np.float32)
travel_dir_world = _unit_vectors(rng, n)
post_pos = pre_pos + step_length[:, None] * travel_dir_world
post_dir_world = _unit_vectors(rng, n)
delta_e = (rng.uniform(0.0, 1.0, n) * pre_E).astype(np.float32)
post_E = pre_E - delta_e
edep = (delta_e * rng.uniform(0.0, 1.0, n)).astype(np.float32)
return pre_pos, pre_dir, pre_E, step_length, post_pos, post_dir_world, post_E, edep
def _expected_raw9(
pre_pos, pre_dir, pre_E, step_length, post_pos, post_dir_world, post_E, edep
):
post_dir_local = local_frame_rotation(pre_dir, post_dir_world)
travel_dir_local = local_frame_rotation(
pre_dir, travel_direction(pre_pos, post_pos)
)
return np.column_stack(
[step_length, pre_E - post_E, edep, post_dir_local, travel_dir_local]
).astype(np.float32)
def _write_truth_parquet(path, n=200, seed=0):
rng = np.random.default_rng(seed)
fields = _make_world_frame_physical(rng, n)
pre_pos, pre_dir, pre_E, step_length, post_pos, post_dir_world, post_E, edep = (
fields
)
table = pa.table(
{
"event_id": rng.integers(0, 20, n),
"pdg": rng.choice([11, -11, 22], n),
"pre_x": pre_pos[:, 0],
"pre_y": pre_pos[:, 1],
"pre_z": pre_pos[:, 2],
"pre_E": pre_E,
"pre_dx": pre_dir[:, 0],
"pre_dy": pre_dir[:, 1],
"pre_dz": pre_dir[:, 2],
"material": rng.choice(["W", "Pb"], n),
"layer_id": rng.integers(0, 10, n).astype(np.int32),
"child_track_ids": [list(range(int(k))) for k in rng.integers(0, 3, n)],
"e_sec": rng.uniform(0.0, 1.0, n).astype(np.float32),
"step_length": step_length,
"post_E": post_E,
"edep": edep,
"post_dx": post_dir_world[:, 0],
"post_dy": post_dir_world[:, 1],
"post_dz": post_dir_world[:, 2],
"post_x": post_pos[:, 0],
"post_y": post_pos[:, 1],
"post_z": post_pos[:, 2],
}
)
pq.write_table(table, path)
return _expected_raw9(*fields)
def _write_rollout_parquet(path, n=150, seed=1, coord=ROLLOUT_COORD_VALUE):
rng = np.random.default_rng(seed)
fields = _make_world_frame_physical(rng, n)
pre_pos, pre_dir, pre_E, step_length, post_pos, post_dir_world, post_E, edep = (
fields
)
table = pa.table(
{
"event_id": rng.integers(0, 20, n),
"track_id": rng.integers(0, 3, n),
"parent_id": np.full(n, -1, dtype=np.int64),
"generation": np.zeros(n, dtype=np.int64),
"step_no": np.zeros(n, dtype=np.int64),
"pdg": rng.choice([11, -11, 22], n),
"pre_x": pre_pos[:, 0],
"pre_y": pre_pos[:, 1],
"pre_z": pre_pos[:, 2],
"pre_E": pre_E,
"pre_dx": pre_dir[:, 0],
"pre_dy": pre_dir[:, 1],
"pre_dz": pre_dir[:, 2],
"post_x": post_pos[:, 0],
"post_y": post_pos[:, 1],
"post_z": post_pos[:, 2],
"post_E": post_E,
"post_dx": post_dir_world[:, 0],
"post_dy": post_dir_world[:, 1],
"post_dz": post_dir_world[:, 2],
"edep": edep,
"step_length": step_length,
"material": rng.choice(["W", "Pb"], n),
"layer_id": rng.integers(0, 10, n).astype(np.int32),
"n_sec_pred": rng.integers(0, 3, n).astype(np.int32),
# "" / "natural_end" mark a real generated step (the latter just
# additionally being a track's last); every row here is a real step,
# so all `n` should survive `load_rollout_vs_truth`'s filtering — see
# `test_load_rollout_vs_truth_drops_synthetic_termination_rows` for
# the escaped/unknown_pdg/energy_cutoff/max_steps bookkeeping rows.
"termination_reason": rng.choice(["", "natural_end"], n),
}
)
if coord is not None:
table = table.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: coord})
pq.write_table(table, path)
return _expected_raw9(*fields)
def test_load_rollout_vs_truth_decodes_raw_targets_correctly(tmp_path):
truth_path = tmp_path / "truth.parquet"
rollout_path = tmp_path / "rollout.parquet"
expected_real = _write_truth_parquet(truth_path, n=200, seed=0)
expected_gen = _write_rollout_parquet(rollout_path, n=150, seed=1)
samples = load_rollout_vs_truth(rollout_path, truth_path)
np.testing.assert_allclose(samples.real_raw, expected_real, atol=1e-4)
np.testing.assert_allclose(samples.gen_raw, expected_gen, atol=1e-4)
def test_load_rollout_vs_truth_allows_unpaired_lengths(tmp_path):
truth_path = tmp_path / "truth.parquet"
rollout_path = tmp_path / "rollout.parquet"
_write_truth_parquet(truth_path, n=200)
_write_rollout_parquet(rollout_path, n=150)
samples = load_rollout_vs_truth(rollout_path, truth_path)
assert samples.real_raw.shape == (200, 9)
assert samples.gen_raw.shape == (150, 9)
assert samples.pdg.shape == (200,)
assert samples.pdg_gen.shape == (150,)
@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
def test_load_rollout_vs_truth_downstream_plots_run_without_error(tmp_path, group_by):
truth_path = tmp_path / "truth.parquet"
rollout_path = tmp_path / "rollout.parquet"
_write_truth_parquet(truth_path, n=200)
_write_rollout_parquet(rollout_path, n=150)
samples = load_rollout_vs_truth(rollout_path, truth_path)
table = marginal_table(samples, group_by=group_by)
assert set(table["dim"]) == set(RAW_TARGET_NAMES)
assert plot_marginals(samples, group_by=group_by) is not None
assert plot_kl_bars(samples, group_by=group_by) is not None
def test_load_rollout_vs_truth_joint_and_constraint_checks_run(tmp_path):
truth_path = tmp_path / "truth.parquet"
rollout_path = tmp_path / "rollout.parquet"
_write_truth_parquet(truth_path, n=200)
_write_rollout_parquet(rollout_path, n=150)
samples = load_rollout_vs_truth(rollout_path, truth_path)
assert plot_correlation_matrices(samples) is not None
assert plot_pairwise(samples) is not None
assert plot_direction_alignment(samples) is not None
assert plot_constraint_violations(samples) is not None
assert constraint_report(samples) is not None
def test_load_rollout_vs_truth_drops_synthetic_termination_rows(tmp_path):
"""Bookkeeping rows for escaped/unknown_pdg/energy_cutoff/max_steps aren't steps.
`rollout.py._terminal_rows` writes one such row per track termination, with
`step_length=0`/`post_pos=pre_pos` and — for every reason but "escaped" —
the track's entire remaining `pre_E` dumped into `edep` so the shower's
total energy still conserves. Mixing these into the per-step comparison
would inject a spurious step_length=0 spike and roughly double the
apparent mean edep from bookkeeping alone, not model behavior (see the
`load_rollout_vs_truth` docstring). This checks they're excluded and the
real steps are decoded unaffected by their presence in the same file.
"""
rng = np.random.default_rng(3)
n_real, n_marker = 60, 40
real_fields = _make_world_frame_physical(rng, n_real)
pre_pos, pre_dir, pre_E, step_length, post_pos, post_dir_world, post_E, edep = (
real_fields
)
marker_pre_pos = rng.uniform(-5.0, 5.0, (n_marker, 3)).astype(np.float32)
marker_pre_dir = _unit_vectors(rng, n_marker)
marker_pre_E = rng.uniform(1.0, 100.0, n_marker).astype(np.float32)
def cat(real_col, marker_col):
return np.concatenate([real_col, marker_col])
n = n_real + n_marker
table = pa.table(
{
"event_id": rng.integers(0, 20, n),
"track_id": rng.integers(0, 3, n),
"parent_id": np.full(n, -1, dtype=np.int64),
"generation": np.zeros(n, dtype=np.int64),
"step_no": np.zeros(n, dtype=np.int64),
"pdg": rng.choice([11, -11, 22], n),
"pre_x": cat(pre_pos[:, 0], marker_pre_pos[:, 0]),
"pre_y": cat(pre_pos[:, 1], marker_pre_pos[:, 1]),
"pre_z": cat(pre_pos[:, 2], marker_pre_pos[:, 2]),
"pre_E": cat(pre_E, marker_pre_E),
"pre_dx": cat(pre_dir[:, 0], marker_pre_dir[:, 0]),
"pre_dy": cat(pre_dir[:, 1], marker_pre_dir[:, 1]),
"pre_dz": cat(pre_dir[:, 2], marker_pre_dir[:, 2]),
# Terminal markers: post_pos == pre_pos (zero-length "step").
"post_x": cat(post_pos[:, 0], marker_pre_pos[:, 0]),
"post_y": cat(post_pos[:, 1], marker_pre_pos[:, 1]),
"post_z": cat(post_pos[:, 2], marker_pre_pos[:, 2]),
"post_E": cat(post_E, np.zeros(n_marker, dtype=np.float32)),
"post_dx": cat(post_dir_world[:, 0], marker_pre_dir[:, 0]),
"post_dy": cat(post_dir_world[:, 1], marker_pre_dir[:, 1]),
"post_dz": cat(post_dir_world[:, 2], marker_pre_dir[:, 2]),
# Terminal markers dump the full remaining pre_E into edep.
"edep": cat(edep, marker_pre_E),
"step_length": cat(step_length, np.zeros(n_marker, dtype=np.float32)),
"material": rng.choice(["W", "Pb"], n),
"layer_id": rng.integers(0, 10, n).astype(np.int32),
"n_sec_pred": rng.integers(0, 3, n).astype(np.int32),
"termination_reason": ([""] * n_real + ["energy_cutoff"] * n_marker),
}
)
table = table.replace_schema_metadata(
{PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE}
)
rollout_path = tmp_path / "rollout_with_markers.parquet"
pq.write_table(table, rollout_path)
truth_path = tmp_path / "truth.parquet"
_write_truth_parquet(truth_path, n=20)
samples = load_rollout_vs_truth(rollout_path, truth_path)
assert samples.gen_raw.shape == (n_real, 9)
np.testing.assert_allclose(samples.gen_raw, _expected_raw9(*real_fields), atol=1e-4)
def test_load_rollout_vs_truth_rejects_wrong_coord_metadata(tmp_path):
truth_path = tmp_path / "truth.parquet"
rollout_path = tmp_path / "rollout.parquet"
_write_truth_parquet(truth_path)
_write_rollout_parquet(rollout_path, coord="local")
with pytest.raises(ValueError, match="not a rollout file"):
load_rollout_vs_truth(rollout_path, truth_path)
# ---------------------------------------------------------------------------
# Tier 4: compute_rollout_observables / compute_truth_observables (unpaired)
# ---------------------------------------------------------------------------
def _make_rollout_style_event_arrays(rng):
"""3 events (3/2/4 steps), each with an unambiguous highest-pre_E row.
Same forced-max-pre_E-row construction as `_make_event_level_arrays`
above, but for the plain world-frame rollout/truth schema, where
edep/step_length are used directly (no energy-simplex decode needed).
"""
event_id = np.array([0, 0, 0, 1, 1, 2, 2, 2, 2], dtype=np.int64)
n = len(event_id)
pre_pos = rng.uniform(-5.0, 5.0, (n, 3)).astype(np.float32)
pre_dir = _unit_vectors(rng, n)
pre_E = rng.uniform(1.0, 50.0, n).astype(np.float32)
pre_E[1] = 100.0 # event 0's entry step
pre_E[3] = 100.0 # event 1's entry step
pre_E[7] = 100.0 # event 2's entry step
step_length = rng.uniform(0.1, 5.0, n).astype(np.float32)
travel_dir = _unit_vectors(rng, n)
post_pos = pre_pos + step_length[:, None] * travel_dir
edep = rng.uniform(0.1, 5.0, n).astype(np.float32)
return event_id, pre_pos, pre_dir, pre_E, post_pos, edep, step_length
def _expected_rollout_style_table(event_id, pre_pos, pre_dir, pre_E, post_pos, edep):
"""Independent re-derivation of total_edep/centroid_depth per event."""
expected = {}
for e in sorted(np.unique(event_id).tolist()):
mask = event_id == e
entry_idx = np.where(mask)[0][np.argmax(pre_E[mask])]
entry_pos = pre_pos[entry_idx]
axis_dir = pre_dir[entry_idx]
disp = post_pos[mask] - entry_pos
depth = disp @ axis_dir
total_edep = float(edep[mask].sum())
centroid = float((edep[mask] * depth).sum() / total_edep)
expected[e] = (total_edep, centroid, int(mask.sum()))
return expected
def _write_rollout_style_event_parquet(path, rng):
event_id, pre_pos, pre_dir, pre_E, post_pos, edep, step_length = (
_make_rollout_style_event_arrays(rng)
)
n = len(event_id)
table = pa.table(
{
"event_id": event_id,
"track_id": np.zeros(n, dtype=np.int64),
"parent_id": np.full(n, -1, dtype=np.int64),
"generation": np.zeros(n, dtype=np.int64),
"step_no": np.arange(n, dtype=np.int64),
"pdg": rng.choice([11, -11, 22], n),
"pre_x": pre_pos[:, 0],
"pre_y": pre_pos[:, 1],
"pre_z": pre_pos[:, 2],
"pre_E": pre_E,
"pre_dx": pre_dir[:, 0],
"pre_dy": pre_dir[:, 1],
"pre_dz": pre_dir[:, 2],
"post_x": post_pos[:, 0],
"post_y": post_pos[:, 1],
"post_z": post_pos[:, 2],
"post_E": np.zeros(n, dtype=np.float32),
"post_dx": pre_dir[:, 0],
"post_dy": pre_dir[:, 1],
"post_dz": pre_dir[:, 2],
"edep": edep,
"step_length": step_length,
"material": rng.choice(["W", "Pb"], n),
"layer_id": rng.integers(0, 10, n).astype(np.int32),
"n_sec_pred": np.zeros(n, dtype=np.int32),
"termination_reason": [""] * n,
}
)
table = table.replace_schema_metadata(
{PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE}
)
pq.write_table(table, path)
return _expected_rollout_style_table(
event_id, pre_pos, pre_dir, pre_E, post_pos, edep
)
def _write_truth_style_event_parquet(path, rng):
event_id, pre_pos, pre_dir, pre_E, post_pos, edep, step_length = (
_make_rollout_style_event_arrays(rng)
)
n = len(event_id)
table = pa.table(
{
"event_id": event_id,
"pdg": rng.choice([11, -11, 22], n),
"pre_x": pre_pos[:, 0],
"pre_y": pre_pos[:, 1],
"pre_z": pre_pos[:, 2],
"pre_E": pre_E,
"pre_dx": pre_dir[:, 0],
"pre_dy": pre_dir[:, 1],
"pre_dz": pre_dir[:, 2],
"material": rng.choice(["W", "Pb"], n),
"layer_id": rng.integers(0, 10, n).astype(np.int32),
"child_track_ids": [[] for _ in range(n)],
"e_sec": np.zeros(n, dtype=np.float32),
"step_length": step_length,
"post_E": np.zeros(n, dtype=np.float32),
"edep": edep,
"post_dx": pre_dir[:, 0],
"post_dy": pre_dir[:, 1],
"post_dz": pre_dir[:, 2],
"post_x": post_pos[:, 0],
"post_y": post_pos[:, 1],
"post_z": post_pos[:, 2],
}
)
pq.write_table(table, path)
return _expected_rollout_style_table(
event_id, pre_pos, pre_dir, pre_E, post_pos, edep
)
def test_compute_rollout_observables_matches_manual_reconstruction(tmp_path):
path = tmp_path / "rollout_events.parquet"
# `centroid_depth` is weighted by *binned* depth (bin centers), not the raw
# continuous depth `_expected_rollout_style_table` computes, so it isn't
# checked here — see test_compute_rollout_and_truth_observables_agree_on_identical_data
# for a same-binning cross-check instead.
expected = _write_rollout_style_event_parquet(path, np.random.default_rng(11))
obs = compute_rollout_observables(path, depth_bins=5, transverse_bins=5)
table = obs.event_table.set_index("event_id")
for eid, (total_edep, _centroid, n_steps) in expected.items():
np.testing.assert_allclose(table.loc[eid, "total_edep"], total_edep, rtol=1e-4)
assert table.loc[eid, "n_steps"] == n_steps
def test_compute_truth_observables_matches_manual_reconstruction(tmp_path):
path = tmp_path / "truth_events.parquet"
expected = _write_truth_style_event_parquet(path, np.random.default_rng(12))
obs = compute_truth_observables(path, depth_bins=5, transverse_bins=5)
table = obs.event_table.set_index("event_id")
for eid, (total_edep, _centroid, n_steps) in expected.items():
np.testing.assert_allclose(
table.loc[eid, "real_total_edep"], total_edep, rtol=1e-4
)
assert table.loc[eid, "n_steps"] == n_steps
def test_compute_rollout_and_truth_observables_agree_on_identical_data(tmp_path):
"""`compute_rollout_observables`/`compute_truth_observables` must treat edep
identically: fed the exact same underlying step data (just written once
through each file's own schema), their depth/transverse profiles and total
per-event edep should come out numerically identical.
"""
rng = np.random.default_rng(13)
event_id, pre_pos, pre_dir, pre_E, post_pos, edep, step_length = (
_make_rollout_style_event_arrays(rng)
)
n = len(event_id)
rollout_table = pa.table(
{
"event_id": event_id,
"track_id": np.zeros(n, dtype=np.int64),
"parent_id": np.full(n, -1, dtype=np.int64),
"generation": np.zeros(n, dtype=np.int64),
"step_no": np.arange(n, dtype=np.int64),
"pdg": np.full(n, 11),
"pre_x": pre_pos[:, 0],
"pre_y": pre_pos[:, 1],
"pre_z": pre_pos[:, 2],
"pre_E": pre_E,
"pre_dx": pre_dir[:, 0],
"pre_dy": pre_dir[:, 1],
"pre_dz": pre_dir[:, 2],
"post_x": post_pos[:, 0],
"post_y": post_pos[:, 1],
"post_z": post_pos[:, 2],
"post_E": np.zeros(n, dtype=np.float32),
"post_dx": pre_dir[:, 0],
"post_dy": pre_dir[:, 1],
"post_dz": pre_dir[:, 2],
"edep": edep,
"step_length": step_length,
"material": np.full(n, "W"),
"layer_id": np.zeros(n, dtype=np.int32),
"n_sec_pred": np.zeros(n, dtype=np.int32),
"termination_reason": [""] * n,
}
)
rollout_table = rollout_table.replace_schema_metadata(
{PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE}
)
rollout_path = tmp_path / "rollout.parquet"
pq.write_table(rollout_table, rollout_path)
truth_table = pa.table(
{
"event_id": event_id,
"pdg": np.full(n, 11),
"pre_x": pre_pos[:, 0],
"pre_y": pre_pos[:, 1],
"pre_z": pre_pos[:, 2],
"pre_E": pre_E,
"pre_dx": pre_dir[:, 0],
"pre_dy": pre_dir[:, 1],
"pre_dz": pre_dir[:, 2],
"material": np.full(n, "W"),
"layer_id": np.zeros(n, dtype=np.int32),
"child_track_ids": [[] for _ in range(n)],
"e_sec": np.zeros(n, dtype=np.float32),
"step_length": step_length,
"post_E": np.zeros(n, dtype=np.float32),
"edep": edep,
"post_dx": pre_dir[:, 0],
"post_dy": pre_dir[:, 1],
"post_dz": pre_dir[:, 2],
"post_x": post_pos[:, 0],
"post_y": post_pos[:, 1],
"post_z": post_pos[:, 2],
}
)
truth_path = tmp_path / "truth.parquet"
pq.write_table(truth_table, truth_path)
rollout_obs = compute_rollout_observables(
rollout_path, depth_bins=5, transverse_bins=5
)
truth_obs = compute_truth_observables(truth_path, depth_bins=5, transverse_bins=5)
np.testing.assert_allclose(rollout_obs.depth_edges, truth_obs.depth_edges)
np.testing.assert_allclose(rollout_obs.transverse_edges, truth_obs.transverse_edges)
np.testing.assert_allclose(rollout_obs.depth_profile, truth_obs.real_depth_profile)
np.testing.assert_allclose(
rollout_obs.transverse_profile, truth_obs.real_transverse_profile
)
rollout_table_sorted = rollout_obs.event_table.sort_values("event_id")
truth_table_sorted = truth_obs.event_table.sort_values("event_id")
np.testing.assert_allclose(
rollout_table_sorted["total_edep"].to_numpy(),
truth_table_sorted["real_total_edep"].to_numpy(),
)
def test_plot_rollout_functions_accept_truth_observables_reference(tmp_path):
rollout_path = tmp_path / "rollout_events.parquet"
truth_path = tmp_path / "truth_events.parquet"
_write_rollout_style_event_parquet(rollout_path, np.random.default_rng(21))
_write_truth_style_event_parquet(truth_path, np.random.default_rng(22))
obs = compute_rollout_observables(rollout_path, depth_bins=5, transverse_bins=5)
reference = compute_truth_observables(truth_path, depth_bins=5, transverse_bins=5)
assert plot_rollout_longitudinal(obs, reference=reference) is not None
assert plot_rollout_transverse(obs, reference=reference) is not None
assert plot_rollout_total_energy(obs, reference=reference) is not None
tests/test_create_root_files.py
+79 -9
View File
@@ -9,6 +9,7 @@ from scripts import create_root_files
parse_detector_spec = create_root_files.parse_detector_spec
next_shard_index = create_root_files.next_shard_index
plan_jobs = create_root_files.plan_jobs
job_seed = create_root_files.job_seed
run_job = create_root_files.run_job
run_all = create_root_files.run_all
SimJob = create_root_files.SimJob
@@ -35,7 +36,13 @@ start = time.time()
time.sleep({sleep})
end = time.time()
payload = json.dumps(
{{"argv": sys.argv[1:], "cwd": os.getcwd(), "start": start, "end": end}}
{{
"argv": sys.argv[1:],
"cwd": os.getcwd(),
"start": start,
"end": end,
"seed": os.environ.get("MINICALOSIM_SEED"),
}}
)
for i in range({output_count}):
with open(f"out_{{i}}.root", "w") as f:
@@ -130,6 +137,47 @@ def test_plan_jobs_multiple_detectors_each_start_independently(tmp_path):
assert by_detector["sampling_fe_scint"] == [0, 1]
def test_job_seed_deterministic():
job = SimJob(detector="pbwo4", config=None, shard_index=3)
assert job_seed("steps", "gen1", job, None) == job_seed("steps", "gen1", job, None)
def test_job_seed_varies_by_shard_index():
a = SimJob(detector="pbwo4", config=None, shard_index=0)
b = SimJob(detector="pbwo4", config=None, shard_index=1)
assert job_seed("steps", "gen1", a, None) != job_seed("steps", "gen1", b, None)
def test_job_seed_varies_by_detector():
a = SimJob(detector="pbwo4", config=None, shard_index=0)
b = SimJob(detector="sampling_pb_scint", config="pb_scint", shard_index=0)
assert job_seed("steps", "gen1", a, None) != job_seed("steps", "gen1", b, None)
def test_job_seed_varies_by_gen():
job = SimJob(detector="pbwo4", config=None, shard_index=0)
assert job_seed("steps", "gen1", job, None) != job_seed("steps", "gen2", job, None)
def test_job_seed_varies_by_energy():
job = SimJob(detector="pbwo4", config=None, shard_index=0)
assert job_seed("steps", "gen1", job, 1.0) != job_seed("steps", "gen1", job, 10.0)
def test_run_job_passes_deterministic_seed_env_var(tmp_path):
fake = _write_fake_executable(tmp_path / "fake_exe.py")
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
tmp_root = tmp_path / ".sim-tmp"
tmp_root.mkdir()
job = SimJob(detector="pbwo4", config=None, shard_index=5)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert result.dest is not None
payload = json.loads(result.dest.read_text())
assert payload["seed"] == str(job_seed("steps", "gen1", job, None))
def test_run_job_moves_output_to_correct_shard_path(tmp_path):
fake = _write_fake_executable(tmp_path / "fake_exe.py")
gen_dir = tmp_path / "raw" / "steps" / "gen1"
@@ -138,7 +186,7 @@ def test_run_job_moves_output_to_correct_shard_path(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="pbwo4", config=None, shard_index=7)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert result.ok
assert result.dest == gen_dir / "pbwo4" / "shard-007.root"
@@ -154,7 +202,7 @@ def test_run_job_passes_config_arg_and_isolates_cwd(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="sampling_pb_scint", config="pb_scint", shard_index=0)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert result.ok
assert result.dest is not None
@@ -172,13 +220,27 @@ def test_run_job_omits_config_arg_when_none(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="pbwo4", config=None, shard_index=0)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert result.dest is not None
payload = json.loads(result.dest.read_text())
assert payload["argv"] == ["10000"]
def test_run_job_appends_energy_arg_when_given(tmp_path):
fake = _write_fake_executable(tmp_path / "fake_exe.py")
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
tmp_root = tmp_path / ".sim-tmp"
tmp_root.mkdir()
job = SimJob(detector="pbwo4_10gev", config=None, shard_index=0)
result = run_job(job, fake, 10000, 10.0, tmp_path, "steps", "gen1", tmp_root)
assert result.dest is not None
payload = json.loads(result.dest.read_text())
assert payload["argv"] == ["10000", "10.0"]
def test_run_job_fails_when_executable_errors(tmp_path):
fake = _write_fake_executable(tmp_path / "fake_exe.py", exit_code=1)
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
@@ -186,7 +248,7 @@ def test_run_job_fails_when_executable_errors(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="pbwo4", config=None, shard_index=0)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert not result.ok
assert "exited 1" in result.message
@@ -199,7 +261,7 @@ def test_run_job_fails_when_no_root_file_produced(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="pbwo4", config=None, shard_index=0)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert not result.ok
assert "found 0" in result.message
@@ -212,7 +274,7 @@ def test_run_job_fails_when_multiple_root_files_produced(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="pbwo4", config=None, shard_index=0)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert not result.ok
assert "found 2" in result.message
@@ -227,7 +289,7 @@ def test_run_job_refuses_to_overwrite_existing_shard(tmp_path):
tmp_root.mkdir()
job = SimJob(detector="pbwo4", config=None, shard_index=0)
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
assert not result.ok
assert "overwrite" in result.message
@@ -242,7 +304,15 @@ def test_run_all_caps_concurrency(tmp_path):
jobs = [SimJob(detector="pbwo4", config=None, shard_index=i) for i in range(6)]
results = run_all(
jobs, fake, 10000, tmp_path, "steps", "gen1", max_workers=2, tmp_root=tmp_root
jobs,
fake,
10000,
None,
tmp_path,
"steps",
"gen1",
max_workers=2,
tmp_root=tmp_root,
)
assert all(r.ok and r.dest is not None for r in results)
tests/test_dataset.py
+24 -58
View File
@@ -1,67 +1,33 @@
import numpy as np
from giant.data.dataset import StepsDataset, train_val_split
from giant.data.dataset import make_event_split
def _dummy(N=500, n_events=20):
def test_make_event_split_sizes():
rng = np.random.default_rng(42)
data = {"event_id": rng.integers(0, n_events, size=N)}
cond_cont = rng.standard_normal((N, 9)).astype(np.float32)
cond_cat = rng.integers(0, 3, size=(N, 2)).astype(np.int64)
target = rng.standard_normal((N, 6)).astype(np.float32)
return data, cond_cont, cond_cat, target
event_ids = rng.integers(0, 50, size=1000)
train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
unique = np.unique(event_ids)
assert len(train_set) + len(val_set) == len(unique)
def test_dataset_length():
data, cond_cont, cond_cat, target = _dummy()
assert len(StepsDataset(cond_cont, cond_cat, target)) == len(target)
def test_dataset_item_shapes():
data, cond_cont, cond_cat, target = _dummy()
c, k, t = StepsDataset(cond_cont, cond_cat, target)[0]
assert c.shape == (9,)
assert k.shape == (2,)
assert t.shape == (6,)
def test_split_sizes_sum_to_total():
data, cond_cont, cond_cat, target = _dummy(N=500)
train_ds, val_ds = train_val_split(
data, cond_cont, cond_cat, target, val_fraction=0.2
)
assert len(train_ds) + len(val_ds) == 500
def test_split_no_empty_sets():
data, cond_cont, cond_cat, target = _dummy(N=500, n_events=20)
train_ds, val_ds = train_val_split(
data, cond_cont, cond_cat, target, val_fraction=0.2
)
assert len(val_ds) > 0
assert len(train_ds) > 0
def test_split_event_leakage():
"""Train and val must not share any event_id."""
N = 1000
n_events = 50
def test_make_event_split_no_overlap():
rng = np.random.default_rng(7)
event_ids = rng.integers(0, n_events, size=N)
data = {"event_id": event_ids}
cond_cont = rng.standard_normal((N, 9)).astype(np.float32)
cond_cat = rng.integers(0, 3, size=(N, 2)).astype(np.int64)
target = rng.standard_normal((N, 6)).astype(np.float32)
event_ids = rng.integers(0, 50, size=1000)
train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
assert train_set.isdisjoint(val_set)
train_ds, val_ds = train_val_split(
data, cond_cont, cond_cat, target, val_fraction=0.2
)
# Recover which event_ids ended up in each split via the indices
# (The dataset doesn't store event_ids, so we check via the original mask logic)
unique_events = np.unique(event_ids)
rng2 = np.random.default_rng(42)
rng2.shuffle(unique_events)
n_val = max(1, int(len(unique_events) * 0.2))
val_events = set(unique_events[:n_val].tolist())
train_events = set(unique_events[n_val:].tolist())
assert val_events.isdisjoint(train_events)
def test_make_event_split_no_empty_sets():
rng = np.random.default_rng(0)
event_ids = rng.integers(0, 20, size=500)
train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
assert len(train_set) > 0
assert len(val_set) > 0
def test_make_event_split_reproducible():
event_ids = np.arange(100)
a_tr, a_val = make_event_split(event_ids, val_fraction=0.1, seed=42)
b_tr, b_val = make_event_split(event_ids, val_fraction=0.1, seed=42)
assert a_tr == b_tr
assert a_val == b_val
tests/test_flow.py
+6 -4
View File
@@ -39,8 +39,9 @@ def test_sample_flow_shape():
B = 6
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.zeros(B, 2, dtype=torch.long)
out = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
assert out.shape == (B, 9)
sample, n_sec = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
assert sample.shape == (B, 9)
assert n_sec.shape == (B,)
def test_ddpm_loss_nonneg():
@@ -55,5 +56,6 @@ def test_sample_ddim_shape():
schedule = CosineSchedule(T=50)
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.zeros(B, 2, dtype=torch.long)
out = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
assert out.shape == (B, 9)
sample, n_sec = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
assert sample.shape == (B, 9)
assert n_sec.shape == (B,)
tests/test_geometry.py
+178
View File
@@ -0,0 +1,178 @@
"""Tests for the geometry oracle (position -> material/layer_id + escape)."""
from pathlib import Path
from unittest.mock import patch
import numpy as np
import pytest
from giant import geometry as g
# scikit-learn is an optional extra; skip the whole module if it's missing.
pytest.importorskip("sklearn")
def _box_batch(n, rng):
"""A labelled point cloud: inside a 100mm box -> PbWO4/0, else AIR/-1."""
pos = rng.uniform(-200, 200, (n, 3)).astype(np.float32)
inside = (np.abs(pos) < 100).all(axis=1)
mat = np.where(inside, "G4_PbWO4", "G4_AIR").astype(object)
lay = np.where(inside, 0, -1).astype(np.int64)
return pos, mat, lay
def _build(subsample=30000, method="knn", escape_factor=5.0):
rng = np.random.default_rng(0)
batches = [_box_batch(20000, rng) for _ in range(3)]
with patch.object(g, "_iter_point_batches", lambda p: iter(batches)):
return g.build_geometry_oracle(
[Path("x")],
method=method,
subsample=subsample,
escape_factor=escape_factor,
)
def test_classes_discovered():
orc = _build()
assert set(orc.classes) == {("G4_PbWO4", 0), ("G4_AIR", -1)}
def test_query_labels_inside_and_outside():
orc = _build()
pos = np.array([[0.0, 0.0, 0.0], [150.0, 150.0, 150.0]])
material, layer_id, _ = orc.query(pos)
assert material[0] == "G4_PbWO4" and layer_id[0] == 0
assert material[1] == "G4_AIR" and layer_id[1] == -1
def test_escape_flag_fires_far_from_data():
orc = _build()
pos = np.array([[0.0, 0.0, 0.0], [1e5, 0.0, 0.0]])
_, _, escaped = orc.query(pos)
assert not escaped[0]
assert escaped[1]
def test_query_empty():
orc = _build()
material, layer_id, escaped = orc.query(np.empty((0, 3)))
assert len(material) == len(layer_id) == len(escaped) == 0
def test_query_bad_shape_raises():
orc = _build()
with pytest.raises(ValueError):
orc.query(np.zeros((4, 2)))
def test_save_load_roundtrip(tmp_path):
orc = _build()
p = tmp_path / "oracle.pkl"
orc.save(p)
loaded = g.GeometryOracle.load(p)
pos = np.array([[0.0, 0.0, 0.0], [150.0, 150.0, 150.0], [1e5, 0.0, 0.0]])
m0, l0, e0 = orc.query(pos)
m1, l1, e1 = loaded.query(pos)
assert (m0 == m1).all() and (l0 == l1).all() and (e0 == e1).all()
assert loaded.escape_threshold == orc.escape_threshold
assert loaded.classes == orc.classes
def test_svm_method_has_escape_tree():
orc = _build(subsample=4000, method="svm")
# SVM cannot answer NN-distance, so a reference tree backs the escape test.
assert orc._ref_tree is not None
_, _, escaped = orc.query(np.array([[1e5, 0.0, 0.0]]))
assert escaped[0]
def _layer_batch(n, rng):
"""Two 100mm slabs along z (with an air gap between/around them), bounded
to a 100x100mm transverse footprint — miniCaloSim's actual layer-stack
shape."""
z = rng.uniform(-20.0, 220.0, n).astype(np.float32)
x = rng.uniform(-50.0, 50.0, n).astype(np.float32)
y = rng.uniform(-50.0, 50.0, n).astype(np.float32)
material = np.full(n, "G4_AIR", dtype=object)
layer_id = np.full(n, -1, dtype=np.int64)
in_l0 = (z >= 0.0) & (z < 100.0)
in_l1 = (z >= 110.0) & (z < 210.0)
material[in_l0] = "G4_PbWO4"
layer_id[in_l0] = 0
material[in_l1] = "G4_W"
layer_id[in_l1] = 1
pos = np.stack([x, y, z], axis=1)
return pos, material, layer_id
def _build_slab(subsample=60000, n_bins=500, escape_factor=5.0):
rng = np.random.default_rng(0)
batches = [_layer_batch(20000, rng) for _ in range(3)]
with patch.object(g, "_iter_point_batches", lambda p: iter(batches)):
return g.build_geometry_oracle(
[Path("x")],
method="slab",
subsample=subsample,
escape_factor=escape_factor,
depth_axis=2,
n_bins=n_bins,
)
def test_slab_is_default_method():
rng = np.random.default_rng(0)
batches = [_layer_batch(20000, rng)]
with patch.object(g, "_iter_point_batches", lambda p: iter(batches)):
orc = g.build_geometry_oracle([Path("x")], subsample=20000)
assert orc.metadata["method"] == "slab"
assert orc._slab is not None
def test_slab_classes_discovered():
orc = _build_slab()
assert set(orc.classes) == {("G4_PbWO4", 0), ("G4_W", 1), ("G4_AIR", -1)}
def test_slab_query_labels_by_depth():
orc = _build_slab()
pos = np.array(
[[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0]]
) # layer 0, gap, layer 1
material, layer_id, escaped = orc.query(pos)
assert list(material) == ["G4_PbWO4", "G4_AIR", "G4_W"]
assert list(layer_id) == [0, -1, 1]
assert not escaped.any()
def test_slab_escape_beyond_depth_range():
orc = _build_slab()
pos = np.array([[0.0, 0.0, 50.0], [0.0, 0.0, 1e5]])
_, _, escaped = orc.query(pos)
assert not escaped[0]
assert escaped[1]
def test_slab_escape_beyond_transverse_radius():
orc = _build_slab()
pos = np.array([[0.0, 0.0, 50.0], [1e5, 1e5, 50.0]])
_, _, escaped = orc.query(pos)
assert not escaped[0]
assert escaped[1]
def test_slab_save_load_roundtrip(tmp_path):
orc = _build_slab()
p = tmp_path / "slab_oracle.pkl"
orc.save(p)
loaded = g.GeometryOracle.load(p)
pos = np.array(
[[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0], [0.0, 0.0, 1e5]]
)
m0, l0, e0 = orc.query(pos)
m1, l1, e1 = loaded.query(pos)
assert (m0 == m1).all() and (l0 == l1).all() and (e0 == e1).all()
assert loaded.escape_threshold == orc.escape_threshold
assert loaded._slab is not None
tests/test_network.py
+4 -2
View File
@@ -39,7 +39,9 @@ def test_denoising_mlp_gradients_flow():
t = torch.rand(B)
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.zeros(B, 2, dtype=torch.long)
loss = model(x_t, t, cond_cont, cond_cat).sum()
loss.backward()
# Both paths must be exercised to get gradients through all parameters.
flow_loss = model(x_t, t, cond_cont, cond_cat).sum()
nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum()
(flow_loss + nsec_loss).backward()
for name, p in model.named_parameters():
assert p.grad is not None, f"no grad for {name}"
tests/test_phase2.py
+360
View File
@@ -0,0 +1,360 @@
"""Tests for Phase 2: secondary particle prediction."""
import numpy as np
import pytest
import torch
from giant.constants import COND_DIM, EMB_DIM, K_MAX, 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
# ── helpers ──────────────────────────────────────────────────────────────────
def _stage1(pdg=3, mat=2):
return DenoisingMLP(pdg_vocab=pdg, mat_vocab=mat, hidden_dim=32, n_blocks=2)
def _sec_decoder(pdg=3, mat=2):
return SecondaryDecoder(pdg_vocab=pdg, mat_vocab=mat, hidden_dim=32, n_blocks=2)
def _cond(B=8, pdg=3, mat=2):
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.stack(
[torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
)
return cond_cont, cond_cat
# ── DenoisingMLP Phase-2 additions ───────────────────────────────────────────
def test_predict_n_sec_shape():
B = 8
model = _stage1()
cond_cont, cond_cat = _cond(B)
logits = model.predict_n_sec(cond_cont, cond_cat)
assert logits.shape == (B, K_MAX + 1)
def test_predict_n_sec_no_nan():
B = 8
model = _stage1()
cond_cont, cond_cat = _cond(B)
logits = model.predict_n_sec(cond_cont, cond_cat)
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)
# ── SecondaryDecoder ──────────────────────────────────────────────────────────
def test_sec_decoder_output_shape():
B = 8
decoder = _sec_decoder()
x_t = torch.randn(B, SEC_DIM)
t = torch.rand(B)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
out = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
assert out.shape == (B, SEC_DIM)
def test_sec_decoder_no_nan():
B = 4
decoder = _sec_decoder()
x_t = torch.randn(B, SEC_DIM)
t = torch.rand(B)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
out = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
assert torch.isfinite(out).all()
def test_sec_decoder_gradients():
B = 4
decoder = _sec_decoder()
x_t = torch.randn(B, SEC_DIM)
t = torch.rand(B)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward()
for name, p in decoder.named_parameters():
assert p.grad is not None, f"no grad for {name}"
# ── masked flow matching loss ─────────────────────────────────────────────────
def test_flow_matching_loss_secondary_scalar():
B, pdg, mat = 8, 3, 2
decoder = _sec_decoder(pdg, mat)
x1 = torch.randn(B, SEC_DIM)
cond_cont, cond_cat = _cond(B, pdg, mat)
stage1_out = torch.randn(B, X_DIM)
sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
loss = flow_matching_loss_secondary(
decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
)
assert loss.shape == ()
assert loss.item() >= 0.0
def test_flow_matching_loss_secondary_mask_zeros_padding():
"""Loss with all-zero mask (no valid secondaries) should be 0."""
B, pdg, mat = 4, 3, 2
decoder = _sec_decoder(pdg, mat)
x1 = torch.randn(B, SEC_DIM)
cond_cont, cond_cat = _cond(B, pdg, mat)
stage1_out = torch.randn(B, X_DIM)
sec_mask = torch.zeros(B, K_MAX, dtype=torch.bool)
loss = flow_matching_loss_secondary(
decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
)
assert loss.item() == pytest.approx(0.0, abs=1e-6)
def test_flow_matching_loss_secondary_has_grad():
B, pdg, mat = 4, 3, 2
decoder = _sec_decoder(pdg, mat)
x1 = torch.randn(B, SEC_DIM)
cond_cont, cond_cat = _cond(B, pdg, mat)
stage1_out = torch.randn(B, X_DIM)
sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
flow_matching_loss_secondary(
decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
).backward()
assert any(p.grad is not None for p in decoder.parameters())
# ── sampling ──────────────────────────────────────────────────────────────────
def test_sample_secondaries_shapes():
B, pdg, mat = 6, 3, 2
decoder = _sec_decoder(pdg, mat)
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(
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_valid.shape == (B, K_MAX)
assert sec_valid.dtype == torch.bool
def test_sample_secondaries_valid_mask_matches_n_sec():
B, pdg, mat = 4, 3, 2
decoder = _sec_decoder(pdg, mat)
cond_cont, cond_cat = _cond(B, pdg, mat)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.tensor([0, 1, 3, K_MAX])
_, _, sec_valid = sample_secondaries(
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
)
for i, n in enumerate(n_sec_pred.tolist()):
assert sec_valid[i, :n].all()
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 ─────────────────────────────────────────────
def test_encode_secondaries_energy_conservation():
"""Decoded stick-breaking fractions must sum to ≈ e_sec."""
from giant.data.transforms import encode_secondaries
rng = np.random.default_rng(42)
N = 50
n_sec = rng.integers(1, 5, size=N)
e_sec = rng.uniform(0.1, 10.0, size=N).astype(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)
for i in range(N):
k = n_sec[i]
energies = rng.dirichlet(np.ones(k)) * e_sec[i]
energies = np.sort(energies)[::-1]
sec_E_list[i, :k] = energies.astype(np.float32)
sec_valid[i, :k] = True
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)
assert np.isfinite(sec_cont).all()
def test_encode_secondaries_direction_encoding():
"""Local-frame secondary directions should be unit vectors for valid slots."""
from giant.data.transforms import encode_secondaries
rng = np.random.default_rng(7)
N = 20
e_sec = np.ones(N, dtype=np.float32) * 5.0
sec_E_list = np.zeros((N, K_MAX), dtype=np.float32)
sec_E_list[:, 0] = 3.0
sec_E_list[:, 1] = 2.0
sec_dir_list = rng.standard_normal((N, K_MAX, 3)).astype(np.float32)
norms = np.linalg.norm(sec_dir_list, axis=-1, keepdims=True)
sec_dir_list /= np.where(norms > 0, norms, 1.0)
sec_valid = np.zeros((N, K_MAX), dtype=bool)
sec_valid[:, :2] = True
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)
# dir columns are sec_cont[:, :, 1:4]
local_dirs = sec_cont[:, :2, 1:4] # (N, 2, 3) — valid slots only
norms_out = np.linalg.norm(local_dirs, axis=-1)
np.testing.assert_allclose(norms_out, 1.0, atol=1e-5)
# ── 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[:, :, 0] *= stick_logit_scale
dirs = sec_cont[:, :, 1:]
dirs /= np.linalg.norm(dirs, axis=-1, keepdims=True)
return sec_cont
def test_decode_secondaries_valid_slots_sum_to_e_sec():
"""The valid slots' energies must sum to exactly e_sec, not just <= e_sec.
Rows with n_sec=0 are excluded: there's no slot to put the budget in, so
valid_sum is correctly 0 regardless of e_sec there (see
test_decode_secondaries_zero_n_sec_has_zero_energy) — the shortfall in
that case is handled downstream (e.g. rollout.py dumps it into edep).
"""
from giant.data.transforms import decode_secondaries
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}
)
valid_sum = (sec_E * sec_valid).sum(axis=1)
has_secondaries = n_sec > 0
np.testing.assert_allclose(
valid_sum[has_secondaries],
e_sec[has_secondaries],
atol=1e-3,
rtol=1e-5,
)
def test_decode_secondaries_zero_n_sec_has_zero_energy():
"""n_sec=0 rows get no secondaries and no forced energy assignment."""
from giant.data.transforms import decode_secondaries
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}
)
assert not sec_valid.any()
np.testing.assert_allclose(sec_E, 0.0)
def test_decode_secondaries_degenerate_row_falls_back_to_even_split():
"""All-zero stick fractions for the valid slots fall back to an even split."""
from giant.data.transforms import decode_secondaries
rng = np.random.default_rng(2)
N = 4
sec_cont = _random_sec_cont(rng, N)
# Drive every valid slot's stick-breaking fraction to ~0 (huge negative logit).
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}
)
for i, k in enumerate(n_sec):
if k == 0:
continue
np.testing.assert_allclose(sec_E[i, :k], e_sec[i] / k, atol=1e-4)
np.testing.assert_allclose(sec_E[i, :k].sum(), e_sec[i], atol=1e-3)
def test_decode_secondaries_rescale_preserves_relative_shares():
"""Rescaling should keep each valid slot's *share* of the budget unchanged.
A shortfall shouldn't get dumped into whichever slot is last by energy
rank — it should be spread proportionally, i.e. sec_E[i] / sec_E[j] for
two valid slots must match before and after the e_sec rescale.
"""
from giant.data.transforms import decode_secondaries
rng = np.random.default_rng(3)
N = 1
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_large, _, _, _ = decode_secondaries(
sec_cont,
sec_pdg_pred,
n_sec,
np.array([50.0], dtype=np.float32),
pre_dir,
{0: 22},
)
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)
tests/test_rollout.py
+163
View File
@@ -0,0 +1,163 @@
"""Tests for the autoregressive shower rollout driver."""
from pathlib import Path
from unittest.mock import patch
import numpy as np
import pytest
import torch
from giant.constants import TERM_ESCAPED, TERM_MAX_STEPS
from giant.data.transforms import Normalizer
from giant.model.network import DenoisingMLP, SecondaryDecoder
from giant.rollout import make_seed_frontier, rollout
pytest.importorskip("sklearn")
from giant import geometry as g # noqa: E402
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)
return s1.eval(), s2.eval()
def _norms():
rng = np.random.default_rng(0)
cond = Normalizer().fit(rng.standard_normal((1000, 8)).astype(np.float32))
tgt = Normalizer().fit(rng.standard_normal((1000, 9)).astype(np.float32))
return cond, tgt
def _oracle():
rng = np.random.default_rng(0)
pos = rng.uniform(-200, 200, (20000, 3)).astype(np.float32)
inside = (np.abs(pos) < 100).all(axis=1)
mat = np.where(inside, "G4_PbWO4", "G4_AIR").astype(object)
lay = np.where(inside, 0, -1).astype(np.int64)
with patch.object(g, "_iter_point_batches", lambda p: iter([(pos, mat, lay)])):
# Pinned to "knn" explicitly: this test's escape-threshold semantics
# (tiny threshold -> escape even at a valid interior point, because no
# training point is that close) are KNN-specific, and the fixture's
# box geometry isn't a layer stack the "slab" method could fit anyway.
return g.build_geometry_oracle([Path("x")], method="knn", subsample=20000)
def _seeds(n=6):
return {
"event_id": np.arange(n, dtype=np.int64),
"pdg": np.full(n, 11, dtype=np.int64),
"pre_pos": np.zeros((n, 3)),
"pre_E": np.linspace(30.0, 90.0, n),
"pre_dir": np.tile([0.0, 0.0, 1.0], (n, 1)),
}
def _run(
escape_threshold=1e9,
energy_cutoff=1.0,
max_steps=30,
max_tracks_per_event=300,
seeds=None,
):
torch.manual_seed(0)
np.random.seed(0)
s1, s2 = _models()
cond, tgt = _norms()
return rollout(
s1,
s2,
_oracle(),
seeds or _seeds(),
cond,
tgt,
PDG_MAP,
MAT_MAP,
energy_cutoff=energy_cutoff,
max_steps=max_steps,
steps=4,
batch_size=128,
max_tracks_per_event=max_tracks_per_event,
escape_threshold=escape_threshold,
)
def test_seed_frontier_track_ids():
seeds = _seeds(3)
fr, counts = make_seed_frontier(**seeds)
assert (fr["track_id"] == [0, 0, 0]).all() # one primary per event -> id 0
assert (fr["parent_id"] == -1).all()
assert (fr["generation"] == 0).all()
assert all(counts[e] == 1 for e in range(3))
# pre_dir is normalised.
np.testing.assert_allclose(np.linalg.norm(fr["pre_dir"], axis=1), 1.0, atol=1e-6)
def test_rollout_terminates_and_has_rows():
rec = _run()
assert len(rec["event_id"]) > 0
# Every seed event appears.
assert set(rec["event_id"].tolist()) == set(range(6))
def test_max_steps_respected():
# Disable the energy cutoff so tracks survive long enough to hit the step cap.
rec = _run(max_steps=5, energy_cutoff=0.0)
assert rec["step_no"].max() <= 5
assert (rec["termination_reason"] == TERM_MAX_STEPS).any()
def test_energy_conserved_deposit_plus_leak():
seeds = _seeds()
rec = _run(seeds=seeds)
for i, ev in enumerate(seeds["event_id"]):
m = rec["event_id"] == ev
dep = rec["edep"][m].sum()
leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)
def test_secondaries_have_valid_parents():
rec = _run()
orphans = 0
for ev in np.unique(rec["event_id"]):
m = rec["event_id"] == ev
tids = set(rec["track_id"][m].tolist())
for pid in rec["parent_id"][m]:
if pid >= 0 and pid not in tids:
orphans += 1
assert orphans == 0
# At least one secondary (generation > 0) is produced by the tiny model.
assert (rec["generation"] > 0).any()
def test_escape_terminates_immediately():
# A tight escape threshold makes even the seed position (origin) escape.
rec = _run(escape_threshold=1e-3)
assert (rec["termination_reason"] == TERM_ESCAPED).all()
assert rec["step_no"].max() == 0
def test_output_schema_complete():
from giant.rollout import _RECORD_KEYS
rec = _run()
assert set(rec.keys()) == set(_RECORD_KEYS)
n = len(rec["event_id"])
assert all(len(v) == n for v in rec.values())
def test_max_tracks_cap_conserves_energy():
# A very small cap forces sub-cap secondaries to deposit in place; energy
# must still balance.
seeds = _seeds()
rec = _run(seeds=seeds, max_tracks_per_event=3)
for i, ev in enumerate(seeds["event_id"]):
m = rec["event_id"] == ev
dep = rec["edep"][m].sum()
leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)
assert len(np.unique(rec["track_id"][m])) <= 3
tests/test_steps_to_parquet.py
+40 -3
View File
@@ -12,13 +12,18 @@ def _frame() -> pl.DataFrame:
"track_id": [1, 1, 2, 1, 2, 3],
"step_no": [0, 1, 0, 0, 0, 0],
"pre_E": [100.0, 80.0, 15.0, 200.0, 20.0, 30.0],
"pdg": [11, 11, 22, 11, 22, 22],
"pre_dx": [0.0, 0.0, 1.0, 0.0, 1.0, 0.0],
"pre_dy": [0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
"pre_dz": [1.0, 1.0, 0.0, 1.0, 0.0, 0.0],
"child_track_ids": [[2], [], [], [2, 3], [], []],
}
)
def test_e_sec_sums_child_first_step_energy():
out = steps_to_parquet._add_secondary_energy(_frame())
out, n_orphaned = steps_to_parquet._add_secondary_attributes(_frame())
assert n_orphaned == 0
e_sec = dict(
zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"])
)
@@ -27,7 +32,7 @@ def test_e_sec_sums_child_first_step_energy():
def test_e_sec_zero_when_no_children():
out = steps_to_parquet._add_secondary_energy(_frame())
out, _ = steps_to_parquet._add_secondary_attributes(_frame())
childless = out.filter(
(pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1)
)
@@ -36,6 +41,38 @@ def test_e_sec_zero_when_no_children():
def test_e_sec_preserves_row_count_and_order():
df = _frame()
out = steps_to_parquet._add_secondary_energy(df)
out, _ = steps_to_parquet._add_secondary_attributes(df)
assert out.height == df.height
assert out["pre_E"].to_list() == df["pre_E"].to_list()
def test_orphaned_child_track_is_dropped_not_nulled():
"""A listed child_track_id with no first step of its own (e.g. absorbed
below the tracking threshold at birth) must not leave a null in
sec_E_list/sec_pdg_list/etc: that null turns into NaN once the parquet
round-trips through the loader, poisoning every later secondary slot in
the step via encode_secondaries' cumulative "remaining budget". It must
also be dropped from child_track_ids itself, so n_sec (len(child_track_ids)
downstream) matches the actual, orphan-free secondary lists."""
df = pl.DataFrame(
{
"event_id": [0, 0, 0],
"track_id": [1, 1, 2],
"step_no": [0, 1, 0],
"pre_E": [100.0, 80.0, 15.0],
"pdg": [11, 11, 22],
"pre_dx": [0.0, 0.0, 1.0],
"pre_dy": [0.0, 0.0, 0.0],
"pre_dz": [1.0, 1.0, 0.0],
# track 3 is listed as a child but never appears with its own step.
"child_track_ids": [[2, 3], [], []],
}
)
out, n_orphaned = steps_to_parquet._add_secondary_attributes(df)
row = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0))
assert n_orphaned == 1
assert row["child_track_ids"].to_list() == [[2]]
assert row["e_sec"].item() == 15.0
assert row["sec_E_list"].to_list() == [[15.0]]
assert None not in row["sec_E_list"].item()
tests/test_transforms.py
+82
View File
@@ -1,6 +1,8 @@
import numpy as np
import pytest
from giant.constants import K_MAX
from giant.data.transforms import (
build_features,
energy_simplex_decode,
energy_simplex_encode,
inv_local_frame_rotation,
@@ -205,3 +207,83 @@ def test_normalizer_serialization():
assert norm2.std is not None and norm.std is not None
np.testing.assert_allclose(norm2.mean, norm.mean)
np.testing.assert_allclose(norm2.std, norm.std)
def test_build_features_clamps_n_sec_label_to_k_max():
"""A step with more secondaries than K_MAX must not overflow the
n_sec classifier's K_MAX+1 classes (regression test: this used to hand
cross_entropy an out-of-range target and crash CUDA training with
'unique_by_key: failed to synchronize: cudaErrorAssert')."""
N = 3
raw_n_sec = np.array([0, 5, K_MAX + 20], dtype=np.int32)
rng = np.random.default_rng(0)
data = {
"pdg": np.array([11, 11, 11], dtype=np.int32),
"material": np.array(["PbWO4", "PbWO4", "PbWO4"], dtype=object),
"pre_pos": rng.standard_normal((N, 3)).astype(np.float32),
"pre_E": np.full(N, 10.0, dtype=np.float32),
"pre_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)),
"layer_id": np.zeros(N, dtype=np.int32),
"n_sec": raw_n_sec,
"e_sec": np.full(N, 1.0, dtype=np.float32),
"step_length": np.full(N, 1.0, dtype=np.float32),
"post_E": np.full(N, 9.0, dtype=np.float32),
"edep": np.full(N, 1.0, dtype=np.float32),
"post_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)),
"post_pos": rng.standard_normal((N, 3)).astype(np.float32),
}
pdg_map = {11: 0}
mat_map = {"PbWO4": 0}
_, _, _, 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])
def _step_data_no_sec_lists(n_sec: np.ndarray) -> dict:
"""Minimal build_features input with n_sec but no per-secondary list columns
(mimics a parquet that skipped the parent->child join)."""
N = len(n_sec)
rng = np.random.default_rng(0)
return {
"pdg": np.full(N, 11, dtype=np.int32),
"material": np.full(N, "PbWO4", dtype=object),
"pre_pos": rng.standard_normal((N, 3)).astype(np.float32),
"pre_E": np.full(N, 10.0, dtype=np.float32),
"pre_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)),
"layer_id": np.zeros(N, dtype=np.int32),
"n_sec": np.asarray(n_sec, dtype=np.int32),
"e_sec": np.full(N, 1.0, dtype=np.float32),
"step_length": np.full(N, 1.0, dtype=np.float32),
"post_E": np.full(N, 9.0, dtype=np.float32),
"edep": np.full(N, 1.0, dtype=np.float32),
"post_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)),
"post_pos": rng.standard_normal((N, 3)).astype(np.float32),
}
def test_build_features_require_secondaries_raises_when_lists_missing():
"""A parquet with n_sec > 0 but no per-secondary list columns was never run
through the parent->child join; require_secondaries must catch it instead of
silently zeroing every Stage-2 target (regression: this collapsed the
secondary species to a single PDG index during training)."""
data = _step_data_no_sec_lists(np.array([0, 2, 1]))
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
with pytest.raises(ValueError, match="per-secondary columns"):
build_features(data, pdg_map, mat_map, require_secondaries=True)
def test_build_features_require_secondaries_ok_when_no_secondaries():
"""require_secondaries only fires when secondaries actually exist; a file
with n_sec == 0 everywhere (e.g. Stage-1-only) must still load."""
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(
data, pdg_map, mat_map, require_secondaries=True
)
assert not sec_cont.any()
assert not sec_pdg_idx.any()
uv.lock
Generated
+124 -2
View File
@@ -470,14 +470,18 @@ dev = [
{ name = "polars" },
{ name = "pytest" },
{ name = "ruff" },
{ name = "scikit-learn" },
{ name = "ty" },
{ name = "uproot" },
]
geometry = [
{ name = "scikit-learn" },
]
[package.metadata]
requires-dist = [
{ name = "awkward", marker = "extra == 'convert'", specifier = ">=2.6,<3" },
{ name = "giant", extras = ["convert", "analysis"], marker = "extra == 'dev'" },
{ name = "giant", extras = ["convert", "analysis", "geometry"], marker = "extra == 'dev'" },
{ name = "ipykernel", marker = "extra == 'analysis'", specifier = ">=7.3.0" },
{ name = "matplotlib", marker = "extra == 'analysis'", specifier = ">=3.8,<4" },
{ name = "numpy", specifier = ">=1.26,<3" },
@@ -488,6 +492,7 @@ requires-dist = [
{ name = "pytest", marker = "extra == 'dev'", specifier = ">=8,<10" },
{ name = "pyyaml", specifier = ">=6,<7" },
{ name = "ruff", marker = "extra == 'dev'", specifier = ">=0.15,<1" },
{ name = "scikit-learn", marker = "extra == 'geometry'", specifier = ">=1.4,<2" },
{ name = "torch", marker = "extra == 'cpu'", specifier = ">=2.3,<2.4", index = "https://download.pytorch.org/whl/cpu", conflict = { package = "giant", extra = "cpu" } },
{ name = "torch", marker = "extra == 'cuda'", specifier = ">=2.3,<2.4", index = "https://download.pytorch.org/whl/cu118", conflict = { package = "giant", extra = "cuda" } },
{ name = "tqdm", specifier = ">=4.60,<5" },
@@ -495,7 +500,7 @@ requires-dist = [
{ name = "typer", specifier = ">=0.12,<1" },
{ name = "uproot", marker = "extra == 'convert'", specifier = ">=5.3,<6" },
]
provides-extras = ["cpu", "cuda", "dev", "convert", "analysis"]
provides-extras = ["cpu", "cuda", "dev", "geometry", "convert", "analysis"]
[[package]]
name = "iniconfig"
@@ -600,6 +605,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" },
]
[[package]]
name = "joblib"
version = "1.5.3"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/41/f2/d34e8b3a08a9cc79a50b2208a93dce981fe615b64d5a4d4abee421d898df/joblib-1.5.3.tar.gz", hash = "sha256:8561a3269e6801106863fd0d6d84bb737be9e7631e33aaed3fb9ce5953688da3", size = 331603, upload-time = "2025-12-15T08:41:46.427Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl", hash = "sha256:5fc3c5039fc5ca8c0276333a188bbd59d6b7ab37fe6632daa76bc7f9ec18e713", size = 309071, upload-time = "2025-12-15T08:41:44.973Z" },
]
[[package]]
name = "jupyter-client"
version = "8.9.1"
@@ -891,6 +905,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" },
]
[[package]]
name = "narwhals"
version = "2.23.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e8/ac/66ed1fc6e38a0c0f330627ec5c5d597990d6159b6712b82af0ad2c65f06c/narwhals-2.23.0.tar.gz", hash = "sha256:13e7ff5b4bb4a2f77b907c2e4d8a76e273dfc1323a3c997440a2f9fd26aed408", size = 656209, upload-time = "2026-07-01T11:21:53.278Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/f4/4e/afc8c31605cb8be1d3bb4438c4d979daa104dab6306cd2b87abe9c3a7299/narwhals-2.23.0-py3-none-any.whl", hash = "sha256:769e7b9ab102c93d8fa019f6b4cd1a657909b04a20bf6210e5a35aae06814ae9", size = 458938, upload-time = "2026-07-01T11:21:51.677Z" },
]
[[package]]
name = "nest-asyncio2"
version = "1.7.2"
@@ -1571,6 +1594,96 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/2d/c7/c53e8dbff9c9dc4b7928773421ae294a5d28fcb8dcda1a089579d3a7e510/ruff-0.15.17-py3-none-win_arm64.whl", hash = "sha256:f3be1fbb34bcdfd146240d8fb92a709d4c2c8191348580a3c044ec60fa0b4456", size = 11355275, upload-time = "2026-06-11T17:54:43.635Z" },
]
[[package]]
name = "scikit-learn"
version = "1.9.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "joblib" },
{ name = "narwhals" },
{ name = "numpy" },
{ name = "scipy" },
{ name = "threadpoolctl" },
]
sdist = { url = "https://files.pythonhosted.org/packages/fa/6f/37092bdb25f712817231799fc5674d8e704066a8a70c1d2d40517e18b4ab/scikit_learn-1.9.0.tar.gz", hash = "sha256:8833266989d3a5110178a9fae30783675460724d0e1efb13b14901d2c660c557", size = 7750767, upload-time = "2026-06-02T11:54:32.706Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/ac/20/75f915ff375d6249e6550ac740fdbbd66159a068fd3af1400ff62036b07a/scikit_learn-1.9.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:2bd41b0d201bc81575531b96b713d3eb5e5f50fb0b82101ff0f92294fdc236ac", size = 8741122, upload-time = "2026-06-02T11:53:24.08Z" },
{ url = "https://files.pythonhosted.org/packages/cc/d5/2b5148f2279196775e1db2aeb85d14b70ac80e7e32b3b28e7ebeafb0901d/scikit_learn-1.9.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:5be45aa4a42a68a533913a6ed736cf309de2226411c79ef8d609a5456f1939b1", size = 8261512, upload-time = "2026-06-02T11:53:27.183Z" },
{ url = "https://files.pythonhosted.org/packages/a0/ee/5adbc77656b71f9456a2f5a7a9fdb4bcf9207a6b962889f1c2f9323afa4e/scikit_learn-1.9.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5e50ed4da51974e86e940690e9a3d82e729b62b5a49f7c9bac534d515d39d86f", size = 8837603, upload-time = "2026-06-02T11:53:30.328Z" },
{ url = "https://files.pythonhosted.org/packages/6c/c2/63fdda36c56437eeb44aaf9493c8bcd62ce230ab1598924fc626ffbfa943/scikit_learn-1.9.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:056c92bb67ad4c28463c2f2653d9701449201e7e7a9e94e321be0f71c4fef2b8", size = 9132097, upload-time = "2026-06-02T11:53:33.456Z" },
{ url = "https://files.pythonhosted.org/packages/83/a4/c8e67227c680e2259c8864ae72ff48b06e16a6f51253a22167aa02a8aa4e/scikit_learn-1.9.0-cp312-cp312-win_amd64.whl", hash = "sha256:4306775fad04cc4b472a1b15af1ae9cede1540fbfcc17fbce3767cd8dc7ae283", size = 8211173, upload-time = "2026-06-02T11:53:36.602Z" },
{ url = "https://files.pythonhosted.org/packages/cf/fd/3c0863792e98e67e9184aa4029288a175935eb65443afcd30d4f143450cf/scikit_learn-1.9.0-cp312-cp312-win_arm64.whl", hash = "sha256:26e22435f63bcdcf396b574273f29f13dd531f5ea035801f5be10ba1540a4e60", size = 7867451, upload-time = "2026-06-02T11:53:39.075Z" },
{ url = "https://files.pythonhosted.org/packages/3c/01/cf3310626b6d48d3e9be69a1223f9180360b5e6edb045f50fade723ce494/scikit_learn-1.9.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:80746d63bd4b6eaca54d36fe5feaf4d28bb38dc6f9470f81c7cad7c40155f119", size = 8705188, upload-time = "2026-06-02T11:53:41.964Z" },
{ url = "https://files.pythonhosted.org/packages/3e/04/5acd7ae280c5f93b6ac5ef6cdec14eef4c8d1cd91d85b3292989c94d96b1/scikit_learn-1.9.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:5b934c45c252844a91d69fda3a34cff5e7307e1db10d77cb10a3980312c74713", size = 8228299, upload-time = "2026-06-02T11:53:44.817Z" },
{ url = "https://files.pythonhosted.org/packages/0c/39/ffe829a5b8ecb40a518724a997794657fdc354ada5e8fe8e64d998c0bac9/scikit_learn-1.9.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:38c3dcb9a1ffb85505ec53d54c7b4aea0cff70050425a7760c2af661ac85df05", size = 8789690, upload-time = "2026-06-02T11:53:47.461Z" },
{ url = "https://files.pythonhosted.org/packages/1f/88/8dab5de10c638c083772a6be83a3d8106ced492f74a928c8693638e5bb50/scikit_learn-1.9.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:da76d09304a4706db7cc1e3ebaa3b6b98a67365cc11d2996c4f1e58ba47df714", size = 9087723, upload-time = "2026-06-02T11:53:50.702Z" },
{ url = "https://files.pythonhosted.org/packages/20/3f/7917ca72464038f6240ec70c29f94862d08a34a74291ae4d4ec5eb8186a0/scikit_learn-1.9.0-cp313-cp313-win_amd64.whl", hash = "sha256:5808d98f15c6bf6d9d96d2348c1997392a5888ce7097e664105f930c4bca1277", size = 8184330, upload-time = "2026-06-02T11:53:53.396Z" },
{ url = "https://files.pythonhosted.org/packages/78/c7/15739eb2f61fda3c54639e9942414e5a19ad8a8d1f5a3266afad7cb7df80/scikit_learn-1.9.0-cp313-cp313-win_arm64.whl", hash = "sha256:d77f54c017633791bc0225a43e2f8d03745fdcfe4880268fcc4df15f505dec2e", size = 7840653, upload-time = "2026-06-02T11:53:56.035Z" },
{ url = "https://files.pythonhosted.org/packages/f4/7d/c9a35cf59b20a86fec24d306f1547b78dec194b08d367ce2a3e4854169d9/scikit_learn-1.9.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:9656acd4e93f74e0b66c8a36c88830a99252dfa900044d36bc2212ae89a47162", size = 8713289, upload-time = "2026-06-02T11:53:58.788Z" },
{ url = "https://files.pythonhosted.org/packages/3c/a7/552a7821597c632b907f7bfe8f36f9f572777af8ef8a48353041cf8e091a/scikit_learn-1.9.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:24360002ae845e7866522b0a5bbf690802e7bc388cac8663502e78aa98598aa2", size = 8245141, upload-time = "2026-06-02T11:54:01.694Z" },
{ url = "https://files.pythonhosted.org/packages/7d/79/f4a0c4fe9711154cddabf913471153af79056382ddc612cfe5ee0ff4b72e/scikit_learn-1.9.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5162ad10a418c8a282dde04c9aa06965de3e9a65f33c1440c0ae69bb1a09d913", size = 8847671, upload-time = "2026-06-02T11:54:04.448Z" },
{ url = "https://files.pythonhosted.org/packages/f0/af/4d72d9e475ac83719160c662619e4bf7b95c19507cd582e7d0167a3c3dae/scikit_learn-1.9.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1fea2cc5677ab49d6f5bade978c866da44957b712d92e9635e8b4f723013c3cb", size = 9118104, upload-time = "2026-06-02T11:54:07.205Z" },
{ url = "https://files.pythonhosted.org/packages/a2/d5/6a58eea2cb9abbb9b3f2bb8b2cfb3243d1152d69f442d256c7af71304769/scikit_learn-1.9.0-cp314-cp314-win_amd64.whl", hash = "sha256:64fa347efc1c839c487433e40c5144d38c336e8a2b59c81aa8660373945c2673", size = 8290674, upload-time = "2026-06-02T11:54:10.087Z" },
{ url = "https://files.pythonhosted.org/packages/65/5b/d4c879cf358f1187141cf90ced473f087183489090244f50c124a2ee478b/scikit_learn-1.9.0-cp314-cp314-win_arm64.whl", hash = "sha256:1b944b6db288f6b926e3650026ddafb988929de95d11fc2cc5fa117773c9ba42", size = 7978807, upload-time = "2026-06-02T11:54:12.769Z" },
{ url = "https://files.pythonhosted.org/packages/8a/43/bfae3121ec67ae09150d453c442c7c1cc166e9aefe056e6ab3b7728a5cfc/scikit_learn-1.9.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:4ccacf04ca5f4b492158a5f28afe0ace43f81b2571e4b9a66d34848b46128949", size = 9031941, upload-time = "2026-06-02T11:54:15.436Z" },
{ url = "https://files.pythonhosted.org/packages/75/b0/20a4546eb17f3b25d3c66df15810411c14ed5065bcfab50b53c96fb627b2/scikit_learn-1.9.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:ee1a8db2c18c08e34c7412d4b10be1cac214cd4ea7dc9715a6a327eb49a37c96", size = 8613528, upload-time = "2026-06-02T11:54:18.842Z" },
{ url = "https://files.pythonhosted.org/packages/18/3c/e440e039bb82cd19004edaaad00acbde0fb9b461083c3ecf37941c557312/scikit_learn-1.9.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:147e9329ef0e39f75d4cffa02b2aa48d827832684926cd5210d9a2cb5c57246b", size = 8855050, upload-time = "2026-06-02T11:54:21.699Z" },
{ url = "https://files.pythonhosted.org/packages/43/26/b341b8dab5998da6270a3a42c2152c578501354d36f944b5856757035ef8/scikit_learn-1.9.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5bad8f8b9950321b54c965fdcbac6c6c55e79e16646b49977bcf3668d3870a1a", size = 9097190, upload-time = "2026-06-02T11:54:24.454Z" },
{ url = "https://files.pythonhosted.org/packages/fb/de/b650b4d69b84468cfa2e28a3ff7b8103743029e6446ce1a97fe060ef688c/scikit_learn-1.9.0-cp314-cp314t-win_amd64.whl", hash = "sha256:78fc56eafd4edb9575d2d8950d1dd152061abb573341a1cb7e099fc40f6c6666", size = 8963204, upload-time = "2026-06-02T11:54:27.428Z" },
{ url = "https://files.pythonhosted.org/packages/ee/f3/ff83d76d7418112e5a61326443cdda87be3545dd8d6599c95b2481a4419e/scikit_learn-1.9.0-cp314-cp314t-win_arm64.whl", hash = "sha256:051075bda8b7aab87b1906ab3d4740a1e1224a19d7b3781a576736edc94e76aa", size = 8222661, upload-time = "2026-06-02T11:54:30.192Z" },
]
[[package]]
name = "scipy"
version = "1.18.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "numpy" },
]
sdist = { url = "https://files.pythonhosted.org/packages/a7/25/c2700dfaf6442b4effaa91af24ebce5dc9d31bb4a69706313aae70d72cd0/scipy-1.18.0.tar.gz", hash = "sha256:67b2ad2ad54c72ca6d04975a9b2df8c3638c34ddd5b28738e94fc2b57929d378", size = 30774447, upload-time = "2026-06-19T15:01:43.456Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/6a/19/ca10ead60b0acc80b2b833c2c4a4f2ff753d0f58b811f70d911c7e94a25c/scipy-1.18.0-cp312-cp312-macosx_10_15_x86_64.whl", hash = "sha256:7bd21faaf5a1a3b2eff922d02db5f191b99a6518db9078a8fb23169f6d22259a", size = 31056519, upload-time = "2026-06-19T14:59:45.203Z" },
{ url = "https://files.pythonhosted.org/packages/96/72/1e6442a00cd2924d361aa1b642ab6373ec35c6fabf311a760be9f76e0f13/scipy-1.18.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:265915e79107de9f946b855e50d7470d5893ec3f54b342e1aa6201cbdcd8bb6b", size = 28681889, upload-time = "2026-06-19T14:59:48.103Z" },
{ url = "https://files.pythonhosted.org/packages/9b/2d/11dd93d21e147a73ba22bd75c0b9208d3a2e0ec76d53170ce7d9029b1015/scipy-1.18.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:9ab7b758be6940954a713ee466e2043e9f6e2ed965c1fce5c91039f4be3d90a9", size = 20423580, upload-time = "2026-06-19T14:59:50.665Z" },
{ url = "https://files.pythonhosted.org/packages/9c/01/93552f75e0d2a7dd115a45e59209c51e8d514daff02fc887d2623be06fe1/scipy-1.18.0-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:97b6cddaaee0a779ef6b5ca83c9604b27cc16b2b8fc22c142652df8793319fb8", size = 23054441, upload-time = "2026-06-19T14:59:53.564Z" },
{ url = "https://files.pythonhosted.org/packages/3c/23/21f5e703643d66f21faa6b4c73195bfcad70c55efcb4f1ab327cd7c4101a/scipy-1.18.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:52a96e21517c7292375c0e27dd796a811f03fcea5fd4d108fdfea8145dcf17ab", size = 33968720, upload-time = "2026-06-19T14:59:56.415Z" },
{ url = "https://files.pythonhosted.org/packages/dd/aa/1b939f6c67ed68635bb538e6752d3dacc02f66535182e939a89581a44e9c/scipy-1.18.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1f55797419e16e7f30cf88ffb3113ce0467f00cfe3f70d5c281730b21769bfc2", size = 35287115, upload-time = "2026-06-19T14:59:59.411Z" },
{ url = "https://files.pythonhosted.org/packages/b6/ff/eec46be7e9234208f801062b53e1983085eddebd693f6c9bfb03b459830d/scipy-1.18.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ad033410e2e0672ffdc1042110cef20e1c46f8fd0616cee1d44d8d58fad8fc11", size = 35577989, upload-time = "2026-06-19T15:00:02.235Z" },
{ url = "https://files.pythonhosted.org/packages/84/ca/210d4759c7210bb7d269437421959b39a33434e2776b60c5cb8a763bb30a/scipy-1.18.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:4a55985d54c769c872e64b7f4c8a81cc30ef700cc04296abbbf3705439c126de", size = 37421717, upload-time = "2026-06-19T15:00:05.102Z" },
{ url = "https://files.pythonhosted.org/packages/2b/54/9a9edb45345bd6744da5ddfb6628e5d5185920494c6a67ec45b6381004cb/scipy-1.18.0-cp312-cp312-win_amd64.whl", hash = "sha256:71ccc8faa2dd16ac310233203474a8b5cb67f10dedd54a3116d34943f4b19132", size = 36597428, upload-time = "2026-06-19T15:00:08.112Z" },
{ url = "https://files.pythonhosted.org/packages/99/0e/33f32a2a58987e26aec0f7df252cbbad1e90ae77bdbc76f40dd4ed0cf0ea/scipy-1.18.0-cp312-cp312-win_arm64.whl", hash = "sha256:d88363fd9d8fbd3511bd273f1a49efb2a540773ddf92a91d57498ce7dd7f3e76", size = 24351481, upload-time = "2026-06-19T15:00:11.103Z" },
{ url = "https://files.pythonhosted.org/packages/05/52/9c0136c2de7ae0779b7b366447766cec6d9f0702c56bb8ffeb04c8fd3af4/scipy-1.18.0-cp313-cp313-macosx_10_15_x86_64.whl", hash = "sha256:09143f676d157d9f546d663504ef9c1becb819824f1afc018814176411942446", size = 31036107, upload-time = "2026-06-19T15:00:14.03Z" },
{ url = "https://files.pythonhosted.org/packages/02/73/0291a64843270f4efb86cdcf2ee0f2048631b65ec6b405398b2b4dbf11bf/scipy-1.18.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:5efe260f69417b97ddae455bfb5a95e8359f7f66ad7fa9522a60feb66f169520", size = 28663303, upload-time = "2026-06-19T15:00:16.819Z" },
{ url = "https://files.pythonhosted.org/packages/d3/0f/10ffa0b697a572f4e0d48b92a88895d366422f019f723e7e14a84c050dac/scipy-1.18.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:68363b7eaacd8b5dd426df56d782cc156468ac79a127a1b87ca597d6e2e82197", size = 20404960, upload-time = "2026-06-19T15:00:19.635Z" },
{ url = "https://files.pythonhosted.org/packages/7e/d2/e896cea21ba8edd6c81d4c55b1ffcc717e79698dcbebf9641b4cfb4c6622/scipy-1.18.0-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:c5557d8be5da8e41353fcd4d21491fdbab83b062fc579e94dc09a7c8ab4f669b", size = 23034074, upload-time = "2026-06-19T15:00:22.107Z" },
{ url = "https://files.pythonhosted.org/packages/ea/b2/e83ea34279a52c03374477c74006256ec78df65fc877baa4617d6de1d202/scipy-1.18.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0d13bca67c096d89fb95ced0d8921807300fce0275643aef9533cc63a0773468", size = 33942038, upload-time = "2026-06-19T15:00:24.964Z" },
{ url = "https://files.pythonhosted.org/packages/f6/af/e8fe5fb136f51e2b01678b92cb4106d10d8cd68ec147ead2e7cb0ac75398/scipy-1.18.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a46f9273dbd0eb1cefba61c9b8648b4dfe3cbc14a080176f9a73e44b8336dc7f", size = 35266390, upload-time = "2026-06-19T15:00:28.059Z" },
{ url = "https://files.pythonhosted.org/packages/3a/49/2c5cbb907b56695fc67517811d1db234dfd83381a84814ec220aded2794d/scipy-1.18.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:5aba46108853ddfc77906b6557aac839d2b52e900c1d72a1180adaaab58d265f", size = 35551324, upload-time = "2026-06-19T15:00:31.014Z" },
{ url = "https://files.pythonhosted.org/packages/bb/73/eda39f7a2d306ff0ffc574afd13c0bbb6d10a603d9a413998ee269487a80/scipy-1.18.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b6f758e35f12757b5d95c00bc6de2438e229c2664b7a92e96f205959d9f2dfa4", size = 37404785, upload-time = "2026-06-19T15:00:34.072Z" },
{ url = "https://files.pythonhosted.org/packages/b7/d2/ae881ee28d014f38e0ccbfd974a06a919ba9af34f1f74bf42b5301891d63/scipy-1.18.0-cp313-cp313-win_amd64.whl", hash = "sha256:1afac4a847207c7ff8efd321734a50b06d0280b3b2a2c0fc2f413101747ad7c7", size = 36554943, upload-time = "2026-06-19T15:00:36.903Z" },
{ url = "https://files.pythonhosted.org/packages/70/3a/21154e2d54eb3639c6bf4dbae2e531c68356bfe95990daa30df33b30d556/scipy-1.18.0-cp313-cp313-win_arm64.whl", hash = "sha256:c5dbddf60e58c2312316d097271a8e73d40eaf2eabfa4d95ed7d3695bbf2ce7b", size = 24350911, upload-time = "2026-06-19T15:00:40.062Z" },
{ url = "https://files.pythonhosted.org/packages/78/b5/915a19b3de2f7430062b509653563db1633ddbb6f021b06731521115d4e2/scipy-1.18.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:4c256ee70c0d1a8a2ace807e199ccd4e3f57037433842abb3fb36bc17eaa9578", size = 31036253, upload-time = "2026-06-19T15:00:43.216Z" },
{ url = "https://files.pythonhosted.org/packages/d7/88/b72def7262e150d16be13fca37a96481138d624e700340bc3362a7588929/scipy-1.18.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:2ef3abc54a4ffc53765374b0d5728532dfdd2585ed23f6b11c206a1f0b1b9af8", size = 28673758, upload-time = "2026-06-19T15:00:46.663Z" },
{ url = "https://files.pythonhosted.org/packages/91/02/2e636a61a525632c373cf6a9c24442a3ffb79e364d38e98b32042964ac32/scipy-1.18.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:f2a6af57bd9e4a75d70e4117e78a1bbee84f79ae3fbb6d0111005d6ebcc4cb8d", size = 20415514, upload-time = "2026-06-19T15:00:49.399Z" },
{ url = "https://files.pythonhosted.org/packages/c9/b6/2135974442f6aba159d9d39d774a1c8cb19947016725d69fecc685df45bf/scipy-1.18.0-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:3f1ac564d3bf6c03d861d2cd87a1bea0da2887136f7fb1bf519c05a8971452d6", size = 23034398, upload-time = "2026-06-19T15:00:51.941Z" },
{ url = "https://files.pythonhosted.org/packages/f6/e6/ba89ec5abf6ee9257c0d1ec985573f3ae32742c24bc03e016388a40b1b15/scipy-1.18.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:40395a5fcd1abee49a5c7aaa98c29db393eedc835138560a588c47ec16156690", size = 33998032, upload-time = "2026-06-19T15:00:54.838Z" },
{ url = "https://files.pythonhosted.org/packages/7f/c4/bc41eb19b0fd0db868f4132920879019318d80cc522ad8f2bca4611af808/scipy-1.18.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8ca01e8ae69f1b18e9a58d91afead31be3cef0dd905a10249dac559ee15460a0", size = 35283333, upload-time = "2026-06-19T15:00:58.152Z" },
{ url = "https://files.pythonhosted.org/packages/53/a4/cbdeef6eb3830a8462a9d4ada814de5fc984345cc9ecf17cbec51a036f1e/scipy-1.18.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7a7f3b01647384dbc3a711e8c6778e0aabbe93959249fef5c7393396bcac0867", size = 35610216, upload-time = "2026-06-19T15:01:01.155Z" },
{ url = "https://files.pythonhosted.org/packages/80/4d/b2b82502b65f661d1b789c1665dcdf315d5f12194e06fc0b37946294ebae/scipy-1.18.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:6aa94e78ec192a30063a5e72e561c28af769dc311190b24fe91774eff1969709", size = 37418960, upload-time = "2026-06-19T15:01:04.155Z" },
{ url = "https://files.pythonhosted.org/packages/93/3e/902d836831474b0ab5a37d16404f7bc5fafd9efba632890e271ba952635f/scipy-1.18.0-cp314-cp314-win_amd64.whl", hash = "sha256:2d8bbdc6c817f5b4006a54d799d4f5bab6f910193cbb9a1ff310833d4d270f61", size = 37288845, upload-time = "2026-06-19T15:01:07.822Z" },
{ url = "https://files.pythonhosted.org/packages/b6/43/8d73b337a3bdb14daa0314f0434210747c02d79d729ce1777574a817dcf6/scipy-1.18.0-cp314-cp314-win_arm64.whl", hash = "sha256:18e9575f1569b2c54174e6159d32942e03731177f63dce7975f0a0c88d102f5b", size = 24988971, upload-time = "2026-06-19T15:01:11.076Z" },
{ url = "https://files.pythonhosted.org/packages/b4/b4/f11918b0508a2787031a0499a03fbe3546f3bb5ca05d01038c45b278c09a/scipy-1.18.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:f351e0dd702687d12a402b867a1b4146a256923e1c38317cbc472f6372b94707", size = 31399325, upload-time = "2026-06-19T15:01:13.723Z" },
{ url = "https://files.pythonhosted.org/packages/7b/d1/1f287b57c0ff0ee5185dff3946d92c8017d39b0e431f0ae79a3ff1859512/scipy-1.18.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:7c7a51b33ce387193c97f228320cf8e87361daa1bba750638677729598b3e677", size = 29092110, upload-time = "2026-06-19T15:01:16.908Z" },
{ url = "https://files.pythonhosted.org/packages/ff/1a/7b74eb6c392fdcb27d414c0e7558a6d0231eb3b6d73571f479bb81ea8794/scipy-1.18.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:84031d7b052a54fae2f8632e0ec802073d385476eb9a63079bce6e23ef9283d4", size = 20833811, upload-time = "2026-06-19T15:01:20.488Z" },
{ url = "https://files.pythonhosted.org/packages/7c/ad/f3941716320a7b9cb4d68734a903b45fe16eff5fb7da7e16f2e619304979/scipy-1.18.0-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:56abf29a7c067dde59be8b9a22d606a4ea1b2f2a4b756d9d903c62818f5dacce", size = 23396644, upload-time = "2026-06-19T15:01:23.364Z" },
{ url = "https://files.pythonhosted.org/packages/22/22/1446b62ffe07f9719b7d9b1b6a4e05a772833ae8f441fe4c22c34c9b250f/scipy-1.18.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ad44305cfa24b1ba5803cbbebf033590ccbac1aa5d612d727b785325ab408b0", size = 34079318, upload-time = "2026-06-19T15:01:26.002Z" },
{ url = "https://files.pythonhosted.org/packages/56/3b/b87da667098bb470fa30c7011b0ba351ee976dd395c78798c66e941665a3/scipy-1.18.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:945c1761b93f38d7f99ae81ae80c63e621471608c7eeead563f6df025585cd58", size = 35324320, upload-time = "2026-06-19T15:01:28.881Z" },
{ url = "https://files.pythonhosted.org/packages/f8/a1/c7932f91909759b0267f75fdea34e91309f96b895757534b76a90b6b4344/scipy-1.18.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:1a4441f15d620578772a49e5ab48c0ee1f7a0220e387110283062729136b2553", size = 35699541, upload-time = "2026-06-19T15:01:31.968Z" },
{ url = "https://files.pythonhosted.org/packages/f7/86/5185061a1fcc41d18c5dc2463969b3a3964b31d9ac67b2fb05d4c7ff7670/scipy-1.18.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:9aac6192fac56bf2ca534389d24623f07b39ff83317d58287285e7fbd622ff76", size = 37472480, upload-time = "2026-06-19T15:01:35.136Z" },
{ url = "https://files.pythonhosted.org/packages/31/8e/f04c68e39919a010d34f2ee1367fd705b0a25a02f609d755f0bfbc0a15fc/scipy-1.18.0-cp314-cp314t-win_amd64.whl", hash = "sha256:e40baea28ae7f5475c779741e2d90b1247c78531207b49c7030e698ff81cee3f", size = 37365390, upload-time = "2026-06-19T15:01:38.091Z" },
{ url = "https://files.pythonhosted.org/packages/d5/19/969dc072906c84dd0a3b05dcf57ea750936087d7873549e408b35cfc3f97/scipy-1.18.0-cp314-cp314t-win_arm64.whl", hash = "sha256:368e0a705903c466aa5f08eefb39e6b1b6b2d659e7352a31fd9e2438365be0f8", size = 25279661, upload-time = "2026-06-19T15:01:40.817Z" },
]
[[package]]
name = "shellingham"
version = "1.5.4"
@@ -1626,6 +1739,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/9b/24/84ce997e8ae6296168a74d0d9c4dde572d90fb23fd7c0b219c30ff71e00e/tbb-2021.13.1-py3-none-win_amd64.whl", hash = "sha256:cbf024b2463fdab3ebe3fa6ff453026358e6b903839c80d647e08ad6d0796ee9", size = 286908, upload-time = "2024-08-07T15:09:05.677Z" },
]
[[package]]
name = "threadpoolctl"
version = "3.6.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/b7/4d/08c89e34946fce2aec4fbb45c9016efd5f4d7f24af8e5d93296e935631d8/threadpoolctl-3.6.0.tar.gz", hash = "sha256:8ab8b4aa3491d812b623328249fab5302a68d2d71745c8a4c719a2fcaba9f44e", size = 21274, upload-time = "2025-03-13T13:49:23.031Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/32/d5/f9a850d79b0851d1d4ef6456097579a9005b31fea68726a4ae5f2d82ddd9/threadpoolctl-3.6.0-py3-none-any.whl", hash = "sha256:43a0b8fd5a2928500110039e43a5eed8480b918967083ea48dc3ab9f13c4a7fb", size = 18638, upload-time = "2025-03-13T13:49:21.846Z" },
]
[[package]]
name = "torch"
version = "2.3.1"