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>
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>
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.
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>
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>
lars