Works through docs/v0.3.0-followups.md item by item, closing the gap
between the design doc and the shipped v0.3.0-stage2-autoregressive code:
1. validate.py: 7-tuple batch unpacking, sample_stage1/sample_stage2
dispatch, stage-2 particle-type-class marginal.
2. Stage-prefixed --stage1-*/--stage2-* CLI flags for train/new-run.
3. Thread stage2_model.k_max through loader/transforms/dataset/pipeline/
train instead of the hardcoded K_MAX constant.
4. Mixed conditioning.particle.type / conditioning.material.type support
end-to-end (data pipeline + dwarf warm-cache).
5. conditioning.share_stages = true: one shared ConditionEncoder instance
across both stages.
6. stage2_model.generator = "ddpm" formally deferred into design doc §11.2
(was silently unimplemented).
7. giant predict/rollout: implement conditioning.*.type = "onehot" via the
checkpoint's saved pdg_topn_map/mat_topn_map.
8. network.py's checkpoint-path model_config migration now fails loudly on
non-zero legacy expert_hidden_dim/expert_n_blocks, matching config.py's
TOML-load path (§4.2).
9. validate_config now rejects stage2_model.n_sec.mode = "truth" for a
rollout-capable checkpoint (§9).
Also cleared all pre-existing `ty check` noise (44 -> 0 diagnostics),
mostly a test-helper dict-unpack pattern that made every unrelated
constructor keyword look like a type error.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Builds the shared top-N-plus-other PDG/material maps (pooling both primary
and secondary occurrences for PDG, directly targeting the meeting's
species-collapse failure mode) and wires up conditioning.{particle,material}
= "onehot" plus stage2_model.particle_type.target in ("onehot", "embedding")
end-to-end: setup-cache persistence, Stage2OneShot's type_head (flow/ddpm)
vs. folded+ST-Gumbel-relaxed adversarial slice (wgan), and the corresponding
CE/MSE training losses. particle_type.target = "physical" stays byte-for-byte
unchanged, keeping the v0.2 migration shim's bit-identical guarantee intact.
giant predict/rollout fail loudly on a onehot/embedding checkpoint until
full decode support lands in step 6.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The whole point of conditioning="physical" is generalizing to a
species/material outside the training menu, but two independent code
paths still hard-required training-vocab membership:
- giant/data/transforms.py: build_cond_features unconditionally raised
KeyError on an out-of-vocab pdg/material. _vectorized_map_lookup
gains a strict=False mode (dummy index instead of raising), used only
under conditioning="physical" where ConditionEncoder never reads
cond_cat anyway; "embedding" mode is untouched and still raises,
since cond_cat IS the conditioning signal there.
- giant/rollout.py: the known_pdg termination gate still killed a track
on step 1 for any pdg outside pdg_map, regardless of conditioning
mode. Now skipped entirely under conditioning="physical".
- giant/model/network.py: PdgRouter/ProcessRouter always build their
own training-vocab nn.Embedding independent of conditioning, silently
reintroducing the same limitation at the routing layer. build_models
now raises loudly if conditioning="physical" is paired with either
router type, rather than silently building a model that can't
generalize the way it claims to.
This unblocks the held-out-species/material generalization experiment
against the multi-material dataset (see CLAUDE.md roadmap). Each fix
has a regression test, including an end-to-end rollout test seeded
with a resolvable-but-out-of-vocab PDG code.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
- log_transform / _validate_unit_pre_dir now raise on non-finite input
instead of letting a NaN row silently poison the persisted normalizer
cache (norm < 1e-6 was always False for NaN, so the existing guard
never caught it).
- encode_secondaries warns when a row's secondary energies cumulatively
exceed e_sec, instead of silently saturating the overflowing slot's
stick-breaking logit via the _EPS floor.
- EVENT_ID_FILE_STRIDE overflow now raises instead of silently colliding
two files' event ids together (reintroducing train/val leakage).
- make_event_split(val_fraction=0.0) now actually holds out nothing,
instead of always forcing at least 1 validation event.
- setup_cache.save() is now serialized with a flock, since two
concurrent writers (a real scenario on this repo's shared
portal/condor machines) could otherwise race and silently drop one
writer's freshly-computed cache section.
- Documented (no behavior change) the pre_dir ≈ -ẑ antipodal rotation
singularity in _rodrigues_axis, which is real but inherent to any
single-valued local-frame convention.
Each fix has a regression test.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The streaming update re-derived two full (B, F) arrays from the
running mean (once before updating it, once after) plus an elementwise
product — five passes over each chunk and three temporary arrays, to
maintain a mean/variance that's tiny in width (COND_DIM=15 at most).
Reformulate as Chan/Golub/LeVeque's parallel-variance algorithm:
compute the chunk's own local mean/M2 (independent of the running
state) and merge it in with an O(F) combination formula. Same
streaming interface and output (identical to ~1e-14, float64 rounding
noise), ~40% faster per update() call on a benchmark chunk.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
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