21 Commits

Author	SHA1	Message	Date
lars	c4b12b5e7a	Pass ConditioningAxisConfig/ParticleTypeConfig themselves instead of raw dicts (gitea #38) ... 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>	2026-08-14 15:03:55 +02:00
lars	1a3c907571	Give Stage1Model/Stage2OneShot/Stage2Autoregressive a shared StageModel base (gitea #39) ... Stage1Model, Stage2OneShot and Stage2Autoregressive each independently implemented ~90 near-identical lines of __init__ scaffolding: build-or-share cond_enc, particle_type_cfg normalisation, objective -> time_emb -> merged_cond_dim -> build_trunk, and the n_sec_head/type_head classifier heads (plus their identical RuntimeError guards). Now unblocked by #33 (trunk registry), #34 (block-conditioning registry) and #36 (build_mlp_head), which settled what belongs in the shared base. Adds StageModel(nn.Module) owning all of that: __init__ builds/shares cond_enc and normalises particle_type_cfg; _build_trunk_and_heads, called by each subclass after it sets up its own conditioning-assembly modules (cond_enc alone for Stage1Model, a context-fusion path for the two Stage2 classes), builds the objective/time embedding/trunk and the n_sec_head/type_head guarded by the shared _require_n_sec_head/ _require_type_head (Stage1Model overrides the n_sec guard since its message points at stage 2, not stage 1). Public __init__ signatures, attribute names, and forward/predict_* behaviour are unchanged. Verified with a pre/post state_dict-key-set diff against the pre-refactor classes (bit-identical) before writing this commit, plus new parametrized tests pinning each class's state_dict key set and the generator -> time_emb contract the base now owns. tests/test_migration_ v02_v03.py's existing bit-identical old-vs-new forward comparison and the rest of tests/test_network.py's per-class coverage pass unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-14 14:37:58 +02:00
lars	4692cee699	Give the cond_cat/cond_cont column layout one owner (gitea #37) ... The conditioning arrays' column order was written down three times — twice in giant/data/transforms.py (build_cond_features and build_features each built cond_cont and cond_cat from scratch) and again in giant/model/encoders.py (cat_col_layout, plus hand-written COND_DIM_BASE + PARTICLE_PHYS_DIM slicing in ConditionEncoder). The three were held in sync only by parallel comments, so a wrong column order produced silently mis-indexed features rather than an exception. The drift had already happened, twice, both times in build_features: - 5b63dfd added per-axis vocab-lookup strictness (an out-of-vocab pdg/material must not KeyError under "physical"/"onehot", where the index is never read) to build_cond_features only. - _cond_normalizer_transform's legacy-normalizer padding, which keeps a pre-physical-conditioning 8-wide cond normalizer loadable, was likewise only wired into build_cond_features — so `giant predict` on such a checkpoint died with a broadcast error. New giant/cond_layout.py holds a frozen CondLayout built from the (particle, material) mode pair, exposing named cond_cont slices (base/particle_phys/material_phys) and cond_cat columns (PDG_COL/MAT_COL/particle_topn_col/material_topn_col/cat_dim). Both builders now share one _build_cond_arrays, ConditionEncoder reads its slices off the same object, and PdgRouter/ProcessRouter use the named dense-vocab columns instead of literal 0/1. CondLayout also absorbs the two duplicated axis-type validations, keeping their message text verbatim. Decisions taken while planning: - Scope is CondLayout only. The issue's second half — a CONDITIONING_AXIS_REGISTRY registering (feature_columns, encoder_module) as a pair — is deferred: it would force ConditioningConfig's fixed particle/material fields into a dynamic axis map and ripple through pipeline.py, checkpoint_io.py and rollout.py, i.e. a config-schema break with no consumer yet. - The two divergences above are unified onto build_cond_features' behaviour rather than preserved as parameters, so the new single source of truth doesn't carry the old split forward. Each gets a regression test that fails before this commit. - cat_col_layout is replaced outright (deleted, dropped from network.py's __all__, its four tests rewritten against CondLayout) rather than kept as a wrapper — two spellings of the same fact is the defect itself. cond_cat's width is now the layout's call rather than "did the caller pass a map", so an "onehot" axis without its top-N map raises instead of yielding a narrower array that ConditionEncoder would index out of bounds. pipeline.py's normalizer-fitting pass reads only cond_cont but had to be handed the maps to satisfy that. No parameter, buffer or state_dict change; existing checkpoints load unchanged, and the protected migration surfaces are untouched. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-14 14:16:20 +02:00
lars	593c5f4d34	Deduplicate n_sec_head/type_head MLPs into build_mlp_head (gitea #36) ... The same two-layer classifier head (Linear(cond_out_dim, hidden_dim // 2) -> SiLU -> Linear(hidden_dim // 2, out_dim)) was hand-rolled five times in giant/model/models.py: Stage1Model.n_sec_head, Stage2OneShot.n_sec_head/ .type_head, and Stage2Autoregressive.n_sec_head/.type_head. The `// 2` ratio and fixed 2-layer depth were undocumented magic numbers, and both n_sec accuracy and secondary-species accuracy are known weak spots that were untunable independently of the trunk they hang off. Adds `build_mlp_head(in_dim, out_dim, hidden, depth, act)` to giant/model/layers.py (depth=1 is a bare Linear; depth>=2 matches the old hardcoded shape exactly), and a new `HeadConfig` (hidden_ratio, depth) dataclass in giant/config.py, wired in as `stage1_model.heads.n_sec` and `stage2_model.heads.{n_sec,type}` — split per head type (not one shared block per stage) since n_sec and species prediction are called out as separate weak spots that may want independent capacity. Defaults (hidden_ratio=0.5, depth=2) reproduce the old hardcoded architecture bit-for-bit, so every existing config.toml and migrated v0.2 checkpoint is unaffected; no changes were needed to migrate_config or the legacy migration surfaces. No new CLI flags, matching how other nested sub-config (router.*, trunk.*) is set via config.toml rather than per-field flags. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-14 10:57:47 +02:00
lars	f301fd98d2	Make HistoryEncoder a pluggable registry, like Router/Objective (gitea #35) ... Stage2Autoregressive.init_history_cache and .history_step both isinstance-checked self.history_encoder against AttentionHistory to decide whether to use its real incremental-cache methods or a no-op fallback, so a third history type couldn't be added without editing Stage2Autoregressive itself. The two-value "markov"/"attention" enum was also independently hardcoded in three places (Stage2Autoregressive's own validation, config.py's validate_config, and AutoregressiveConfig.from_dict's default). Mirrors the Router (giant/model/routers.py) and Objective (giant/model/objectives.py, gitea #32) pattern: HistoryEncoder now declares working O(1) init_cache/step defaults (init_cache -> None, step -> one forward() call), so every registered history type satisfies the incremental interface without opting in; AttentionHistory overrides both with its real KV-cache versions since its forward() needs the full prefix. Added HISTORY_REGISTRY/register_history/build_history, registered "markov" and "attention", and deleted both isinstance checks in models.py. Per user decision during planning, config.py's validate_config now imports HISTORY_REGISTRY and checks membership dynamically instead of keeping its own hardcoded tuple, making the registry the single source of truth end to end (verified no import cycle: config.py had no prior dependency on giant.model, and giant.model.history has none on giant.config). No config-schema change and no checkpoint impact — this is a pure internal-interface refactor. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-14 10:35:46 +02:00
lars	32aa5a5f92	Decouple secondary-species vocabulary from conditioning.particle.emb_dim (gitea #29) ... conditioning.particle.emb_dim and stage2_model.particle_type.target="onehot"'s class count were silently the same number everywhere (pipeline.py's PDG top-N map build, Stage2OneShot/Stage2Autoregressive's type head, StageSpec's training loss width, the checkpoint's shared pdg_topn_map), fixing the secondary-species vocabulary at whatever width the unrelated physical-conditioning MLP happened to use — the exact vocabulary the v0.3.0 pivot exists to fix. Adds stage2_model.particle_type.n_classes (default 0 = inherit conditioning.particle.emb_dim, preserving today's behavior and every existing checkpoint) and a single resolve_type_n_classes helper used everywhere the coupling used to be implicit. Splits the checkpoint's shared pdg_topn_map into a conditioning-only pdg_topn_map and a new sec_type_topn_map, built independently through the existing (axis, n_classes)-keyed setup cache (no extra scan when they still resolve to the same N) and threaded through giant predict/giant rollout's decode path. A checkpoint with no sec_type_topn_map key (pre-#29) falls back to reusing pdg_topn_map, reproducing the old shared behavior exactly. Decided with the user during planning: commit directly on this branch; represent the split as an additive sec_type_topn_map checkpoint key rather than conditionally reusing pdg_topn_map; build the two top-N maps independently rather than the issue's proposed build-at-max-and-slice, since the setup cache already avoids redundant scans across runs. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-13 16:11:14 +02:00
lars	da717971b6	Honour wgan.critic_hidden_dim/critic_n_res_blocks in build_critics (gitea #28) ... build_critics always sized a WGAN critic off the generator's own hidden_dim/n_res_blocks, silently discarding the documented 0=inherit sentinel on stage{1,2}_model.wgan.critic_hidden_dim/critic_n_res_blocks (the same convention critic_lr already honoured). Now both keys are read with the 0 -> inherit fallback, and stage-scoped-only CLI flags (--stage{1,2}-critic-hidden-dim/--stage{1,2}-critic-n-res-blocks) are added -- no shared alias, since critic sizing is an architectural per-stage knob like --hidden-dim/--n-res-blocks, not a shared training hyperparameter like --n-critic/--gp-weight/--noise-dim/--critic-lr. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-13 15:36:26 +02:00
lars	9bf5874308	Make config dataclasses the single source of truth for DEFAULT_CONFIG ... DEFAULT_CONFIG and build_models/build_critics/StageSpec.from_config's inline .get(key, default) fallbacks had already drifted: two keys (stage2_model.decoder, stage2_model.particle_type.target) resolved differently depending on whether a config dict came from merge_cli_overrides (fully populated, correct) or was hand-built and partial (fell back to stale v0.2-shaped literals). Introduce frozen dataclasses (GiantConfig and its nested blocks) in giant/config.py as the actual single declaration of every default; DEFAULT_CONFIG is now generated from them instead of hand-maintained, and build_models, build_critics, and StageSpec.from_config consume the dataclasses instead of duplicating literal fallbacks, so this class of drift can't recur. Router/n_sec sub-blocks keep an `extra` catch-all for their genuinely dynamic keys (composed-router axes, runtime-seeded centers_init, legacy_owner). Fixing the fallback surfaced the same latent bug in two existing partial-config callers that had been silently depending on it: a test fixture in test_train.py and scripts/warm_setup_cache.py's minimal cfg (now merged against DEFAULT_CONFIG instead of hand-rolled, closing the gap for good). See issues.md Issue 1. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-12 14:35:14 +02:00
lars	55332db67a	Bump ruff line-length to 120 and reformat ... Rejoins lines that only wrapped because they exceeded the old 88-char limit; ruff check and the full test suite (725 passed) are unaffected.	2026-08-12 13:33:09 +02:00
lars	878e9ddca3	Delete docs/v0.3.0-design.md and strip all references to it ... 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>	2026-08-10 11:19:02 +02:00
lars	da7cde3ef9	v0.3.0 post-implementation audit: resolve all 9 tracked discrepancies ... 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>	2026-08-07 16:12:58 +02:00
lars	200c6d243b	v0.3.0 step 7: AttentionHistory (KV-cached) + scheduled/never teacher forcing ... AttentionHistory (giant/model/network.py) adds causal self-attention over the emitted-secondary prefix as the alternative to MarkovHistory, with a parallel forward() for training and an init_cache()/step() KV-cache path for sample.py's per-slot AR inference loop, wired into Stage2Autoregressive via history="attention". giant/train.py adds _stage2_tf_prob and _assemble_stage2_ar_inputs_scheduled, mixing ground-truth history with a detached sample_secondaries_ar self-sample per slot so teacher_forcing="scheduled"/"never" close the train/inference gap teacher_forcing="always" always avoided; wired into both stage-2 AR trainers. config.py's validate_config no longer rejects these two previously unimplemented schema values. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-07 13:15:56 +02:00
lars	c9d255b1c5	v0.3.0 step 5: Stage2Autoregressive (history=markov) + §11.4 grad instrumentation ... Replaces the Stage2Autoregressive stub with a real per-token secondary decoder: MarkovHistory summarizes the previous secondary, remaining-energy fraction and slot index round out the per-token conditioning, and the existing Trunk/MonolithicTrunk/RoutedTrunk machinery is reused unchanged by batching all K_MAX tokens together under teacher forcing (one parallel pass, no new trunk code). build_models/build_critics wire it in; the WGAN critic stays whole-sequence, so build_critics needs no AR-specific path. train.py's FlowDDPMStageTrainer/WGANStageTrainer gain a decoder branch, sharing optimizer/EMA/checkpoint machinery with the one-shot path. _assemble_stage2_real is now defined in terms of the new unflattened _assemble_stage2_ar_target helper, removing a near-duplicate branch. Also lands the §11.4 differentiability validation-obligation instrumentation (trunk-gradient norm from the particle-type slice vs. the continuous slices, for generator=wgan + particle_type.target=onehot) via backward hooks in _relax_onehot_type_slice, decoder-agnostic and surfaced as two new metrics.csv columns. This also fixes the standing regression where any config not explicitly overriding decoder="one_shot" crashed at build_models, since stage2_model.decoder defaults to "autoregressive" — confirmed by removing tests/test_pipeline.py's now-stale override so the default config runs end-to-end against real synthetic data. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-07 09:36:49 +02:00
lars	4fc15ecdfc	v0.3.0 step 4: type map + particle_type.target = "onehot"/"embedding" ... 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>	2026-08-06 15:43:48 +02:00
lars	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>	2026-08-06 10:55:29 +02:00
lars	e6e0eb22bf	Implement Phase 2: secondary particle prediction ... Two-stage factorisation: Stage 1 predicts 9D primary kinematics + n_sec classification head (COND_DIM reduced to 8, dropping n_sec/e_sec inputs); Stage 2 (SecondaryDecoder) generates K_MAX=15 secondary slots via masked flow matching over (stick_logit, local_dir, type_emb) conditioned on Stage 1 output. Joint training with combined loss L_s1 + λ_nsec*L_nsec + λ_s2*L_s2. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>	2026-06-29 11:34:31 +02:00
lars	8475199609	Encode edep/secondary/post energy as a conservation-constrained simplex ... 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.	2026-06-25 16:01:13 +02:00
lars	8cebc4809d	Apply ruff format and document lint/type tooling in CLAUDE.md ... First repo-wide ruff format pass, plus a note in CLAUDE.md to run ruff and ty periodically. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>	2026-06-18 17:40:27 +02:00
lars	53fd2e4405	Add ruff and ty as dev dependencies, fix lint/type findings ... 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>	2026-06-18 17:38:40 +02:00
lars	72bd65ff9f	Add post_pos as a model target via travel_dir decomposition ... 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>	2026-06-18 10:36:55 +02:00
lars	9277d79dff	Implement Phase 1: full data pipeline, model, training, and config support ... - Data pipeline: loader (parquet→numpy), transforms (log, local-frame Rodrigues rotation, Normalizer), StepsDataset with event-ID-based split - Model: SinusoidalEmbedding, ConditionEncoder, ResBlock, DenoisingMLP - Schedule: cosine DDPM and conditional flow matching loss (Lipman 2022) - Samplers: flow (Euler ODE), DDPM ancestral, DDIM deterministic - Training loop: AdamW + cosine LR, grad clipping, best-val checkpoint - Validation: per-dimension marginal summary (normalised space) - CLI: TOML config support with CLI-overrides; hyperparam-encoded output directory; config.toml with git hash saved into each run's checkpoint dir - 21 unit tests covering transforms, network, flow/DDPM losses, dataset splits Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>	2026-06-17 10:48:03 +02:00