Commit Graph

120 Commits

Author SHA1 Message Date
lars b6ed618517 Run rollout validation for the mixture of experts at 50 GeV.
CI / Lint (ruff check) (push) Successful in 1m6s
CI / Format (ruff format) (push) Failing after 1m4s
CI / Type check (ty) (push) Successful in 1m9s
CI / Tests (push) Successful in 2m33s
CI / Bump version, build & publish wheel (push) Has been skipped
2026-07-24 09:59:30 +02:00
lars 4f785c43e6 Merge pull request 'Draft: Resolve "Wasserstein-GAN throwaway experiment (fast single-pass eval)"' (#14) from 3-wasserstein-gan-throwaway-experiment-fast-single-pass-eval into master
CI / Lint (ruff check) (push) Successful in 1m6s
CI / Format (ruff format) (push) Successful in 1m1s
CI / Type check (ty) (push) Successful in 1m6s
CI / Tests (push) Successful in 2m34s
CI / Bump version, build & publish wheel (push) Failing after 26s
Reviewed-on: #14
2026-07-23 15:32:48 +02:00
lars 8d46f27d1b Reformat rollout_validation.ipynb with ruff format
CI / Lint (ruff check) (push) Successful in 1m1s
CI / Format (ruff format) (push) Successful in 1m4s
CI / Type check (ty) (push) Successful in 1m2s
CI / Tests (push) Successful in 2m36s
CI / Bump version, build & publish wheel (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 58s
CI / Format (ruff format) (pull_request) Successful in 1m5s
CI / Type check (ty) (pull_request) Successful in 1m1s
CI / Tests (pull_request) Successful in 2m32s
CI / Bump version, build & publish wheel (pull_request) Has been skipped
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-23 14:38:49 +02:00
lars 9162ba7ae0 Split CI lint stage into parallel jobs
ruff check, ruff format check, ty check, and pytest now run as four
independent jobs instead of one sequential lint job gating test —
faster wall-clock CI since none of these checks depend on each other.
build still waits on all four before bumping/publishing.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-23 14:38:46 +02:00
lars faac19edf3 Quote the on: key in the CI workflow
CI / Lint & type check (push) Failing after 1m11s
CI / Tests (push) Has been skipped
CI / Lint & type check (pull_request) Failing after 59s
CI / Tests (pull_request) Has been skipped
CI / Bump version, build & publish wheel (push) Has been skipped
CI / Bump version, build & publish wheel (pull_request) Has been skipped
Unquoted on: gets parsed as the YAML 1.1 boolean true instead of the
string "on", so Gitea couldn't find a trigger key and didn't register
the workflow at all.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-23 14:35:55 +02:00
lars 4e77e3fc4b Add Gitea Actions CI pipeline
Three-stage workflow: lint/format/type-check (ruff, ty), pytest,
and (master pushes only) a version-bump + uv build + publish to the
Gitea package registry. Bump commit is tagged [skip ci] to avoid
retriggering the workflow.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-23 14:33:15 +02:00
lars 0d2967fa6b Add router gating diagnostic for MoE checkpoints
plot_router_gating visualizes soft expert gate weights vs. a continuous
routing axis (e.g. pre-step energy), binned into equal-population
quantiles and stacked to show the router's soft decision boundaries.
Wired into rollout_validation.ipynb as a new notebook-only section
that loads a checkpoint's Router directly, since gate weights aren't
present in rollout/predict parquet output.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-23 14:15:39 +02:00
lars 2e5268f403 Ignore the scratchpad working directory
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-22 11:22:39 +02:00
lars 44b0a92e67 Add WGAN-GP mode as a throwaway fast-eval experiment
Adds --mode wgan alongside flow/ddpm: both stages get a WGAN-GP
generator/critic pair (giant.model.wgan) instead of flow matching, so
inference is a single forward pass per stage rather than a 10-step ODE
integration — the fast-eval architecture noted in the roadmap.
predict/rollout auto-detect the mode from the checkpoint's model_config.
Best-checkpoint selection for wgan uses marginal-KL against the EMA
generators every epoch, since a critic loss isn't a monotone quality
signal. --router is not supported together with --mode wgan.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-22 11:20:28 +02:00
lars 4d190696fe Merge branch '7-condition-on-material-particle-physical-properties-not-learned-embeddings' into 'master'
Resolve "Condition on material + particle physical properties, not learned embeddings"

Closes #7

See merge request lbogner/giant!5
2026-07-22 10:38:59 +02:00
lars 06c9ad8e5f Fix crashes in physical-property conditioning edge cases
- make_seed_frontier only resolves particle mass/charge in "physical"
  mode, so "embedding"-mode rollouts no longer crash on a seed PDG code
  giant.particles can't resolve (the TERM_UNKNOWN_PDG gate now handles it).
- nearest_known_pdg skips unresolvable candidate PDG codes instead of
  raising and killing the whole rollout/predict run.
- predict/rollout fail with a clear message when a checkpoint predates
  the sec_phys normalizer, instead of a bare KeyError.
- validate_marginals' phys_kl degrades to NaN (matching the
  energy_fraction_kl pattern) instead of crashing when a validated batch
  has zero secondaries on either side.
- Correct CLAUDE.md's stale claim that the materials table is unfilled.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-20 10:36:25 +02:00
lars 68fb99bed8 Condition on material/particle physical properties instead of learned embeddings
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>
2026-07-17 15:12:54 +02:00
lars a6bb142a40 Merge branch 'analysis-streaming-rewrite' into 'master'
Rewrite analysis module as a lean, fully-streaming pipeline

See merge request lbogner/giant!4
2026-07-17 12:34:05 +02:00
lars 098dacb8cd Merge branch '4-prototype-a-mixture-of-experts-routing-tree-architecture' into 'master'
Add mixture-of-experts routing prototype for Stage 1 and Stage 2

Closes #4

See merge request lbogner/giant!3
2026-07-17 12:32:45 +02:00
lars 81ec225b86 Reimplement rollout-vs-truth comparison on the streaming analysis module
Merged in every non-analysis change from the MoE-prototype branch (routing,
training, data pipeline, streaming rollout output), keeping this branch's
lean streaming giant/analysis.py and rebuilding the rollout-vs-truth feature
natively on it instead of resurrecting the old numpy SampleCollection path.

- Add RolloutVsTruth, accepted anywhere Tier 1-3 functions take a predict-parquet
  source: decodes a giant rollout file and a held-out truth file into
  RAW_TARGET_NAMES space via a polars port of the forward local-frame rotation,
  fully streaming (no SampleCollection, no eager materialization).
- Add compute_rollout_vs_truth_observables_pl for Tier 4, reusing
  EventObservables (now backed by independent real_table/gen_table to support
  unequal rollout/truth event counts) so every existing shower-observable plot
  function works unchanged for both one-step and full-rollout comparisons.
- Update analysis/rollout_validation.ipynb to the new API and CLAUDE.md's
  architecture description; add test coverage for the new source type.
- Fix a pre-existing return-type mismatch in giant.rollout.rollout() (found by
  `ty check`): the on_chunk summary-dict branch didn't match the declared
  dict[str, np.ndarray] return type, now expressed as a RolloutSummary TypedDict.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-17 12:32:30 +02:00
lars 6744d05948 Merge branch '4-prototype-a-mixture-of-experts-routing-tree-architecture' into analysis-streaming-rewrite
# Conflicts:
#	giant/analysis.py
#	tests/test_analysis.py
2026-07-17 12:06:10 +02:00
lars b658f4d40e Merge branch 'phase2-secondary-prediction' into 'master'
Implement Phase 2: secondary particle prediction

See merge request lbogner/giant!1
2026-07-17 11:04:22 +02:00
lars 5a0e98c5d1 Add EMA weights, weight decay, step-based LR schedule, and grad-norm logging to training
Gives flow-matching sampling a cleaner EMA shadow copy to draw from (--ema-decay,
--weights raw|ema in predict/rollout), fixes the LR warmup/cosine schedule stepping
once per epoch even when an epoch is tens of thousands of steps, and caps the
per-epoch val-loss pass (--max-val-batches) so large val sets don't dominate epoch time.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-17 11:02:12 +02:00
lars 24a83486ec Scale auto batch-size estimate by MoE expert count during training
Training runs the full soft mixture (every expert over the whole batch),
so routed activation memory scales with the expert count; the old estimate
used one expert's dims and would overshoot free VRAM by a factor of
n_experts. Fold the expert count into n_blocks for the training path
(inference's top-1 dispatch still just partitions the batch, so one
expert's dims bound it).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-15 16:22:25 +02:00
lars f3fec8bcb3 Add ComposedRouter for multi-axis MoE gating
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>
2026-07-15 13:53:21 +02:00
lars 550dc679c7 Stream giant rollout output instead of buffering the whole run
_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>
2026-07-15 11:10:04 +02:00
lars ac5fbd14b4 Add PdgRouter for particle-type-based expert gating
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>
2026-07-15 10:13:07 +02:00
lars 05d5dee606 Apply ruff format after merging phase2-secondary-prediction
The merged proc_idx/proc_map plumbing wasn't run through ruff format
before merging; reflow only, no logic changes.
2026-07-15 10:00:36 +02:00
lars a5683517be Merge branch 'phase2-secondary-prediction' into 4-prototype-a-mixture-of-experts-routing-tree-architecture
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.
2026-07-15 09:59:43 +02:00
lars ea22b85df9 Run full rollout validation. 2026-07-13 14:32:36 +02:00
lars c81076c324 Support --energy-gev in dwarf make-root for the new minicalosim energy arg
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>
2026-07-13 13:59:09 +02:00
lars 670f57c309 Rescale secondary energies to exactly consume the e_sec budget
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>
2026-07-13 13:45:13 +02:00
lars 4c1250e246 Fix rollout edep mismatch and add truth overlay to Tier 4 observables
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>
2026-07-13 12:38:44 +02:00
lars 36bdbf7cc1 Add load_rollout_vs_truth to compare rollouts against held-out truth data
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>
2026-07-13 10:58:42 +02:00
lars 55c676fb9b Rewrite analysis module as a lean, fully-streaming pipeline
Replace giant/analysis.py's dual numpy-SampleCollection + polars paths with a
single polars-streaming implementation that produces the validation notebook's
plots directly from a `giant predict --coord local` parquet, sized for files
larger than RAM.

- Drop the numpy SampleCollection path (load_predicted_local, marginal_table,
  correlation_matrices, direction_alignment, constraint_report, plot_kl_bars)
  and the rollout observables; the 5 remaining plotters now take a parquet
  path / LazyFrame and stream internally.
- Rewrite compute_event_observables_pl to aggregate in parallel streaming
  polars (post-pos reconstruction as expressions) instead of a serial
  pyarrow-batch + numpy loop, fixing a pre-existing OOM (holistic median +
  323M-row join in the bin-edge sizing). Medians are approximated from a
  streaming log-bin histogram with within-bin interpolation.
- Keep every full-file scan narrow (few columns): on a file larger than RAM,
  peak mmap memory, not scan count, is the binding constraint. Marginals run
  one dim at a time (~15GB peak) rather than a combined all-dims pass (OOM).
- Update analysis/validation.ipynb to the path-based API; delete the
  analysis/export_*.py and compare_ode_steps_*.py one-off scripts.
- Rewrite tests/test_analysis.py around parquet fixtures with an inline numpy
  oracle; add correlation/streaming-plotter and approx-median coverage.

Verified end-to-end on the 32GB predict file: full notebook completes at
~25GB peak (no OOM); event rollup runs at ~13 cores.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-10 14:47:46 +02:00
lars 740ebdf6b6 Drop orphaned child tracks instead of nulling secondary targets
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>
2026-07-09 14:37:28 +02:00
lars 028fa13b7b Drop orphaned child tracks instead of nulling secondary targets
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>
2026-07-09 14:36:06 +02:00
lars a8fc019e32 Derive a unique per-job seed for minicalosim shard generation
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>
2026-07-09 10:09:23 +02:00
lars 5dc086e35a Derive a unique per-job seed for minicalosim shard generation
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>
2026-07-09 10:08:04 +02:00
lars 339d2fb473 Error on missing secondary lists instead of silently zeroing Stage-2 targets
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>
2026-07-09 09:09:14 +02:00
lars f387178dbf Error on missing secondary lists instead of silently zeroing Stage-2 targets
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>
2026-07-09 09:02:51 +02:00
lars 0b3ece52ed Add ProcessRouter for physics-process-based expert gating
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.
2026-07-08 16:15:09 +02:00
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 a14a4f973a Apply ruff format across the codebase
Whitespace-only reflow (line wrapping, blank lines between defs); no
logic changes.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-08 14:44:53 +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