# GIANT reference baseline (v0.3 schema).
#
# The fixed comparison point every future architecture variant is measured
# against. Chosen so that each experimental axis the roadmap cares about
# (routed trunk, WGAN generators, attention history, shared conditioning,
# embedding/onehot conditioning) is a *single* edit away from this file.
#
# Rationale for the choices below, from the runs already on record
# (analysis_runs/ + the `giant` W&B project):
#
#   * flow, not wgan, for both stages. Ranking the five existing rollouts by
#     mean Jensen-Shannon divergence against the Geant4 reference, the plain
#     non-routed flow model wins (0.172) over the routed flow runs
#     (0.197/0.200) and both WGAN runs (0.218/0.234) — and it beats them by
#     ~7x on per-event total deposited energy and by 3-10x on every
#     per-PDG marginal. WGAN stays a variant, not the reference.
#
#   * no router. The routed runs are not better, and soft-mixing 10 small
#     experts costs ~10x per-pass throughput at train time (29k samples/s vs
#     the WGAN runs' 52-116k), which is what made those runs take ~110 h for
#     30 epochs.
#
#   * hidden_dim 512 / 6 blocks per stage. The best-scoring rollout so far
#     was hidden_dim 1024, but at 4x the trunk FLOPs of 512. 512/6 sits in
#     the same weight class as the variants it will be compared against and
#     leaves headroom to train it properly rather than cheaply.
#
#   * dropout 0.0. Training set is ~5e8 steps against <1e7 parameters;
#     capacity overfitting is not the binding constraint, and every recent
#     run used 0.0.
#
# Known weak spots this baseline is expected to *exhibit* (they are the
# reason for the comparisons, not a reason to retune this file): every model
# on record under-produces steps per event by ~2x (rollout ~7e4 vs Geant4
# ~1.4e5) and secondaries per event by 2-3.5x (~2-3e4 vs 7.2e4), and n_sec
# head accuracy sits at 0.863-0.867 regardless of size or objective.

[meta]
# REQUIRED. Without it config.migrate_config reads this file as v0.2 and
# rewrites it from V02_FIXED_FACTS — silently forcing decoder = "one_shot",
# particle_type.target = "physical" and the v0.2 default sizes, while still
# passing validate_config.
config_version = 3

[conditioning]
# Physical-property MLPs rather than learned vocab embeddings: computable for
# any PDG code / material, which is what the held-out-species and
# held-out-material generalization comparisons need.
out_dim      = 128
share_stages = false

# n_layers = 2 rather than the v0.3 default of 1: v0.2's conditioning MLP was
# always 2 deep (see _migration.V02_FIXED_FACTS), so this keeps the encoder
# identical to the architecture that produced the results cited above.
[conditioning.particle]
type     = "physical"
emb_dim  = 16
n_layers = 2

[conditioning.material]
type     = "physical"
emb_dim  = 16
n_layers = 2

[stage1_model]
generator    = "flow"
hidden_dim   = 512
n_res_blocks = 6
dropout      = 0.0

[stage2_model]
# The v0.3 pivot: autoregressive in descending-energy order with a
# categorical species target, which is the agreed response to the 2026-08-03
# secondary-species failure. Flow (not the schema default wgan) so the
# baseline varies only the decoder relative to the best v0.2 result.
#
# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated:
# sample.sample_secondaries_ar loops `for k in range(k_max)` unconditionally
# — all 15 slots regardless of predicted n_sec — so a flow AR token costs
# k_max * steps = 150 stage-2 calls per physics step. That makes this block
# the dominant cost on both sides:
#   training   flow AR 29.5k samp/s  vs flow one-shot 190.7k samp/s  (6.5x)
#   inference  flow AR  8.5k step/s  vs flow one-shot  68.7k step/s  (8.1x)
# Accepted deliberately: one-shot is the configuration whose secondary
# species distribution failed, and that failure is what v0.3 exists to fix.
decoder      = "autoregressive"
generator    = "flow"
hidden_dim   = 512
n_res_blocks = 6
dropout      = 0.0
k_max        = 15

[stage2_model.autoregressive]
history         = "markov"
teacher_forcing = "always"

[stage2_model.particle_type]
target = "onehot"
# Decoupled from conditioning.particle.emb_dim (gitea #29). 32 classes + the
# "other" bucket keeps essentially all real secondary species out of "other"
# without making the head expensive.
n_classes    = 32
other_policy = "sample"

[train]
epochs        = 50
# Sized for ONE NVIDIA L40S on deepthought2 (46068 MiB; the box has two, and
# CLAUDE.md's shared-machine rule allows a single GPU). From a measured
# linear fit of this exact config's training step on the local RTX 4070:
#   peak reserved MiB = 0.9736 * batch_size + 115
# so 36864 reserves ~36.0 GiB, i.e. 78% of the card, leaving ~10 GiB of
# headroom for fragmentation and the CUDA context. Throughput is already
# flat above bs~4096 on the 4070, so this is chosen for occupancy on the
# larger card, not for step efficiency — and it sits next to the 43008/32768
# of the runs lr = 3e-4 was proven at.
batch_size    = 36864
lr            = 3e-4
warmup_epochs = 3
weight_decay  = 0.01
ema_decay     = 0.9999
val_fraction  = 0.1
num_workers   = 4
seed          = 0
# The marginal/KL pass is expensive (~5000 s on top of an epoch), so keep it
# to every 10th epoch; the cheap per-epoch val loss still runs every epoch.
validate_every = 10
validate_steps = 10
wandb          = true
wandb_project  = "giant"