Fits normalizers over multi-hundred-million-row datasets, so the setup
pass's per-row Python overhead compounds fast: encode_secondaries
recomputed an O(K) prefix sum from scratch on every one of its 15
stick-breaking iterations, np.isin re-sorted the full train-event-id
array on every chunk, and pdg/material/process index lookups ran a
Python dict lookup per row. The normalizer-fit pass also computed
encode_secondaries's stick-logit and direction-rotation blocks in full
even though it only ever reads the mass/charge columns.
Replace the prefix-sum recompute with a single np.cumsum, add a
sorted_membership helper (searchsorted-based) in place of np.isin at
both the setup-pass and per-epoch call sites, vectorize the index
lookups via _vectorized_map_lookup, and add an opt-in phys_only path
so the setup pass skips the stick-breaking/rotation work it discards
anyway. All four changes are output-identical performance refactors,
backed by new unit tests plus the existing suite.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Checkpoints trained before commit 68fb99b (physical-property
conditioning, COND_DIM 8->15) saved a COND_DIM_BASE-wide cond
normalizer, fit before build_cond_features grew the extra physical
columns. Any inference against such a checkpoint under current code
(predict/rollout/router_gating) crashed broadcasting a 15-wide
cond_cont against an 8-wide mean/std.
In "embedding" mode those physical columns are never read by
ConditionEncoder, so padding the missing entries with mean=0/std=1 is
a safe no-op. "physical" mode reads them directly, so a mismatch there
still raises instead of silently normalizing garbage.
Adds model.conditioning = "physical" | "embedding": physical mode routes
particle mass/charge and material Z_eff/A_eff/density/X0/lambda_int through
small MLPs to replace the learned PDG/material embedding tables, so the
surrogate generalizes to PDG codes/materials outside the training vocab
instead of memorizing it. "embedding" stays available as the comparison
baseline (old checkpoints without the key default to it).
Stage 2 now regresses a secondary's mass/charge directly against a fixed
physics-derived target instead of a learned/snapped embedding, and uses no
snapping at inference — the model's raw predicted (mass, charge) is the
secondary's physical identity, including for its own further rollout steps.
A separate reporting-only nearest-known-PDG lookup (never fed back into the
model) populates output pdg columns / the embedding-mode rollout fallback.
giant/materials.py's table is populated with Geant4's own built-in NIST
constants (Z_eff, A_eff, density, X0, lambda_int), extracted directly from
the Geant4 11.4.1 build vendored in minicalosim via G4NistManager rather
than hand-typed literature values. G4_LYSO is left unfilled: confirmed (both
by runtime lookup and by searching minicalosim's history) that it's never
actually a constructed Geant4 material there, only documentation/UI color-map
text.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Brings in the rollout-validation fixes developed alongside Phase 2
(exact e_sec budget rescaling in decode_secondaries, filtering
synthetic termination rows out of load_rollout_vs_truth, Tier 4 truth
overlay, --energy-gev support in dwarf make-root) and reconciles them
with this branch's mixture-of-experts routing work: build_features/
build_models/dataset plumbing keep the ProcessRouter's proc_map/
proc_idx threading, and create_root_files.py's job_seed folds in both
the per-job seed derivation and the new energy_gev component.
A parquet that carries child_track_ids/e_sec but was never run through the
parent->child join lacks the per-secondary columns (sec_E_list/sec_pdg_list/
sec_dir_list). build_features would fall back to all-zero sec_cont/sec_pdg_idx,
collapsing every secondary to PDG index 0 and a constant energy fraction — a
broken Stage 2 that trained with no error (single-species validation tables).
Add an opt-in require_secondaries flag that raises when n_sec > 0 but the lists
are absent, and enable it on the training paths (StreamingStepsDataset and the
normalizer-fit pass). giant predict keeps the default False for Stage-1-only use.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
A parquet that carries child_track_ids/e_sec but was never run through the
parent->child join lacks the per-secondary columns (sec_E_list/sec_pdg_list/
sec_dir_list). build_features would fall back to all-zero sec_cont/sec_pdg_idx,
collapsing every secondary to PDG index 0 and a constant energy fraction — a
broken Stage 2 that trained with no error (single-species validation tables).
Add an opt-in require_secondaries flag that raises when n_sec > 0 but the lists
are absent, and enable it on the training paths (StreamingStepsDataset and the
normalizer-fit pass). giant predict keeps the default False for Stage-1-only use.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Routes on the physics process (Compton, phot, brems, ...) that ends a
step, supervised by a small classifier since process is a post-step
outcome unobservable at gate time. Threads a process label end-to-end
through the data pipeline (loader, build_features, dataset batches,
training loss/checkpointing) alongside the existing EnergyRouter.
Real data has steps with up to ~37 secondaries, but the n_sec head only
has K_MAX+1=16 classes. The unclamped label occasionally overflowed
cross_entropy's valid range and crashed CUDA training with
"unique_by_key: failed to synchronize: cudaErrorAssert". The
continuous secondary targets were already truncated to K_MAX slots;
only this label was missed.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Removes unused imports and an ambiguous variable name, narrows
Optional types before use so ty's flow analysis is satisfied, swaps
sum() over polars expressions for pl.sum_horizontal to avoid the
Literal[0] fallback type, and converts numpy bin edges to plain lists
before passing to matplotlib's hist (whose stub only accepts
Sequence[float]). Also applies ruff format across the repo, which had
drifted out of sync with the formatter.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
- energy_simplex_encode: warn when clipping post_E to pre_E discards
recorded edep/e_sec instead of silently zeroing them
- local/inv_local_frame_rotation: validate and normalize pre_dir instead
of silently assuming unit norm; raise on near-zero-norm rows
- train(): make --lr authoritative on resume instead of being silently
overwritten by the checkpoint's optimizer/scheduler state; print and
exit cleanly instead of silently training zero epochs when the
checkpoint already meets --epochs; truncate metrics.csv on a fresh
run instead of always appending
- dwarf update-manifest: check file existence for every manifest line,
not just ones whose gen/schema actually changed
- pyproject.toml: dev extra now pulls in convert+analysis so the
documented `uv sync --extra cpu --extra dev` + `pytest` actually
passes collection
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Replaces the independent log_delta_e/log_edep targets with 2 additive-log-ratio
coordinates over the deposit/secondary/post-energy simplex (fractions of pre_E
summing to 1), so edep + e_sec + post_E == pre_E holds by construction after
decoding (softmax) rather than being learned approximately. Requires e_sec
(secondary energy) as a new conditioning input and a steps_to_parquet.py pass
to derive it from child track first-step energies.
ruff removed unused imports across analysis.py and several test files.
ty caught a wrong dict[int, int] annotation on StreamingStepsDataset's
mat_map (materials are strings) and a real bug in steps_to_parquet.py
where --compression none passed None to polars' write_parquet, which
only accepts the literal "uncompressed". Also narrows a few
Optional-typed attributes (ddpm_schedule, Normalizer.mean/std) with
asserts and aligns __getitem__'s parameter name with torch's Dataset
base class.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
step_length already encodes |post_pos - pre_pos| by definition, so a raw
post_pos target would duplicate that magnitude and could drift inconsistent
with step_length during sampling. Instead add travel_dir, a unit vector
(local frame) giving only the direction of pre_pos->post_pos; post_pos is
reconstructed at inference as pre_pos + step_length * travel_dir, keeping
the two self-consistent. Target grows from 6D to 9D.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
lars