Commit Graph

82 Commits

Author SHA1 Message Date
lars bac541240f Add mixture-of-experts routing prototype for Stage 1 and Stage 2
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>
2026-07-08 15:17:19 +02:00
lars af8dce53a7 Update CLAUDE.md and README for the implemented Phase 2 model
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>
2026-07-08 13:42:53 +02:00
lars 3faa272562 Add fast slab lookup for the GeometryOracle, replacing knn as the default
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>
2026-07-08 11:42:12 +02:00
lars 436d9fa4d4 Recalibrate batch-size estimate for the post-Phase-2 model size
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>
2026-07-08 11:42:03 +02:00
lars 980b6ae7da Weight Stage-2 secondary loss equally between direction and type-embedding dims
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>
2026-07-08 11:41:24 +02:00
lars 4ee75d0042 Clamp n_sec classification label to K_MAX
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>
2026-07-08 10:11:17 +02:00
lars 26a176aeaa Add autoregressive shower rollout driver
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>
2026-07-08 10:10:36 +02:00
lars 549c051417 Detach Stage-2 type-embedding target to stop self-referential collapse
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>
2026-07-06 14:19:37 +02:00
lars 9a9e165b7d Wire up n_sec/species/energy-fraction validation for Stage 2
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>
2026-07-06 14:06:21 +02:00
lars 0c96ddf765 Wire up predict CLI to load and run the Stage-2 sec_decoder
`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>
2026-07-06 14:06:07 +02:00
lars 4ba419ebe4 Merge energy-conservation-poc into phase2-secondary-prediction
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>
2026-07-06 12:18:30 +02:00
lars 1a22ae022c Add energy-conservation PoC ODE-step comparison scripts
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>
2026-07-06 11:58:25 +02:00
lars f0cb41477f Add --comment option to predict, recorded in YAML sidecar
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-06 09:37:48 +02:00
lars 25718f175e Fix ruff, ty, and pytest failures; apply ruff format
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>
2026-07-02 16:59:22 +02:00
lars e5bf7c51cb Fix silent failure modes surfaced by extensive code review
- 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>
2026-07-02 16:54:02 +02:00
lars 857d8b315f Show VERSIONS.md reason extracts in dwarf status
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>
2026-07-02 11:13:35 +02:00
lars 57e044e347 Add file counts and reference tracking to dwarf status
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>
2026-07-02 11:09:47 +02:00
lars 7761cd21b6 Color-code dwarf status output by tree level
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>
2026-07-02 10:34:06 +02:00
lars c0746c40e0 Add disk usage summary to dwarf status
Shows per-schema, per-gen (raw/processed split), per-kind, and grand
total sizes so the dataset tree's footprint is visible at a glance.
2026-07-02 10:10:51 +02:00
lars b6e0fa59a4 Merge origin/energy-conservation-poc (--to/--gen flags, train date-prefix)
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.
2026-07-02 09:57:20 +02:00
lars 917835e182 Fold --to/--gen dataset-versioning flags into the dwarf CLI
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>
2026-07-02 09:57:17 +02:00
lars d5853d5a75 Unify dataset/tooling scripts into a single dwarf Typer CLI
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>
2026-07-02 09:52:37 +02:00
lars 305e436783 Prefix default train output dir with current date
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-07-01 16:15:19 +02:00
lars 091b23a60f Add --to flag for bump-gen/bump-schema and --gen flag for update-manifest
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>
2026-07-01 15:58:26 +02:00
lars 7b37b284f8 WIP: energy-conservation PoC analysis/transforms updates 2026-07-01 13:58:44 +02:00
lars b25c6967ab Fix giant.analysis import after Phase 2 dataset API changes
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>
2026-06-29 14:17:40 +02:00
lars e6e0eb22bf Implement Phase 2: secondary particle prediction
Two-stage factorisation: Stage 1 predicts 9D primary kinematics + n_sec
classification head (COND_DIM reduced to 8, dropping n_sec/e_sec inputs);
Stage 2 (SecondaryDecoder) generates K_MAX=15 secondary slots via masked
flow matching over (stick_logit, local_dir, type_emb) conditioned on Stage 1
output. Joint training with combined loss L_s1 + λ_nsec*L_nsec + λ_s2*L_s2.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-29 11:34:31 +02:00
lars c627142135 Route predict output to UUID-named parquet with YAML reference sidecar
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>
2026-06-29 09:57:14 +02:00
lars caaf3c6102 Restrict holdout overlap check to holdout vs dev/full only
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>
2026-06-26 11:24:05 +02:00
lars f56a45690d Add update-manifest and create-manifest subcommands to bump_dataset_version
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>
2026-06-26 11:03:32 +02:00
lars 6b85e2c90b Add --copy mode to migrate_geant_steps.py
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>
2026-06-25 17:20:13 +02:00
lars 5be91e0d17 Expose dataset/conversion scripts as uv entry points
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>
2026-06-25 16:51:47 +02:00
lars 320365606a Add tooling for a versioned geant_steps dataset layout
Introduces raw/<kind>/<gen>/<detector>/shard-NNN.root and
processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet as the dataset
convention, plus scripts to operate on it: migrate_geant_steps.py for the
one-time move into this layout, bump_dataset_version.py to cut new
gen/schema versions with a logged reason, steps_to_parquet_parallel.py to
convert ROOT shards to parquet in parallel and place them correctly, and
create_root_files.py to generate new ROOT shards via a minicalosim
executable. The loader gains .manifest file support so pools/ (train/dev/
holdout shard lists) can be passed straight to `giant train`.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-25 16:46:48 +02:00
lars 8475199609 Encode edep/secondary/post energy as a conservation-constrained simplex
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.
2026-06-25 16:01:13 +02:00
lars 64c6bd1cef Add photon edep export scripts and per-step presentation plots
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-25 13:36:54 +02:00
lars bdac36b203 Allow steps_to_parquet.py to accept multiple ROOT input files 2026-06-24 19:28:50 +02:00
lars eea9a24d9f Add export script for ETP group-update presentation plots
One-off script (like export_validation_plots.py / export_event_observables.py)
that writes a 3x3 marginals grid, post_dir/travel_dir norm histograms, and
vector-PDF copies of the KL-bars/photon-edep/event-level plots directly into
the thesis-presentations repo's images/ folder.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-24 18:18:26 +02:00
lars b4ce04e772 Add mean/median deposited energy and step length plots per event
Move ipykernel into the analysis extra instead of a separate
dependency group, since it's needed wherever analysis plotting runs.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-24 12:46:17 +02:00
lars c802a59033 Add export script for Tier 4 event-level/pdg-share plots
One-off script mirroring export_validation_plots.py, used to export the
new event-level and pdg-contribution-share plots into the knowledge-base
attachments folder.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-24 12:42:49 +02:00
lars 2dba1edda6 Run validation incompletely 2026-06-24 10:53:38 +02:00
lars 139d95631f Add pdg energy/length contribution pie plots
pdg_contribution_table_pl sums real/generated total deposited energy and
total step_length per pdg species over the whole file (pure lazy polars
group_by, no post_pos reconstruction needed for these scalars). Adds
plot_pdg_energy_share/plot_pdg_length_share, each rendering two pies
(real vs generated) so the per-species breakdown can be compared directly,
plus a matching notebook section.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-22 16:55:51 +02:00
lars 453f9f9e20 Add total length traveled per event to event observables
sum(step_length) per event_id, alongside the existing total deposited
energy, since path length and energy deposit aren't interchangeable once
tracks scatter. Adds plot_total_length and a matching notebook cell.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-22 16:47:51 +02:00
lars 6d839a983f Rework validation notebook with markdown sections and Tier 4 plots
Adds section headers and explanatory comments throughout the existing
tiers, plus new cells running compute_event_observables_pl and the four
event-level shower plots. Notebook outputs reflect the user's own re-run.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-22 16:42:03 +02:00
lars f81263fa6e Add event-level shower observables to giant.analysis
Aggregates giant predict --coord local output per event_id into total
deposited energy, longitudinal/transverse shower profiles, and shower-max
depth, reconstructed into world-frame physical units (mm, MeV). Streams the
file in two polars passes rather than building a SampleCollection, since
per-event sums would be corrupted by row subsampling on these large files.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-22 15:52:47 +02:00
lars 6fc68fe1aa Export plots for knowledge base 2026-06-22 14:41:45 +02:00
lars 936249a8a0 Rerun validation 2026-06-22 14:40:58 +02:00
lars ae1b565ca3 Add KL bar plots and sample_frac to load_predicted_local; ignore root parquet scratch files
Adds plot_kl_bars/plot_kl_bars_pl (numpy/polars variants) for ranking which
target dimension or pdg/material stratum drives KL regressions, with the
same kl*n group capping as plot_marginals. Switches existing histogram
plots to step-type/log-scale. Adds sample_frac to load_predicted_local for
subsampling large predict parquets.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-22 11:52:41 +02:00
lars 9d7b1ca9bb Buffer predict rows across row-group boundaries before inference
Inference batches were capped at the source parquet's row-group size
(e.g. 122,880 rows) because the old loop only sliced within a single
chunk read from disk. --batch-size (including auto) had no effect once
it exceeded that, leaving most estimated GPU memory unused. Accumulate
rows across row groups and files into a buffer and slice exactly
batch-size pieces off it, so inference always uses the requested batch
size regardless of how the file happens to be chunked.
2026-06-22 09:19:56 +02:00
lars 903227d2df Skip rows with unknown PDG codes during predict
Checkpoints have a fixed PDG embedding vocab sized at train time, so a
code unseen during training has no embedding index. Drop those rows
and report a per-code skip count instead of raising a KeyError.
2026-06-22 09:06:05 +02:00
lars c057693d90 Add tqdm progress bar to predict 2026-06-22 08:53:04 +02:00