b1cf9d345dcfe491dc569a575eded8a8a561002d
13 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
|
|
878e9ddca3 |
Delete docs/v0.3.0-design.md and strip all references to it
CI / Format (ruff format) (push) Failing after 28s
CI / Lint (ruff check) (push) Successful in 29s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 33s
CI / Type check (ty) (push) Successful in 37s
CI / Format (ruff format) (pull_request) Failing after 37s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 37s
CI / Tests (pull_request) Successful in 2m49s
CI / Tests (push) Successful in 2m55s
The design doc and its followups doc are no longer needed as a live reference now that the v0.3.0 redesign is implemented — comments and docstrings across the codebase cited it extensively (file path, "design doc §X.Y", "decision N", or bare "§X.Y" section numbers) as design rationale. Removed docs/ and edited every citing comment/docstring to drop the now-dangling reference while keeping the substantive explanation next to it. CLAUDE.md's v0.3.0 roadmap bullet loses its trailing pointer to the deleted file. Verified: no remaining "docs/v0.3.0", "design doc", "decision N", or "§N.N" references (repo-wide grep); ruff and ty clean; full test suite on the heaviest-touched modules (network, sample, rollout, migration, config, train) passes. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> |
||
|
|
93b19911f8 |
v0.3.0 step 6: sample.py/rollout.py AR generation + class->PDG decode
CI / Format (ruff format) (push) Successful in 36s
CI / Lint (ruff check) (push) Successful in 38s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Failing after 45s
CI / Lint (ruff check) (pull_request) Successful in 41s
CI / Tests (push) Has been skipped
CI / Format (ruff format) (pull_request) Successful in 35s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Failing after 37s
CI / Tests (pull_request) Has been skipped
- giant/sample.py: fix every sampler's call convention against
Stage1Model/Stage2OneShot's actual forward signatures (was still
calling model(x, t, cond_cont, cond_cat) positionally); add
sample_secondaries_ar (free-running AR loop, unsnapped history feature)
and sample_stage1/sample_stage2/resolve_n_sec dispatch helpers that read
each stage's generator_kind/decoder off the model instance itself.
- giant/particles.py: decode_topn_class (argmax + other_policy) and
decode_embedding_nearest (L1-snap + distance) turn a secondary's
"onehot"/"embedding" type prediction into a concrete PDG.
- giant/rollout.py: decode_secondary_identity routes all three
particle_type.target values to real mass/charge; per-stage generator
dispatch (drops the single shared `mode` string, adds ddpm support);
L1DistCollector accumulates the §11.3 embedding-distance diagnostic.
- giant/cli.py: drop the onehot/embedding-target rejection gate (narrowed
to the still-unimplemented conditioning.particle/material.type=onehot
axis); fix the dead model_cfg.get("mode") bug in predict/rollout.
- giant/analysis/: new type_embedding_l1_distance PlotSpec, wired through
the rollout YAML sidecar (no live-model call needed, unlike
router_gating -- the histogram is already pre-aggregated at rollout
time).
- Un-xfail every test that was blocked on this step (test_rollout.py,
test_flow.py, test_wgan.py, test_phase2.py, test_router.py,
test_validate.py); add test_sample.py, test_type_embedding_distance.py.
Known follow-up: giant/validate.py still unpacks the training val-batch
as a stale 6-tuple and doesn't use the new per-stage dispatch, so
marginal validation during training degrades gracefully with a warning
rather than working -- not in this step's scope.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
|
||
|
|
9112e845e0 |
v0.3.0 step 3: per-stage train.py trainers + pipeline.py/cli.py rewrite
CI / Lint (ruff check) (push) Successful in 27s
CI / Format (ruff format) (push) Successful in 28s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 23s
CI / Tests (push) Successful in 1m1s
CI / Lint (ruff check) (pull_request) Successful in 26s
CI / Format (ruff format) (pull_request) Successful in 27s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 27s
CI / Tests (pull_request) Successful in 1m2s
Replaces train.py's single global training loop with a StageTrainer hierarchy (FlowDDPMStageTrainer, WGANStageTrainer) — one per active stage, each owning its own optimizer/LR schedule/EMA and reading only the shared batch tuple (stage 2 always teacher-forces on the ground-truth x1_s1, so stages never need each other's output at train time). Supports every stage1/stage2 generator combination, including the design doc's headline mixed case (stage1=flow + stage2=wgan) and its reverse, plus stage1-only/stage2-only ablation runs, routed+gumbel stages, and checkpoint save/resume. metrics.csv/wandb logging are stage-prefixed. validate_marginals calls are guarded with a one-time warning and a Wasserstein-magnitude fallback for wgan best-checkpoint selection, since giant/sample.py still assumes stage1 always owns n_sec_head (decision 1 moved it to stage 2 by default) — deferred to design doc step 6, not silently papered over. pipeline.py's run_setup_stage/run_train_job now read the new nested config directly; the dangling resolve_expert_dims call and the --mode wgan --router rejection are both gone (routed WGAN works). cli.py's train/new-run build correctly-shaped config overrides (architecture flags -> stage1_model only per the approved decision; --mode/--n-critic/--gp-weight/--critic-lr broadcast to both stages, matching migrate_config's own precedent and avoiding a regression on the common --mode case); predict/rollout's dangling build_models tuple-unpack is fixed; new-run now tags config_version, fixing a bug where a re-loaded v0.3 config.toml would have been silently corrupted by migrate_config mistaking it for v0.2. config.py's validate_config rejects mixed particle/material conditioning types for now (ConditionEncoder supports it, the data pipeline in giant/data/transforms.py doesn't yet). analysis/render.py and router_gating.py handle both the new nested model_config shape and legacy flat checkpoints. scripts/warm_setup_cache.py updated for run_setup_stage's new signature. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> |
||
|
|
9ce55e5013 |
v0.3.0 step 2: network.py refactor to composable stage models
Decomposes the ten permutation classes in giant/model/network.py into the reusable parts from docs/v0.3.0-design.md §5: ConditionEncoder (now independently configurable per particle/material axis), ContextAdapter, Trunk/MonolithicTrunk/RoutedTrunk/ExpertTrunk, and the stage classes Stage1Model/Stage2OneShot/CriticModel (Stage2Autoregressive stubbed, raises NotImplementedError until step 4/5). build_models/build_critics now return a dict keyed by stage and accept the new nested config shape, with routed WGAN reachable for the first time (the old --mode wgan --router rejection is gone) and stage2_model.router.tie_to_stage1 sharing a literal Router instance. A v0.2 checkpoint's flat model_config auto-migrates via _migrate_legacy_model_config + migrate_legacy_state_dict, preserving the n_sec_head's attachment to Stage1Model (legacy_owner="stage1", design doc §4.1). tests/test_migration_v02_v03.py proves this bit-identical against a frozen v0.2 snapshot (tests/legacy/network_v02_snapshot.py) for both flow and wgan, both conditioning modes. scripts/check_migration_v02_v03.py is the real-checkpoint counterpart for a portal machine with /ceph access. giant/model/schedule.py's flow-matching/DDPM loss helpers are updated to the new model-call convention (t as a keyword). giant/sample.py, giant/rollout.py, and giant/validate.py are not yet updated (deferred to design doc step 6) — their exercising tests are marked xfail with that reasoning rather than silently broken. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> |
||
|
|
5b63dfd588 |
Fix conditioning="physical" so it can actually generalize past training vocab
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> |
||
|
|
b51eafcfa5 |
Add opt-in straight-through Gumbel-softmax combine weights to MoE router
CI / Lint (ruff check) (push) Successful in 28s
CI / Format (ruff format) (push) Successful in 28s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 21s
CI / Lint (ruff check) (pull_request) Successful in 25s
CI / Format (ruff format) (pull_request) Successful in 30s
CI / Tests (push) Successful in 1m37s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 31s
CI / Tests (pull_request) Successful in 55s
Trains the routed trunk's forward combination as a hard one-hot sample (matching eval-time top-1 dispatch exactly) while keeping a smooth gradient on the backward pass, targeting the train/eval mismatch identified as a likely contributor to experts overlapping instead of partitioning in the first energy-router rollout benchmark. Off by default (model.router.gumbel); existing routed configs/checkpoints are unaffected. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> |
||
|
|
da5f54ea1c |
Add learnable per-expert width and shared temperature to EnergyRouter
CI / Lint (ruff check) (push) Successful in 28s
CI / Format (ruff format) (push) Successful in 29s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 27s
CI / Tests (push) Successful in 1m1s
CI / Lint (ruff check) (pull_request) Successful in 30s
CI / Format (ruff format) (pull_request) Successful in 27s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 27s
CI / Tests (pull_request) Successful in 59s
EnergyRouter's gate sharpness was a single fixed temperature shared by every expert, with no way for an expert to independently learn how much of the energy axis it covers. Adds two mutually exclusive, default-off modes: learn_width (per-expert learnable width) and learn_temperature (single learnable shared scalar), both bounded via a sigmoid interpolation warm-started to reproduce today's fixed-temperature gate exactly at init, to compare against each other without risking the unbounded-width collapse failure mode. Also promotes gate_stats's entropy into a generic, optional Router.entropy_loss (lambda_entropy) as a secondary guard against all experts' widths co-inflating together. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> |
||
|
|
ee29b9a303 |
router: seed EnergyRouter centers from data quantiles instead of a fixed linspace
CI / Lint (ruff check) (push) Successful in 1m1s
CI / Format (ruff format) (push) Successful in 1m6s
CI / Type check (ty) (push) Successful in 59s
CI / Tests (push) Successful in 1m45s
CI / Lint (ruff check) (pull_request) Successful in 1m4s
CI / Format (ruff format) (pull_request) Successful in 1m5s
CI / Type check (ty) (pull_request) Successful in 1m4s
CI / Tests (pull_request) Successful in 1m55s
CI / Bump version, build & publish wheel (push) Has been skipped
CI / Bump version, build & publish wheel (pull_request) Has been skipped
The 2026-07-22 rollout benchmark's router_gating diagnostic showed the 10-expert EnergyRouter's default linspace(-2, 2, n_experts) init assumes a roughly uniform z-normalized energy distribution, leaving experts heavily overlapping instead of partitioning the range. Add an optional centers_init kwarg (backward compatible, defaults to the old linspace) and have giant train estimate it from a reservoir sample of the real energy column, collected during the existing normalizer-fitting pass. |
||
|
|
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> |
||
|
|
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>
|
||
|
|
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> |
||
|
|
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. |
||
|
|
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> |