Stage 2's autoregressive decoder still predicted multiplicity the v0.2 way:
a one-shot n_sec_head classifier over conditioning alone, run before any
secondary token existed, with the AR loop then always executing k_max slots
and discarding the tail. This adds a real per-slot EOS mechanism instead:
- Stage2Autoregressive gains a stop_head (build_stop_head=True) that predicts
P(n_sec == k | prefix) at each slot, mutually exclusive with n_sec_head
(n_sec.mode = "stop_token" builds no n_sec_head at all).
- sample_secondaries_ar accepts n_sec_pred=None to drive generation off the
stop head instead of a pre-resolved count: each row stops the first slot
its stop logit fires (stage2_model.n_sec.stop_sampling = "greedy" — the
default, threshold at 0 — or "sample", a Bernoulli draw), and the whole
batch loop breaks once every row has stopped, so cost scales with the
realized n_sec instead of a fixed k_max. Passing n_sec_pred explicitly
(the scheduled-sampling self-sample path) is unchanged.
- resolve_n_sec returns None for a stop-token decoder instead of raising;
rollout.py/cli.py/validate.py now derive the realized count from
sample_stage2's returned sec_valid (sec_valid.sum(-1)) after sampling,
rather than resolving it up front — a no-op reordering under every other
n_sec.mode, where sec_valid was already built from n_sec_pred.
- Training: _stop_target_and_mask (giant/training/stage2_inputs.py) builds
the per-slot target/mask (one slot wider than the existing token-content
sec_mask, since the stop slot itself needs supervision) and
StageTrainer._stop_loss trains it with masked BCE, gated on stop_head
exactly like _n_sec_loss gates on n_sec_head. Wired into both the
flow/ddpm trainer and the WGAN trainer (whose skip_g_step now also checks
stop_head), weighted by the existing stage2_model.n_sec.lambda — the stop
head replaces n_sec_head under this mode, so no new weight key.
- validate_config now accepts stop_token (requires decoder="autoregressive"
and n_sec.owner="stage2") instead of always rejecting it.
Decisions made during planning: stop_sampling defaults to "greedy" for
deterministic rollouts; the stop head reuses stage2_model.heads.n_sec's
HeadConfig shape and stage2_model.n_sec.lambda's weight rather than adding
new config keys, since the two heads never coexist.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
build_models/build_critics parsed model_config into frozen dataclasses
(ConditioningConfig, Stage2ModelConfig, ...) but then threw the parsed
sub-objects away and passed the original raw dicts (conditioning["particle"],
s2_spec.particle_type.to_dict()) down into ConditionEncoder/StageModel/etc,
which re-read them with their own hardcoded .get(key, default) fallbacks —
each an independent copy of a fact the dataclass already stated once. Worst
instance: giant/training/trainers.py:236 converted an already-parsed
ParticleTypeConfig back into a dict for no reason.
Threads ConditioningAxisConfig (particle_cfg/material_cfg) and
ParticleTypeConfig (particle_type_cfg) as the actual dataclass instances
through every signature that used to type them dict: ConditionEncoder,
StageModel/CriticModel, resolve_type_n_classes/stage2_type_dim/
stage2_trunk_sec_dim, giant/model/builders.py, giant/sample.py,
giant/training/stage2_inputs.py, giant/training/trainers.py (StageSpec/
StageTrainer), giant/pipeline.py, giant/rollout.py, giant/validate.py — so ty
now catches a misspelled field instead of it silently falling back. No
config-schema change: config.toml/checkpoint model_config keep the same
nested-dict shape; only what happens after the existing X.from_dict(...)
parse changes.
User-confirmed scope decision: both axes (particle_cfg/material_cfg and
particle_type_cfg), not just the more heavily-duplicated particle_type_cfg
axis, and not stopping at the two most literal parse-then-discard round
trips — matching the issue's own proposal.
Preserved-default decision: StageModel's particle_type_cfg=None sentinel
(hit only by direct/test construction — build_models always passes an
explicit particle_type) still resolves to ParticleTypeConfig(target=
"physical"), not ParticleTypeConfig()'s own target="onehot" config-file
default — switching it would have silently grown an unused, gradient-less
type_head on every test that constructs Stage2OneShot/Stage2Autoregressive
without particle_type_cfg=, breaking their "every param has a grad" checks.
New tests in tests/test_network.py: ConditionEncoder/StageModel store the
exact ConditioningAxisConfig/ParticleTypeConfig instance passed in (identity,
not just equality) — no internal dict round-trip — and build_models's output
carries real dataclass instances end to end, not the plain dicts it produced
before this fix.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
build_features (transforms.py) now returns StepFeatures and
StreamingStepsDataset (dataset.py) now yields StepBatch, both NamedTuples
with the same field order as the tuples they replace, so ty can catch a
dropped/added field at every consuming call site instead of a silent
positional-tuple mismatch. Converted the unreadable throwaway-heavy unpacks
in cli.py, pipeline.py, validate.py, and dataset.py to named attribute
access; gave the WGAN path's derived 5-element batch its own
_Stage2RealFakeBatch NamedTuple; updated the two test batch-construction
helpers to build real StepBatchs.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
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>
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>
The per-dim print loop lacked the empty-array guard already used for
the KL computation right above it and the sec-slot loop further down,
so an all-zero-secondaries validation batch (e.g. early/unstable
training) triggered numpy RuntimeWarnings from .mean()/.std() on
empty arrays.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
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>
- 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>
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>
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.
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>
validate_marginals and collect_samples could already vary flow ODE
steps for inference (giant predict --steps), but training-time
marginal validation and DDIM evaluation were stuck at hardcoded
defaults. Add a validate_steps config/CLI option and forward steps to
sample_ddim consistently.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
validate_marginals now estimates a per-dimension KL(real || generated) via
a shared histogram, alongside the existing mean/std comparison, so
distribution-shape drift shows up even when the first two moments match.
Wire it into giant/train.py: every validate_every epochs (default 10, 0
disables), the training loop runs validate_marginals against val_loader and
prints the table. validate_every flows through DEFAULT_CONFIG/config.toml
and is exposed as --validate-every on both giant train and scripts/train.py.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
cli.py and scripts/train.py duplicated ~140 lines of training setup and had
drifted (scripts/train.py forgot to save model_config, breaking predict on
those checkpoints). Extract shared logic into giant/constants.py (X_DIM,
target names), giant/config.py (device/git/TOML/seeding helpers, run
metadata), and giant/pipeline.py (the actual training-job orchestration),
so both entry points become thin CLI wrappers around the same code path.
Also adds --seed/--resume support (checkpoints now carry optimizer/scheduler
state, epoch, and best_val_loss), a richer [meta] section in the saved
config.toml (git hash, seed, versions, timestamp, invocation, dataset
stats), and a metrics.csv (train/val loss, lr, epoch time) written every
epoch and append-safe across resumes.
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