Route on several independent axes at once (e.g. energy x pdg), each with
its own expert count and hyperparameters. The joint gate is the outer
product of per-axis softmax gates, so it stays a partition of unity and
top1/balance_loss factor per-axis. Config uses flat axis{i}_{field} keys
in model.router (TOML/CLI friendly), also settable via repeatable
--router-axis flags.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
_Recorder previously accumulated every generated step across all events/
tracks/steps in Python lists, materialised once at the end and written
via a single pq.write_table — memory scaled with n_events * max_steps *
avg_tracks_per_event. rollout() now takes an optional on_chunk callback
that streams each non-empty batch immediately (fixed per-key dtypes via
_RECORD_DTYPES keep every chunk's table schema identical, which
pq.ParquetWriter requires across writes); giant rollout wires this to an
incrementally-written ParquetWriter, mirroring the row-group streaming
giant predict already does on its input side. Without on_chunk, rollout()
keeps its old buffered return for existing callers/tests.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Routes on the pre-step PDG code, which — unlike ProcessRouter's process
label — is already known at gate time (a conditioning input), so no
supervision is needed and classify_loss falls back to the zero default.
Generalizes EnergyRouter's soft-turn-on-then-Voronoi trick from a 1-D
distance to a small learned PDG embedding space: its own embedding table
maps each PDG code to a point, and n_experts learnable centers partition
that space.
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.
run_pbwo4/run_sampling now accept a trailing energy_GeV positional arg;
thread it through plan/run/seed so datasets like pbwo4_10gev can be
generated at non-default beam energies.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
decode_secondaries's stick-breaking only guarantees valid secondary slots
sum to <= e_sec, leaving a shortfall that rollout.py silently dumped into
that step's edep. Rescale the valid slots by one common per-row factor
instead, so they sum to exactly e_sec whenever n_sec > 0: this spreads any
shortfall proportionally across all secondaries rather than concentrating
it in whichever slot is last by energy rank (which would let that one
low-energy secondary balloon and distort the shower's topology). Rows
where every valid slot decodes to ~zero fall back to an even split.
n_sec == 0 rows are unchanged (still nothing to carry the budget, so
rollout.py's edep top-up still applies there) — narrowed the related
caveat in load_rollout_vs_truth's docstring to just that case.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
load_rollout_vs_truth was including rollout.py's synthetic termination-
bookkeeping rows (escaped/unknown_pdg/energy_cutoff/max_steps) unfiltered:
these carry step_length=0 and edep=pre_E dumped in one row for shower-level
energy conservation, not a real per-step value, and nearly doubled the
apparent mean edep in a repro. _load_world_frame_side now drops them, keeping
only real generated steps (continuing or natural_end). Also documents the
remaining, unfixable difference: rollout's edep on real steps absorbs any
secondary-energy budget Stage 2 didn't allocate, which truth's edep never does.
Adds compute_truth_observables, the truth-schema counterpart to
compute_rollout_observables, so the Tier 4 event-level plots
(plot_rollout_longitudinal/transverse/total_energy) can overlay a real
reference computed directly from load_rollout_vs_truth's own truth file,
without needing a separate paired giant predict --coord local file. Shares
the depth/transverse binning core with compute_rollout_observables via a new
_event_axis_depth_transverse helper.
Updates rollout_validation.ipynb's Tier 4 section to use this reference and
points ROLLOUT_FILE/TRUTH_FILE at a real prediction/shard pair.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Extends the Tier 1-3 SampleCollection diagnostics (marginals, correlations,
pairwise, direction alignment, constraints) to work on a full autoregressive
giant rollout shower checked against an independent ground-truth steps file,
rather than only paired giant predict --coord local output. The two files
are unpaired (different lengths, own conditioning), so SampleCollection
gains optional *_gen fields and _group_labels/marginal_table/plot_marginals/
plot_pairwise build independent real/gen masks instead of assuming one.
Adds analysis/rollout_validation.ipynb, a sibling of validation.ipynb built
around this workflow.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
A listed child_track_id can fail to match any first-step row (e.g. a
secondary absorbed below the tracking threshold at birth). The
parent->child left join in _add_secondary_attributes left these as
nulls, which silently became NaN once the parquet round-tripped
through the loader's float32 padding — poisoning every later secondary
slot in that step via the cumulative "remaining budget" in
encode_secondaries, while e_sec quietly undercounted and n_sec (from
len(child_track_ids)) overcounted relative to the actual lists.
Drop orphans from both the per-secondary lists and child_track_ids
itself so downstream counts stay consistent, and thread the per-file
orphaned count back through convert_steps_to_parquet so both the
sequential and --jobs>1 batch paths in `dwarf convert` can report an
aggregate total instead of relying on grepping printed output.
Also floors encode_secondaries' slot-0 budget to _EPS (matching the
i>0 branch), fixing a harmless but noisy 0/0 divide warning on
zero-secondary steps.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
A listed child_track_id can fail to match any first-step row (e.g. a
secondary absorbed below the tracking threshold at birth). The
parent->child left join in _add_secondary_attributes left these as
nulls, which silently became NaN once the parquet round-tripped
through the loader's float32 padding — poisoning every later secondary
slot in that step via the cumulative "remaining budget" in
encode_secondaries, while e_sec quietly undercounted and n_sec (from
len(child_track_ids)) overcounted relative to the actual lists.
Drop orphans from both the per-secondary lists and child_track_ids
itself so downstream counts stay consistent, and thread the per-file
orphaned count back through convert_steps_to_parquet so both the
sequential and --jobs>1 batch paths in `dwarf convert` can report an
aggregate total instead of relying on grepping printed output.
Also floors encode_secondaries' slot-0 budget to _EPS (matching the
i>0 branch), fixing a harmless but noisy 0/0 divide warning on
zero-secondary steps.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Concurrent job launches in create_root_files.py can start within the
same wall-clock second, and minicalosim's default seed falls back to
time(NULL) in that case — so two "independent" shards could silently
get identical RNG state and produce byte-identical physics. Requires
the companion MINICALOSIM_SEED env-var support in the minicalosim repo.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Concurrent job launches in create_root_files.py can start within the
same wall-clock second, and minicalosim's default seed falls back to
time(NULL) in that case — so two "independent" shards could silently
get identical RNG state and produce byte-identical physics. Requires
the companion MINICALOSIM_SEED env-var support in the minicalosim repo.
Co-Authored-By: Claude Sonnet 5 <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>
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.
Both stages can now route through a pluggable Router (EnergyRouter as the
first implementation, a soft turn-on gate over pre-step log-energy) into
several small ExpertTrunks instead of one monolithic trunk. Trains as a
differentiable soft mixture and dispatches to a single expert per row at
eval time, which is the source of the per-call speedup this prototype is
after (issue #5's ~10x native-Geant4 budget). Disabled by default, so
existing configs/checkpoints are unaffected; build_models() centralizes
routed-vs-monolith construction across train/predict/rollout.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Bring the docs in line with the current two-stage code: energy ALR
simplex output, 8D conditioning (n_sec/e_sec now predicted, not given),
the SecondaryDecoder stage, and the shower rollout + geometry oracle.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
miniCaloSim's detector is a stack of planar layer slabs along one axis, so
material/layer_id are a pure function of depth. The new "slab" method
exploits this with an exact O(log #segments) binary search over
depth-axis segment boundaries, instead of a nearest-neighbour search over
hundreds of thousands of reference points — much cheaper per call, which
matters since the oracle is queried on every autoregressive rollout step.
"knn"/"svm" remain as fallbacks for non-slab geometries.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The estimate_batch_size(training=True) calibration point was measured on
the pre-Phase-2 architecture (hidden_dim=512). Re-measured against the
current hidden_dim=1024 stack (Stage-2 secondary decoder + n_sec head
included): ~29696 batch size at ~7683 MiB VRAM.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The masked flow-matching loss for the secondary decoder averaged uniformly
over all 20 per-slot dims, letting the 16 type-embedding dims outvote the
4 physically-interesting ones (stick-break logit + direction). Split the
two blocks and average each over its own width before summing, so they
contribute with equal weight regardless of EMB_DIM.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
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>
Closes the loop from single-step prediction into full showers:
- giant/geometry.py + `dwarf build-geometry-oracle`: learn position ->
(material, layer_id) from data (KNN/SVM) to supply the conditioning the
surrogate does not predict; flag detector escape by NN distance.
- giant/rollout.py: breadth-first batched frontier that steps all active
tracks, spawns secondaries as new tracks, and terminates 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 exactly.
- `giant rollout` CLI: seed from real events (argmax pre_E), load checkpoint,
write a world-frame steps parquet + YAML sidecar.
- giant/analysis.py: compute_rollout_observables + plot_rollout_* for
single-sided longitudinal/transverse/total-energy shower profiles;
analysis/export_rollout_observables.py driver.
- scikit-learn added as an optional `geometry` extra (lazy-imported).
- Tests: tests/test_geometry.py, tests/test_rollout.py.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The shared PDG embedding table was used, un-detached, as the regression
target for the Stage-2 flow-matching loss. Since that tensor becomes x1
in u_t = x1 - x0, gradients could pull the embedding table itself toward
the decoder's predictions instead of the decoder learning to match the
table, risking species-embedding collapse and degrading the
nearest-neighbor species decode at inference.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
validate_marginals only ever checked Stage-1 primary marginals.
Extend it to optionally accept sec_decoder and report n_sec
classification accuracy + count distribution, secondary species
distribution, and per-slot energy-fraction marginals (real vs.
generated, each restricted to its own valid-slot mask). train.py's
periodic validation call now passes sec_decoder through.
Also fixes build_features looking up a "sec_pdg_idx" key that nothing
ever populated (the loader only ever produces "sec_pdg_list", raw PDG
codes) — the condition gating real secondary-target encoding was
therefore always false, so Stage 2 has been training on all-zero
sec_cont/sec_pdg_idx targets. Maps sec_pdg_list through pdg_map to
build sec_pdg_idx properly; this is also what makes the new species
validation meaningful rather than trivially degenerate.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
`giant predict` only ever ran Stage 1, echoing ground-truth n_sec instead
of predicting it — Phase 2 training already produced a joint checkpoint
but nothing consumed the sec_decoder half of it. Loads sec_decoder
alongside the Stage-1 model (filtering model_config per-model, since
splatting it whole into either constructor breaks on the other's
sec_slot_dim/k_max-only keys), runs sample_secondaries + PDG snapping in
--coord global mode, and appends predicted n_sec/species/energy/direction
columns to the output parquet.
Also fixes decode_secondaries rotating raw (non-unit) flow output straight
into world frame without normalizing first — a rotation preserves
magnitude, so un-normalized ODE output produced non-unit secondary
directions, caught via an end-to-end smoke test.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Brings the energy-conservation PoC work (dwarf CLI unification, dwarf
status improvements, predict --comment, ODE-step comparison scripts,
predict-parquet-only analysis refactor) onto the Phase 2 branch.
Conflict resolution:
- giant/analysis.py: took the energy-conservation-poc version wholesale.
That branch deliberately removed the live checkpoint+sampler diagnostics
path (ModelBundle/load_model_bundle/make_val_loader/collect_samples) in
favor of reading `giant predict --coord local` parquet output. Phase 2's
only edits to this file adapted the removed path to the new dataset API,
so nothing Phase-2-specific is lost; no external code called those funcs.
Fixes for pre-existing breakage surfaced by the merge (both predate it):
- giant/cli.py: predict's `_process` unpacked build_features into 5 values,
but Phase 2 made it return 8 (added n_sec/sec_cont/sec_pdg_idx). Expanded
the unpack; `giant predict --coord local` would have crashed otherwise.
- tests/test_steps_to_parquet.py: Phase 2 renamed _add_secondary_energy ->
_add_secondary_attributes without updating this test. Renamed the calls
and extended the fixture with the pdg/pre_d{x,y,z} columns the expanded
function reads; e_sec assertions unchanged.
- analysis/compare_ode_steps_energy_conservation.py: E731 lambda assignment
(added in the un-linted final PoC commit) rewritten as a def.
ruff, ty, and pytest (179 passed) all green.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add analysis scripts for the 10-vs-20 flow-matching ODE-step ablation on the
energy-conservation PoC predict outputs:
- compare_ode_steps_energy_conservation.py: per-event energy-budget table +
20-step plots and the 10-vs-20 overlay.
- compare_ode_steps_kl.py: per-step marginal KL(real||gen) per target dim over
fixed shared bins, so the two runs are directly comparable dim-by-dim.
Also commit export_energy_conservation_poc.py (the baseline event-level budget
export) and repoint validation.ipynb at the PoC predict file at sample_frac=1.0.
Co-Authored-By: Claude Opus 4.8 <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>
Parses the gen/schema log lines apply_bump() writes to VERSIONS.md and
prints a truncated reason under each gen/schemaN row, so `dwarf status`
answers "why does this version exist" without opening the changelog.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Shows per-directory file counts throughout the tree, plus a referenced
count for raw/ (matched against any same-named parquet under processed/)
and each schemaN dir (matched against pools/*.manifest entries).
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Each row (kind header, gen, raw/processed, schema, root totals) gets a
distinct ANSI color so the hierarchy is easier to scan. Disabled when
stdout isn't a TTY or NO_COLOR is set.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Both commits' changes were already folded in by hand in the previous
commit; this merge just records the shared history so the branches
reconcile cleanly.
origin/energy-conservation-poc grew bump-gen/bump-schema --to and
update-manifest --gen flags (091b23a) plus a train output-dir date
prefix (305e436) after the dwarf unification was written locally.
Reconcile: bring plan_bump_gen/plan_bump_schema/plan_update_manifest's
target/target_gen support into the plain-function (argparse-free) form,
thread --to/--gen through scripts/dwarf.py's bump-gen/bump-schema/
update-manifest commands, and take giant/cli.py's date-prefix change
and the associated tests as-is.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Replace the five separately-hyphenated uv entry points (steps-to-parquet,
steps-to-parquet-parallel, migrate-geant-steps, bump-dataset-version,
create-root-files) plus the unregistered hparam_scan.py with one `dwarf`
command exposing convert/migrate/bump-gen/bump-schema/status/
update-manifest/create-manifest/make-root/hparam-scan as subcommands.
Each scripts/*.py module now only holds argparse-free business logic;
scripts/dwarf.py wires it up with Typer, matching giant/cli.py's style.
`dwarf convert` merges the old serial/parallel conversion scripts behind
a --jobs flag (default 1: sequential with plain -o; >1: dataset-layout
fan-out via subprocess).
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
bump-gen and bump-schema now accept --to genN/schemaN to target a specific
version instead of always auto-incrementing. update-manifest gains --gen genN
to repoint the gen component of manifest paths (combinable with --schema).
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
train_val_split was removed from giant.data.dataset in favor of
make_event_split + StreamingStepsDataset (event-based split, streaming
batches), and build_features grew secondary-prediction outputs.
make_val_loader and collect_samples still referenced the old API.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
When no --out is given, predictions are written to a UUID-named parquet
(/ceph central store for ceph inputs, sibling directory otherwise) and a
YAML reference file is created next to the checkpoint recording the output
path, dataset, checkpoint, and timestamp. Adds pyyaml as a core dependency
and unit tests for the two new helper functions.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
dev and full are allowed to share files — only holdout must be strictly
isolated. When creating dev or full, only compare against holdout.manifest;
when creating holdout, compare against all other manifests in the dir.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
update-manifest rewrites the schemaN component in existing manifest files to a
specified or auto-detected highest schema, verifying all target files exist before
writing. create-manifest builds a new manifest from explicit parquet file paths,
supporting --pool/--type (full|holdout|dev) to derive the output path from root,
and enforcing holdout isolation by checking for cross-manifest overlap whenever a
holdout manifest is involved.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Lets the migration run while another process still has the original files
open for reading: --copy uses shutil.copy2 instead of move, and skips the
now-empty-directory cleanup since the legacy train/ etc. dirs stay populated
by design.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
scripts/ is now a proper package (scripts/__init__.py, added to the wheel's
packages), with each script registered under [project.scripts] using its
bare dashed name (e.g. `uv run migrate-geant-steps`). Tests now import these
modules normally instead of loading them by file path.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
lars