Compare commits
5 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 gitea-actions
|
c613588a70 | ||
|
gitea-actions
|
56642ebd2c | ||
 lars
|
9a03f4552a | ||
|
lars
|
5c576fa8f3 | ||
|
lars
|
bf3271f09e |
+1
-1
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.3.17"
|
||||
current_version = "0.3.18"
|
||||
parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
|
||||
serialize = ["{major}.{minor}.{patch}"]
|
||||
search = "{current_version}"
|
||||
|
||||
@@ -173,6 +173,52 @@ jobs:
|
||||
git push origin "refs/tags/$TAG"
|
||||
fi
|
||||
|
||||
update-badges:
|
||||
name: Update README badges (version, test count)
|
||||
needs: [ruff-check, ruff-format, type-check, test, bump-version]
|
||||
if: github.ref == 'refs/heads/master' && github.event_name == 'push'
|
||||
runs-on: ubuntu-latest
|
||||
container:
|
||||
image: docker.gitea.com/runner-images:ubuntu-latest
|
||||
volumes:
|
||||
- /srv/act-runner-cache/uv:/uv-cache
|
||||
steps:
|
||||
# ref: master (not the triggering SHA) so this picks up whatever
|
||||
# bump-version just pushed, rather than badging the pre-bump commit.
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
token: ${{ secrets.CI_TOKEN }}
|
||||
ref: master
|
||||
- uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
enable-cache: false
|
||||
- run: |
|
||||
echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
|
||||
echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
|
||||
- run: uv sync --extra cpu --extra dev
|
||||
- name: Compute version and test count
|
||||
id: stats
|
||||
run: |
|
||||
VERSION=$(uv version --short)
|
||||
TEST_COUNT=$(uv run pytest --collect-only -q 2>/dev/null | grep -oE '^[0-9]+ tests? collected' | grep -oE '^[0-9]+')
|
||||
echo "version=$VERSION" >> "$GITHUB_OUTPUT"
|
||||
echo "test_count=$TEST_COUNT" >> "$GITHUB_OUTPUT"
|
||||
- name: Rewrite badge lines in README.md
|
||||
run: |
|
||||
sed -i -E "s|badge/version-[^-]+-informational|badge/version-${{ steps.stats.outputs.version }}-informational|" README.md
|
||||
sed -i -E "s|badge/tests-[0-9]+%20passing-brightgreen|badge/tests-${{ steps.stats.outputs.test_count }}%20passing-brightgreen|" README.md
|
||||
- name: Commit and push if the badges actually changed
|
||||
run: |
|
||||
git config user.name "gitea-actions"
|
||||
git config user.email "actions@git.larsbogner.de"
|
||||
git add README.md
|
||||
if ! git diff --cached --quiet -- README.md; then
|
||||
git commit -m "chore: update README badges (version ${{ steps.stats.outputs.version }}, ${{ steps.stats.outputs.test_count }} tests)"
|
||||
git push origin HEAD:master
|
||||
else
|
||||
echo "Badges already up to date"
|
||||
fi
|
||||
|
||||
sync-version-on-tag:
|
||||
name: Sync project version with tag
|
||||
if: startsWith(github.ref, 'refs/tags/')
|
||||
|
||||
@@ -1,5 +1,13 @@
|
||||
# Changelog
|
||||
|
||||
## [0.3.18] - 2026-09-03
|
||||
|
||||
### Changed
|
||||
|
||||
- Docs: rewrite README, keep version/test badges live via Gitea Actions
|
||||
|
||||
- Perf: compact Stage-2 AR inference loop to active rows only
|
||||
|
||||
## [0.3.17] - 2026-09-02
|
||||
|
||||
### Changed
|
||||
|
||||
@@ -1,203 +1,295 @@
|
||||
# giant
|
||||
<div align="center">
|
||||
|
||||
**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate.
|
||||
# GIANT
|
||||
|
||||
A conditional generative model that replaces the Geant4 step function: given a pre-step particle state it samples a post-step outcome — the primary's continuation plus its secondary particles — and autoregressively rolls that out into full showers. Trained entirely from parquet dumps of the miniCaloSim steps tree; no Geant4 runtime dependency.
|
||||
### **G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate
|
||||
|
||||
A conditional generative model that replaces the Geant4 step function — sample a
|
||||
post-step outcome instead of simulating one, then roll that out into full
|
||||
calorimeter showers.
|
||||
|
||||
[](pyproject.toml)
|
||||
[](pyproject.toml)
|
||||
[](CHANGELOG.md)
|
||||
[](tests/)
|
||||
[](https://git.larsbogner.de/lars/giant/actions)
|
||||
[](#license)
|
||||
|
||||
</div>
|
||||
|
||||
---
|
||||
|
||||
## The idea
|
||||
|
||||
Geant4's step function is the innermost loop of detector simulation — for every
|
||||
particle, at every step, it stochastically samples where the particle goes next,
|
||||
how much energy it deposits, and what secondaries it spawns. GIANT learns that
|
||||
function instead of running it: given a pre-step particle state (position,
|
||||
energy, direction, particle species, material), a two-stage model samples a
|
||||
post-step outcome — including the variable-length list of secondaries — and
|
||||
autoregressively rolls that out into whole showers. It trains entirely from
|
||||
parquet dumps of a Geant4 steps tree; nothing downstream needs a Geant4 runtime.
|
||||
|
||||
Two guarantees are architectural, not learned:
|
||||
|
||||
- **Energy is conserved by construction.** Stage 1 decodes deposit / secondary /
|
||||
post-step energy through a softmax simplex that sums to the pre-step energy
|
||||
exactly; Stage 2's secondaries stick-break that same energy budget.
|
||||
- **No shower leaks across the train/val split.** Steps are split by `event_id`,
|
||||
never by row, so correlated steps from the same shower can't appear on both
|
||||
sides.
|
||||
|
||||
## How a step becomes a shower
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A["pre-step state\nposition · energy · direction\nspecies · material"] --> B["ConditionEncoder\nphysical / embedding / onehot"]
|
||||
B --> C["Stage 1\n9D post-step outcome"]
|
||||
C --> D["Stage 2 (autoregressive)\nsecondaries, descending energy"]
|
||||
D --> E["rollout step"]
|
||||
E -->|"primary continues"| F["GeometryOracle\nposition to material, layer"]
|
||||
E -->|"secondaries pushed"| G["track queue"]
|
||||
F --> A
|
||||
G --> A
|
||||
E -->|"terminated"| H["deposited shower"]
|
||||
```
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
uv sync --extra cpu # install deps (CPU torch; use --extra cuda for GPU)
|
||||
|
||||
giant new-run --hidden-dim 512 --lr 3e-4 # scaffold config.toml + run dir
|
||||
giant model summary --config config.toml # parameter counts + which config keys actually bite
|
||||
giant train path/to/steps.parquet # train (flow + wgan by default)
|
||||
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
|
||||
uv sync --extra cpu # install deps (CPU torch; --extra cuda for GPU)
|
||||
|
||||
giant new-run --hidden-dim 512 --lr 3e-4 # scaffold config.toml + run dir
|
||||
giant train path/to/steps.parquet # train (flow stage 1 + wgan stage 2, by default)
|
||||
dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl # needed for rollout
|
||||
giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
|
||||
|
||||
giant analyze prep rollout.yaml && giant analyze render <run_dir> --gallery # rollout-vs-Geant4 diagnostics
|
||||
```
|
||||
|
||||
Every command takes `--help` for the full flag list, and `--config config.toml` for anything not exposed as a flag.
|
||||
Every command takes `--help` for its full flag list, and `--config config.toml`
|
||||
for anything not exposed as a flag.
|
||||
|
||||
## Architecture
|
||||
---
|
||||
|
||||
A **two-stage model**, checkpointed together. Either stage's outcome can be produced by one of three interchangeable generative objectives (`--stage1-generator`/`--stage2-generator`, or `--mode` to set both at once): `flow` (conditional flow matching, ODE-sampled in ~10 steps), `ddpm` (denoising diffusion), or `wgan` (single-pass WGAN-GP generator/critic).
|
||||
<details>
|
||||
<summary><h2 style="display:inline">Architecture</h2></summary>
|
||||
|
||||
**Stage 1 — primary step.** Predicts the 9D post-step outcome (`giant/constants.py:LOCAL_TARGET_NAMES`) from the pre-step conditioning:
|
||||
**Stage 1 — primary step.** Predicts the 9D post-step outcome
|
||||
(`giant/constants.py:LOCAL_TARGET_NAMES`) from the pre-step conditioning:
|
||||
|
||||
| Index | Variable | Encoding |
|
||||
|-------|----------|----------|
|
||||
| 0 | `step_length` [mm] | log |
|
||||
| 1–2 | `edep_logit`, `sec_logit` | ALR coords of the deposit/secondary/post-energy simplex |
|
||||
| 3–5 | `post_dir` in local frame | unit vector |
|
||||
| 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector |
|
||||
| 1–2 | `edep_logit`, `sec_logit` | ALR coordinates of the deposit / secondary / post-energy simplex |
|
||||
| 3–5 | `post_dir` | unit vector, local frame (`pre_dir = ẑ`) |
|
||||
| 6–8 | `travel_dir` (`post_pos − pre_pos`) | unit vector, local frame |
|
||||
|
||||
- Energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — conservation is architectural, not learned.
|
||||
- `post_dir`/`travel_dir` live in the frame where `pre_dir = ẑ`. `post_pos` isn't a target — it's reconstructed as `pre_pos + step_length * world_frame(travel_dir)`.
|
||||
Energy logits decode via `softmax([edep_logit, sec_logit, 0]) × pre_E`, so
|
||||
`edep + e_sec + post_E == pre_E` holds exactly. `post_pos` is not itself a
|
||||
target — it's reconstructed as `pre_pos + step_length · world_frame(travel_dir)`,
|
||||
since duplicating that magnitude in a second target would let the two drift out
|
||||
of sync.
|
||||
|
||||
**Stage 2 — secondaries.** Conditioned on the pre-step state and Stage 1's outcome, it generates the variable-length list of secondary particles. Two decoding strategies (`--stage2-decoder`):
|
||||
**Stage 2 — secondaries.** Conditioned on the pre-step state and Stage 1's
|
||||
outcome, it generates the variable-length secondary list, one token at a time in
|
||||
descending-energy order (`autoregressive`, default) or all `K_MAX` slots in one
|
||||
masked pass (`one_shot`). Autoregressive tokens condition on a running history —
|
||||
`markov` (previous token only) or `attention` (causal self-attention, KV-cached
|
||||
at inference). Either way, secondary energies stick-break the `e_sec` budget
|
||||
handed down from Stage 1, so the whole chain conserves energy. A secondary's
|
||||
species is represented `onehot` (categorical, top-N PDG codes + "other"),
|
||||
`physical` (continuous log-mass/charge), or `embedding` (nearest-neighbour
|
||||
lookup).
|
||||
|
||||
- `autoregressive` — emits secondaries one at a time in descending-energy order, each token conditioned on a running history of prior tokens (`markov`: previous token only, or `attention`: causal self-attention, KV-cached at inference)
|
||||
- `one_shot` — all `K_MAX` slots generated in a single forward pass, masked past the predicted `n_sec`
|
||||
**Conditioning (15D).** Pre-step position / energy / direction / layer, plus
|
||||
particle mass/charge and material Z_eff/A_eff/density/X0/λ_int — encoded the
|
||||
same three ways as secondary species above, configured *independently* per axis
|
||||
(`conditioning.particle.type` / `conditioning.material.type`). `physical`
|
||||
computes rather than looks up, so it generalizes to species and materials
|
||||
outside the training menu; that's the default. `n_sec`/`e_sec` are always model
|
||||
outputs, never conditioning inputs.
|
||||
|
||||
Either way, secondary energies stick-break the `e_sec` budget handed down from Stage 1, so the full chain conserves energy. A secondary's particle identity is represented as `onehot` (categorical, top-N PDG codes + "other"), `physical` (continuous log-mass/charge), or `embedding` (nearest-neighbour lookup).
|
||||
**Composable by design** — every stage assembles from small registries, so
|
||||
swapping one axis doesn't touch the others:
|
||||
|
||||
**Conditioning.** Pre-step position/energy/direction/layer, plus particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, encoded the same three ways as particle identity above. The particle and material axes are configured independently (`conditioning.particle.type` / `conditioning.material.type`; `--conditioning` sets both at once) and may mix — the `physical` representation generalizes to species/materials outside the training menu since it's computed rather than looked up. `n_sec`/`e_sec` are always model outputs, never conditioning inputs.
|
||||
| Registry | Choices |
|
||||
|---|---|
|
||||
| Objective | `flow` (matching, ~10-step ODE sample) · `ddpm` (denoising diffusion) · `wgan` (single-pass GAN) |
|
||||
| Trunk | `resmlp` · `none`, optionally MoE-routed (`RoutedTrunk`) |
|
||||
| Router | `energy` · `pdg` · `process` · `composed` · `none` — soft-mixed at train time, **top-1 dispatched at eval time**, which is the actual inference-speed win |
|
||||
| History (stage 2 AR) | `markov` · `attention` · `none` |
|
||||
|
||||
**MoE routing** (`--router`): a pluggable `Router` (`energy`/`pdg`/`process`/`composed` axes) top-1-dispatches each row to one of several small expert trunks at eval time, instead of running one monolithic trunk. The CLI flags configure Stage 1's router; Stage 2 has its own `stage2_model.router` block, config-file only.
|
||||
</details>
|
||||
|
||||
## Data
|
||||
<details>
|
||||
<summary><h2 style="display:inline">Configuration</h2></summary>
|
||||
|
||||
- Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from ROOT via `dwarf convert`. One row = one Geant4 step.
|
||||
- **Conditioning (pre-step) columns:** `event_id`, `pdg`, `pre_x`/`pre_y`/`pre_z`, `pre_E`, `pre_dx`/`pre_dy`/`pre_dz` (direction), `material`, `layer_id`.
|
||||
- **Primary outcome (post-step) columns:** `post_x`/`post_y`/`post_z`, `post_E`, `post_dx`/`post_dy`/`post_dz`, `step_length`, `edep` (energy deposited in this step), `e_sec` (total energy carried off by secondaries), `child_track_ids` (its length gives `n_sec`).
|
||||
- **Secondary columns**, one variable-length list per step: `sec_pdg_list`, `sec_E_list`, `sec_dx_list`/`sec_dy_list`/`sec_dz_list` — padded/truncated to `K_MAX` (15) slots on load, ordered by descending energy.
|
||||
- **Optional:** `process` — the physics process that produced the step (e.g. `compt`, `phot`, `eBrem`); a post-step label used only as classifier supervision (`ProcessRouter`), never as conditioning.
|
||||
- Train/val split is by `event_id` (`--seed`-controlled), not row shuffle, so correlated steps from the same shower never leak across the split.
|
||||
- Loading a directory or `.manifest` of multiple parquet files (each one Geant4 job, `event_id` restarting from 0) offsets each file's `event_id`s by a fixed per-file stride so ids stay globally unique across files.
|
||||
Every default lives in one place: frozen dataclasses in `giant/config.py`,
|
||||
composed into `GiantConfig` (`conditioning` / `stage1_model` / `stage2_model` /
|
||||
`train`). `DEFAULT_CONFIG` is *generated* from `GiantConfig().to_dict()` rather
|
||||
than hand-maintained, so the dataclasses can't drift from what actually gets
|
||||
merged. TOML config keys are validated against that shape — an unknown key is
|
||||
rejected with a did-you-mean suggestion. Precedence: CLI flag > `--config` file
|
||||
> default.
|
||||
|
||||
## Project structure
|
||||
```toml
|
||||
# config.toml — resolved shape of the four blocks
|
||||
[conditioning]
|
||||
particle.type = "physical"
|
||||
material.type = "physical"
|
||||
|
||||
```
|
||||
giant/
|
||||
├── giant/
|
||||
│ ├── data/
|
||||
│ │ ├── loader.py # parquet → numpy arrays (incl. streaming/chunked reads)
|
||||
│ │ ├── transforms.py # log transforms, local-frame rotation, energy simplex, secondary encode/decode
|
||||
│ │ ├── dataset.py # StepsDataset / StreamingStepsDataset (PyTorch)
|
||||
│ │ └── setup_cache.py # sidecar cache for the pre-epoch setup scan (vocab/split/normalizers)
|
||||
│ ├── model/
|
||||
│ │ ├── models.py # Stage1Model, Stage2OneShot, Stage2Autoregressive, CriticModel
|
||||
│ │ ├── builders.py # build_models / build_critics — config dict → assembled stage models
|
||||
│ │ ├── encoders.py # ConditionEncoder (physical / embedding / onehot, per axis)
|
||||
│ │ ├── layers.py # ResBlock/AdaLNResBlock registry, SinusoidalEmbedding, MLP heads
|
||||
│ │ ├── trunks.py # trunk registry (resmlp, none) + RoutedTrunk (MoE expert bodies)
|
||||
│ │ ├── routers.py # Router registry: energy / pdg / process / composed / none
|
||||
│ │ ├── history.py # stage-2 AR history encoders: markov / attention (KV-cached) / none
|
||||
│ │ ├── objectives.py # flow / ddpm / wgan objective registry
|
||||
│ │ ├── schedule.py # CosineSchedule (DDPM) and flow matching utilities
|
||||
│ │ ├── wgan.py # WGAN-GP gradient penalty / critic / generator losses
|
||||
│ │ ├── summary.py # build-only introspection behind `giant model summary`
|
||||
│ │ ├── _legacy.py # v0.2 checkpoint model_config/state-dict migration
|
||||
│ │ └── network.py # re-export shim over all of the above
|
||||
│ ├── constants.py # output/conditioning dims, K_MAX, secondary slot layout, schema keys
|
||||
│ ├── cond_layout.py # single source of truth for the cond_cont/cond_cat column layout
|
||||
│ ├── particles.py # PDG → (mass, charge) decode, incl. nuclear/ion codes; onehot/embedding secondary-identity decode
|
||||
│ ├── materials.py # material name → (Z_eff, A_eff, density, X0, λ_int)
|
||||
│ ├── config.py # default hyperparameters, TOML config merging, device autodetect
|
||||
│ ├── pipeline.py # builds datasets/normalizers and kicks off a training run (with setup-stage caching)
|
||||
│ ├── training/ # two-stage training: loop, per-stage trainers, metrics, checkpointing
|
||||
│ │ ├── loop.py # epoch loop, graceful shutdown, best-checkpoint selection
|
||||
│ │ ├── trainers.py # StageSpec + flow/ddpm and WGAN-GP per-stage trainers
|
||||
│ │ ├── stage2_inputs.py# ground-truth stage-2 targets + autoregressive/teacher-forcing inputs
|
||||
│ │ ├── metrics.py # MetricsCollector: metrics.csv columns, W&B logging, progress/summary
|
||||
│ │ ├── amp.py # bf16 autocast (`train.precision`)
|
||||
│ │ ├── plots.py # training-progress plots (`giant analyze metrics`)
|
||||
│ │ └── checkpoint.py # checkpoint assembly/restore (format unchanged since v0.2)
|
||||
│ ├── sample.py # DDPM / DDIM / flow matching / WGAN samplers + secondary sampling
|
||||
│ ├── checkpoint_io.py # checkpoint → ready-to-run models/normalizers (predict + rollout)
|
||||
│ ├── geometry.py # GeometryOracle: position → (material, layer_id, escaped) for rollout
|
||||
│ ├── rollout.py # autoregressive shower rollout driver
|
||||
│ ├── validate.py # step-level marginal + KL-divergence validation
|
||||
│ ├── _migration.py # shared v0.2 → v0.3 facts used by both migration surfaces
|
||||
│ ├── analysis/ # rollout-vs-reference analysis pipeline (see `giant analyze` below)
|
||||
│ │ ├── sources.py # canonical LazyFrames + secondary view
|
||||
│ │ ├── variables.py # per-step value expressions shared by range sizing and the catalog
|
||||
│ │ ├── reduce.py # streaming reduction primitives (hist1d, per-event scalars, profiles, ...)
|
||||
│ │ ├── grouping.py # fixed bin edges + energy/pdg/material group sets
|
||||
│ │ ├── context.py # resolves grouping into `shared.json` once per run
|
||||
│ │ ├── reduced.py # Partial/Reduced — the compact JSON a compute job emits
|
||||
│ │ ├── catalog.py # declarative PlotSpec registry (`giant analyze list`)
|
||||
│ │ ├── router_gating.py / type_embedding_distance.py # checkpoint-bound diagnostics
|
||||
│ │ ├── runtime_estimate.py # per-(plot, chunk) walltime estimates for submit
|
||||
│ │ ├── condor.py # prep / compute-one / merge / submit-description plumbing
|
||||
│ │ └── render.py # PDFs + HTML gallery (only module importing plotstyle/LaTeX)
|
||||
│ └── cli.py # `giant train` / `new-run` / `model summary` / `predict` / `rollout` / `analyze`
|
||||
├── giant/tools/ # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
|
||||
│ ├── dwarf.py # Typer app: convert, migrate, bump-gen, bump-schema, status,
|
||||
│ │ # update-manifest, create-manifest, make-root,
|
||||
│ │ # build-geometry-oracle, warm-cache, hparam-scan
|
||||
│ ├── steps_to_parquet.py # ROOT → parquet conversion (uproot/awkward/polars) — `dwarf convert`
|
||||
│ ├── steps_to_parquet_parallel.py # fan out conversion over several ROOT files — `dwarf convert --jobs N`
|
||||
│ ├── migrate_geant_steps.py # one-time move into the raw/processed/pools/derived layout — `dwarf migrate`
|
||||
│ ├── bump_dataset_version.py # cut a new raw gen or parquet schema, with a logged reason —
|
||||
│ │ # `dwarf bump-gen` / `bump-schema` / `status` / `update-manifest` / `create-manifest`
|
||||
│ ├── create_root_files.py # generate new ROOT shards via a minicalosim executable — `dwarf make-root`
|
||||
│ ├── geometry_oracle.py # fit a position → (material, layer_id) oracle — `dwarf build-geometry-oracle`
|
||||
│ ├── warm_setup_cache.py # precompute `giant train`'s setup-stage sidecar — `dwarf warm-cache`
|
||||
│ ├── hparam_scan.py # hyperparameter grid scan over `giant train` runs — `dwarf hparam-scan`
|
||||
│ └── profile_analysis_costs.py # profiling helper for the `giant analyze` reduction pipeline
|
||||
└── tests/
|
||||
[stage1_model]
|
||||
generator = "flow"
|
||||
|
||||
[stage2_model]
|
||||
generator = "wgan"
|
||||
decoder = "autoregressive"
|
||||
|
||||
[train]
|
||||
epochs = 100
|
||||
batch_size = 4096
|
||||
lr = 3e-4
|
||||
```
|
||||
|
||||
## Setup
|
||||
Some knobs only exist in the config file, with no CLI flag:
|
||||
`stage2_model.autoregressive.teacher_forcing`/`.history`,
|
||||
`stage2_model.particle_type.target`/`.class_weighting`,
|
||||
`stage2_model.n_sec.mode`/`.owner`, `conditioning.share_stages`,
|
||||
`stage2_model.router.*`, and the finer `router` knobs (`lambda_balance`,
|
||||
`gumbel`, `learn_width`, …).
|
||||
|
||||
`configs/` holds kept reference configs — `baseline.toml` is the fixed
|
||||
comparison point every experimental variant (routed trunk, WGAN, attention
|
||||
history, embedding conditioning) is a single edit away from. v0.2 flat-schema
|
||||
configs and checkpoints load and auto-migrate.
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><h2 style="display:inline">Data</h2></summary>
|
||||
|
||||
Input is parquet — one row per Geant4 step — from
|
||||
[miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from
|
||||
ROOT via `dwarf convert`.
|
||||
|
||||
| Group | Columns |
|
||||
|---|---|
|
||||
| **Conditioning (pre-step)** | `event_id`, `pdg`, `pre_x`/`pre_y`/`pre_z`, `pre_E`, `pre_dx`/`pre_dy`/`pre_dz`, `material`, `layer_id` |
|
||||
| **Primary outcome (post-step)** | `post_x`/`post_y`/`post_z`, `post_E`, `post_dx`/`post_dy`/`post_dz`, `step_length`, `edep`, `e_sec`, `child_track_ids` (length → `n_sec`) |
|
||||
| **Secondaries** (variable-length lists) | `sec_pdg_list`, `sec_E_list`, `sec_dx_list`/`sec_dy_list`/`sec_dz_list` — padded/truncated to `K_MAX = 15` slots, descending energy |
|
||||
| **Optional** | `process` — physics-process label, classifier supervision only (`ProcessRouter`), never conditioning |
|
||||
|
||||
Train/val split is by `event_id` (`--seed`-controlled), not row shuffle, so a
|
||||
shower's correlated steps never straddle the split. Loading a directory or
|
||||
`.manifest` of several parquet files offsets each file's `event_id`s by a
|
||||
per-file stride so ids stay globally unique. The pre-epoch setup scan (vocab
|
||||
maps, event split, normalizer stats) persists to a sidecar cache
|
||||
(`--cache-setup`/`--rebuild-setup-cache`), precomputable ahead of time via
|
||||
`dwarf warm-cache`.
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><h2 style="display:inline">CLI reference</h2></summary>
|
||||
|
||||
**`giant`** — train, run, and analyze the surrogate:
|
||||
|
||||
| Command | Does |
|
||||
|---|---|
|
||||
| `new-run` | scaffold a `config.toml` + run directory from flags |
|
||||
| `train DATA` | train the two-stage model |
|
||||
| `model summary` | build-only parameter counts, without training |
|
||||
| `predict DATA --checkpoint …` | per-step predictions from a checkpoint |
|
||||
| `rollout DATA --checkpoint … --geometry …` | full autoregressive shower rollout |
|
||||
| `analyze prep/submit` | build a run dir; `submit` also queues HTCondor compute jobs |
|
||||
| `analyze compute-one` / `merge-one` | one plot × chunk reduction / merge (what a condor job runs) |
|
||||
| `analyze render <run_dir> --gallery` | merge chunks → styled PDFs + HTML gallery (local, needs LaTeX) |
|
||||
| `analyze metrics <train_run_dir>` | training-progress plots from `metrics.csv` |
|
||||
| `analyze list` | every catalog plot id |
|
||||
|
||||
**`dwarf`** — dataset/tooling CLI:
|
||||
|
||||
| Command | Does |
|
||||
|---|---|
|
||||
| `convert` | ROOT Steps tree → parquet (`--jobs N` fans out) |
|
||||
| `migrate` | one-time move into the raw/processed/pools/derived layout |
|
||||
| `bump-gen` / `bump-schema` / `status` | dataset versioning |
|
||||
| `update-manifest` / `create-manifest` | point/build a manifest of parquet files |
|
||||
| `make-root` | generate new ROOT shards via a minicalosim executable |
|
||||
| `build-geometry-oracle` | fit position → (material, layer_id) for rollout |
|
||||
| `warm-cache` | precompute `giant train`'s setup-stage sidecar |
|
||||
| `hparam-scan` | grid-scan dropout × n_blocks × hidden_dim |
|
||||
|
||||
Worth knowing on `giant train` (full surface behind `--help`):
|
||||
`--mode {flow,ddpm,wgan}` / `--stage1-generator` / `--stage2-generator`,
|
||||
`--stage2-decoder {autoregressive,one_shot}`, `--conditioning
|
||||
{physical,embedding,onehot}`, `--router` / `--router-type` / `--n-experts` /
|
||||
`--router-axis`, `--stage{1,2}-init-from` + `--stage{1,2}-freeze` (retrain one
|
||||
stage against a fixed other one), `--precision {fp32,bf16}`, `--wandb`.
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><h2 style="display:inline">Rollout & analysis</h2></summary>
|
||||
|
||||
`giant rollout` seeds showers from each event's highest-energy entry step, then
|
||||
autoregressively steps the model to completion — advancing all active tracks
|
||||
breadth-first, batched — pushing secondaries as new tracks and looking up
|
||||
`material`/`layer_id` from the geometry oracle each step. Tracks terminate on
|
||||
one of six reasons (energy cutoff, max steps, detector escape, natural end,
|
||||
unknown pdg, max tracks); every reason but escape deposits the remaining energy
|
||||
locally, so showers conserve energy by construction — only `escaped` counts as
|
||||
leakage.
|
||||
|
||||
`giant analyze` compares one or more rollouts against a single held-out
|
||||
reference: `prep` resolves shared bin edges/groups once, `submit`/`compute-one`
|
||||
run each (plot, `event_id`-disjoint chunk) pair as a polars/numpy-only HTCondor
|
||||
job, `render` merges the chunks and produces the styled PDFs + HTML gallery
|
||||
locally (the only step that needs LaTeX). Each rollout gets its own colored
|
||||
series against one shared reference line. `giant analyze metrics` is a separate
|
||||
entry point — training-progress plots straight from a run's `metrics.csv`.
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><h2 style="display:inline">Install</h2></summary>
|
||||
|
||||
| Extra | Adds | For |
|
||||
|---|---|---|
|
||||
| `cpu` **or** `cuda` | torch 2.3.x | required — mutually exclusive, pick one |
|
||||
| `geometry` | scikit-learn | `dwarf build-geometry-oracle`, rollout |
|
||||
| `analysis` | matplotlib, plotstyle | `giant analyze render` |
|
||||
| `convert` | uproot, awkward | `dwarf convert` |
|
||||
| `wandb` | wandb | `giant train --wandb` |
|
||||
| `dev` | pytest, ruff, ty, + all of the above | development |
|
||||
|
||||
```bash
|
||||
uv sync --extra cpu # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
|
||||
uv sync --extra cpu --extra dev # add dev tools (pytest, ruff, ty)
|
||||
uv sync --extra cpu --extra geometry # add scikit-learn, for `dwarf build-geometry-oracle` / rollout
|
||||
uv sync --extra cpu --extra analysis # matplotlib/polars/plotstyle, for `giant analyze render`
|
||||
uv sync --extra cpu --extra convert # uproot/awkward/polars, for `dwarf convert`
|
||||
uv sync --extra cpu --extra wandb # W&B logging (`giant train --wandb`)
|
||||
uv sync --extra cpu --extra dev # everything needed to develop
|
||||
```
|
||||
|
||||
The `dev` extra pulls in `convert`, `analysis`, `geometry` and `wandb` as well.
|
||||
Plain `uv sync` with no extra installs **no torch at all** — always include
|
||||
`--extra cpu` or `--extra cuda`.
|
||||
|
||||
`cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x). Plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
|
||||
</details>
|
||||
|
||||
## Training, prediction, rollout
|
||||
<details>
|
||||
<summary><h2 style="display:inline">Development</h2></summary>
|
||||
|
||||
```bash
|
||||
giant new-run --hidden-dim 512 --lr 3e-4 --comment "..." # scaffold a config.toml + run dir
|
||||
giant train path/to/steps.parquet # train (flow stage 1 + wgan stage 2, default)
|
||||
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
|
||||
|
||||
dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl # position → material/layer_id
|
||||
giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
|
||||
```
|
||||
|
||||
Useful flags on `giant train`:
|
||||
|
||||
- `--mode {flow,ddpm,wgan}` sets both stages' objective at once; `--stage1-generator`/`--stage2-generator` override per stage
|
||||
- `--stage2-decoder {autoregressive,one_shot}` — Stage 2 decoding strategy (see Architecture)
|
||||
- `--conditioning {physical,embedding,onehot}` — conditioning representation
|
||||
- `--router` / `--router-type` / `--n-experts` / `--router-axis` — MoE routing
|
||||
- `--stage2-stage1-context {truth,sampled}` — feed Stage 2 the ground-truth or the model's own sampled Stage-1 outcome (annealable via `stage2_model.ctx_p_start`/`ctx_p_end`)
|
||||
- `--precision {fp32,bf16}` — bf16 autocast in the training loop
|
||||
- `--wandb` — log per-epoch metrics to Weights & Biases (needs `uv sync --extra wandb`); metric names are `<stage>/<split>/<metric>` plus an unprefixed run-level tail, all derived from `giant/training/trainers.py` `MetricSpec`s
|
||||
- `--no-cache-setup` / `--rebuild-setup-cache` — control the setup-stage sidecar cache (vocab maps, event split, normalizer stats); `dwarf warm-cache` precomputes it
|
||||
- `--stage1-init-from`/`--stage2-init-from` (checkpoint `.pt`) + `--stage1-freeze`/`--stage2-freeze` — load a stage's weights from another checkpoint and never update them, so the other stage can be retrained alone against a fixed, known-good one while still producing a complete, rollout-capable checkpoint
|
||||
|
||||
Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_model.autoregressive.teacher_forcing`/`.history`, `stage2_model.particle_type.target`/`.class_weighting`, `stage2_model.n_sec.mode`/`.owner`, `conditioning.share_stages`, `stage*_model.trunk.*` and the finer `router` knobs (`lambda_balance`, `gumbel`, `learn_width`, …). `configs/` holds kept reference configs. v0.2 flat-schema configs and checkpoints load fine (auto-migrated).
|
||||
|
||||
`giant rollout` seeds showers from each event's highest-energy entry step, then autoregressively steps the model to completion, pushing secondaries as new tracks and looking up `material`/`layer_id` from the geometry oracle each step. Tracks terminate on energy cutoff, max steps, detector escape, or natural end; energy is deposited locally on every stop except escape, so showers conserve energy by construction.
|
||||
|
||||
## Validation and analysis
|
||||
|
||||
- `giant.validate.validate_marginals` — step-level marginal + KL-divergence checks during training (`--validate-every`)
|
||||
- `giant analyze` — deeper rollout-vs-reference diagnostics (marginals by energy/pdg/material, per-event totals, shower profiles, species share, leakage, secondaries):
|
||||
|
||||
```bash
|
||||
giant analyze submit rollout.yaml --accounting-group cms # prep + one HTCondor job per plot × chunk (compute only)
|
||||
giant analyze submit a.yaml b.yaml --accounting-group cms --label flow --label wgan # N rollouts vs one shared reference
|
||||
giant analyze render <run_dir> --gallery # local: merge chunks, then styled PDFs + HTML gallery (needs LaTeX)
|
||||
|
||||
giant analyze list # every catalog plot id
|
||||
giant analyze prep rollout.yaml --chunks 8 # just the run directory, no submission
|
||||
giant analyze compute-one --id marginal_edep --run-dir <run_dir> --chunk 0 # what a condor job runs
|
||||
giant analyze merge-one --id marginal_edep --run-dir <run_dir> # merge one plot's chunks (debugging)
|
||||
```
|
||||
|
||||
`<run_dir>` defaults to `<cwd>/analysis_runs/analysis_<id>` (`--run-dir` overrides it; `prep`/`submit` print it). Multiple rollout YAMLs must all name the same reference (`dataset`) file; each renders as its own colored series against one reference line/panel. Compute jobs are polars/numpy only; only `render` needs LaTeX, so it always runs locally.
|
||||
|
||||
Separately, `giant analyze metrics <train_run_dir>` renders training-progress plots (loss/lr/accuracy/grad-norm/router/wgan/throughput) straight from a training run's `metrics.csv`.
|
||||
|
||||
## Development
|
||||
|
||||
```bash
|
||||
uv run pytest # run tests
|
||||
uv run pytest # 964 tests
|
||||
uv run ruff check . # lint
|
||||
uv run ruff format . # format
|
||||
uv run ty check . # type check
|
||||
```
|
||||
|
||||
Gitea Actions (`.gitea/workflows/ci.yml`) runs lint + format-check + type-check
|
||||
+ tests on every push and PR; merges to `master` auto-bump the patch version
|
||||
and regenerate `CHANGELOG.md` — don't hand-edit either.
|
||||
|
||||
</details>
|
||||
|
||||
## License
|
||||
|
||||
Not yet decided — treat this repository as all-rights-reserved until a
|
||||
`LICENSE` file is added.
|
||||
|
||||
+20
-5
@@ -84,15 +84,30 @@ dropout = 0.0
|
||||
# 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:
|
||||
# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated — but see
|
||||
# the row-compaction note below, which changes the INFERENCE side of this:
|
||||
# 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.
|
||||
#
|
||||
# Row compaction (landed after the above measurement): at inference,
|
||||
# sample.sample_secondaries_ar used to loop `for k in range(k_max)`
|
||||
# unconditionally — all 15 slots regardless of predicted n_sec — so a flow
|
||||
# AR token cost k_max * steps = 150 stage-2 calls per physics step. It now
|
||||
# drops a row from the batch the moment its own secondary count is
|
||||
# exhausted, so the real inference cost is ~n_sec * steps stage-2 calls
|
||||
# (this checkpoint's own rollout measured 0.382 secondaries/step — see
|
||||
# giant-baseline-flow-ar-rollout-validation.md), not k_max * steps. A CPU
|
||||
# micro-benchmark at that multiplicity (giant/model/history.py's
|
||||
# hidden_dim=512/6-block shape, k_max=15, batch 512) measured 17.6-22.9x
|
||||
# fewer wall-clock seconds for the AR loop alone (markov/attention history
|
||||
# respectively) — directional only (CPU, synthetic n_sec distribution, not
|
||||
# an end-to-end rollout); the 8.1x inference ratio above is now stale and
|
||||
# should be re-measured on GPU via a real rollout + `eval_cost_per_step`
|
||||
# once one is run against this checkpoint. Training cost (the 6.5x/29.5k
|
||||
# figures) is untouched by this: teacher_forcing = "always" here never
|
||||
# calls the AR sampler at train time (see [stage2_model.autoregressive]).
|
||||
decoder = "autoregressive"
|
||||
generator = "flow"
|
||||
hidden_dim = 512
|
||||
|
||||
@@ -4,6 +4,7 @@ for the `HISTORY_REGISTRY`/`build_history` factory, which mirrors
|
||||
`giant.model.routers`'s `Router`/`ROUTER_REGISTRY` pattern (gitea #35)."""
|
||||
|
||||
import inspect
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
@@ -34,6 +35,17 @@ class HistoryEncoder(nn.Module):
|
||||
def step(self, feat: torch.Tensor, has_prev: torch.Tensor, cache: object) -> tuple[torch.Tensor, object]:
|
||||
return self.forward(feat, has_prev), cache
|
||||
|
||||
def select_cache(self, cache: object, idx: torch.Tensor) -> object:
|
||||
"""Row-compacts an inference cache (`init_cache`/`step`'s state) down
|
||||
to `idx` — used by `giant.sample.sample_secondaries_ar`'s row
|
||||
compaction to keep a shrinking active-row set's cache aligned as rows
|
||||
finish generating. Default here matches `init_cache`/`step`'s O(1)
|
||||
default: `cache` is always `None`, so there's nothing to index —
|
||||
correct for any encoder whose per-step state doesn't carry a batch
|
||||
dimension (`MarkovHistory` has no cache at all; its running state is
|
||||
`prev_repr`/`remaining`, compacted directly by the caller)."""
|
||||
return cache
|
||||
|
||||
|
||||
HISTORY_REGISTRY: dict[str, type[HistoryEncoder]] = {}
|
||||
|
||||
@@ -215,3 +227,13 @@ class AttentionHistory(HistoryEncoder):
|
||||
x, kv_new = block.step(x, kv)
|
||||
new_cache.append(kv_new)
|
||||
return x, new_cache
|
||||
|
||||
def select_cache(self, cache: object, idx: torch.Tensor) -> list[torch.Tensor | None]:
|
||||
"""Row-compacts every block's `(B, T, dim)` KV cache down to `idx`
|
||||
along its batch dimension — see `HistoryEncoder.select_cache`. `idx`
|
||||
may be a long index tensor or a boolean mask (`giant.sample`'s AR
|
||||
loop uses both). `None` entries (a block that has never seen a
|
||||
`step` call yet) stay `None`."""
|
||||
assert isinstance(cache, list)
|
||||
cache_t = cast("list[torch.Tensor | None]", cache)
|
||||
return [None if kv is None else kv[idx] for kv in cache_t]
|
||||
|
||||
@@ -614,6 +614,14 @@ class Stage2Autoregressive(StageModel):
|
||||
slot — see `AttentionHistory.step`'s docstring."""
|
||||
return self.history_encoder.step(token_feat, has_prev, cache)
|
||||
|
||||
def select_history_cache(self, cache, idx: torch.Tensor):
|
||||
"""Row-compacts `cache` (from `init_history_cache`/`history_step`)
|
||||
down to `idx` — see `HistoryEncoder.select_cache`. Used by
|
||||
`giant.sample.sample_secondaries_ar`'s active-row compaction to keep
|
||||
the cache aligned with a shrinking batch as rows finish generating
|
||||
across AR slots."""
|
||||
return self.history_encoder.select_cache(cache, idx)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x_t: torch.Tensor,
|
||||
|
||||
+110
-50
@@ -231,6 +231,7 @@ def sample_secondaries_ar(
|
||||
stage1_out: torch.Tensor,
|
||||
n_sec_pred: torch.Tensor | None,
|
||||
steps: int = 10,
|
||||
full_length: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""`Stage2Autoregressive` inference loop: one token at a time, in
|
||||
descending-energy slot order, up to `k_max` sequential calls. Unlike
|
||||
@@ -242,19 +243,18 @@ def sample_secondaries_ar(
|
||||
expressiveness.
|
||||
|
||||
A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass —
|
||||
the "K sequential forwards" cost applies per-token here, not
|
||||
once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
|
||||
physics step (or ~`n_sec * steps` under `n_sec_pred=None` below, once
|
||||
every row in the batch has stopped).
|
||||
the "K sequential forwards" cost applies per-token here, not once, so a
|
||||
flow/ddpm AR run costs ~`n_sec * steps` model calls per physics step
|
||||
(measured on `configs/baseline.toml`: 0.382 secondaries/step at rollout
|
||||
time), not `k_max * steps` — see the row-compaction paragraph below.
|
||||
|
||||
`n_sec_pred`, if given, fixes each row's secondary count up front (as
|
||||
resolved by `resolve_n_sec` — `n_sec.mode` in `("head", "truth")`, or a
|
||||
stop-token decoder driven by `_assemble_stage2_ar_inputs_scheduled`'s
|
||||
ground-truth `n_sec`, which must run the *full* `k_max`-length free-
|
||||
running self-sample regardless of the decoder's own stop head — the
|
||||
scheduled-sampling training contract does not truncate). This always
|
||||
runs the full `k_max`-iteration loop, masking by the given count at the
|
||||
end exactly as before.
|
||||
`n_sec_pred`, if given (as resolved by `resolve_n_sec` — `n_sec.mode` in
|
||||
`("head", "truth")`, or a stop-token decoder driven by
|
||||
`_assemble_stage2_ar_inputs_scheduled`'s ground-truth `n_sec`, which must
|
||||
run the *full* `k_max`-length free-running self-sample regardless of the
|
||||
decoder's own stop head — the scheduled-sampling training contract does
|
||||
not truncate, see `full_length` below) fixes each row's secondary count
|
||||
up front.
|
||||
|
||||
`n_sec_pred=None` is only valid when `sec_decoder.stop_head` is set
|
||||
(`n_sec.mode = "stop_token"`): before generating each slot's token, that
|
||||
@@ -263,11 +263,35 @@ def sample_secondaries_ar(
|
||||
why this needs no extra state) decides whether generation should have
|
||||
already stopped, per `sec_decoder.n_sec_sampling` ("greedy": threshold at
|
||||
0; "sample": a Bernoulli draw at `sigmoid(logit)`). A row's own
|
||||
`n_sec_pred` is the first slot index where this fires; once every row in
|
||||
the batch has fired, the loop breaks before spending a model call on the
|
||||
next slot's token — the average-case cost win the docstring above
|
||||
describes. A row that never fires within `k_max` is capped there
|
||||
(`K_MAX` stays a safety cap, not a modeling ceiling).
|
||||
`n_sec_pred` is the first slot index where this fires. A row that never
|
||||
fires within `k_max` is capped there (`K_MAX` stays a safety cap, not a
|
||||
modeling ceiling).
|
||||
|
||||
**Row compaction.** A row that has already produced its `n_sec_pred`
|
||||
tokens (or, under `stop_token`, has already fired its stop logit) has
|
||||
nothing left to contribute — every later slot of that row is masked out
|
||||
of `sec_valid` on return, and downstream consumers (`giant/rollout.py`,
|
||||
`giant/cli.py`) never read it. So unless `full_length=True`, this
|
||||
function drops such rows from the active set entirely instead of running
|
||||
the model on them: `active_idx` starts at every row with `n_sec_pred > 0`
|
||||
(or, under `stop_token`, every row — the first stop decision can fire at
|
||||
slot 0) and only shrinks as rows finish, so slot `k`'s model calls cost
|
||||
`O(active rows)` not `O(B)`. Slots a row never reaches keep their `0.0`
|
||||
zero-init in `sec_cont`/`sec_type` — masked by `sec_valid`, identical to
|
||||
what a full, uncompacted run would have written there before masking.
|
||||
`AttentionHistory`'s KV cache is kept aligned to the shrinking active set
|
||||
via `Stage2Autoregressive.select_history_cache`
|
||||
(`giant.model.history.HistoryEncoder.select_cache`) every time the set
|
||||
shrinks; `MarkovHistory`'s O(1) state (`prev_repr`/`remaining`, carried
|
||||
directly rather than through a cache) is compacted the same way.
|
||||
|
||||
`full_length=True` disables all of the above: every row runs the full
|
||||
`k_max`-iteration loop regardless of `n_sec_pred`/stop decisions, exactly
|
||||
reproducing the pre-compaction behaviour. Required by
|
||||
`_assemble_stage2_ar_inputs_scheduled`'s scheduled-sampling self-sample,
|
||||
whose training contract needs a real prediction at every slot up to
|
||||
`k_max` (mixed per-slot against ground truth) even past a row's own
|
||||
`n_sec` — see that function's docstring.
|
||||
|
||||
Under `history="attention"` the history encoding is computed once per
|
||||
slot via `Stage2Autoregressive.history_step` (a KV-cache append)
|
||||
@@ -312,11 +336,6 @@ def sample_secondaries_ar(
|
||||
sec_cont = torch.zeros(B, k_max, CONT_SLOT_DIM, device=device)
|
||||
sec_type = torch.zeros(B, k_max, type_dim, device=device)
|
||||
|
||||
# Running per-token state, threaded from one slot to the next.
|
||||
prev_repr = torch.zeros(B, CONT_SLOT_DIM + type_dim, device=device)
|
||||
remaining = torch.ones(B, device=device)
|
||||
history_cache = sec_decoder.init_history_cache()
|
||||
|
||||
use_stop_token = n_sec_pred is None
|
||||
if use_stop_token:
|
||||
assert getattr(sec_decoder, "stop_head", None) is not None, (
|
||||
@@ -324,21 +343,46 @@ def sample_secondaries_ar(
|
||||
"with no stop_head — only valid under stage2_model.n_sec.mode = "
|
||||
"'stop_token'"
|
||||
)
|
||||
finished = torch.zeros(B, dtype=torch.bool, device=device)
|
||||
derived_n_sec = torch.full((B,), k_max, dtype=torch.long, device=device)
|
||||
# Tracks which rows have already recorded a stop, globally by
|
||||
# original batch index — needed even under compaction's own
|
||||
# never-revisit guarantee, because `full_length=True` keeps every
|
||||
# row in `active_idx` for the whole loop, so a row whose stop logit
|
||||
# fires once but flips back below threshold at a later slot (a real
|
||||
# possibility for an untrained/lightly-trained stop_head) must not
|
||||
# have `derived_n_sec` overwritten by that later, spurious re-fire.
|
||||
finished = torch.zeros(B, dtype=torch.bool, device=device)
|
||||
|
||||
# `active_idx`: rows still contributing tokens, indexed into the
|
||||
# original batch. Only ever shrinks (never full_length) or stays fixed
|
||||
# at arange(B) (full_length) — see the row-compaction docstring section.
|
||||
active_idx = torch.arange(B, device=device)
|
||||
if not full_length and not use_stop_token:
|
||||
active_idx = active_idx[n_sec_pred > 0]
|
||||
|
||||
# Running per-token state, already compacted to `active_idx`.
|
||||
prev_repr = torch.zeros(active_idx.numel(), CONT_SLOT_DIM + type_dim, device=device)
|
||||
remaining = torch.ones(active_idx.numel(), device=device)
|
||||
history_cache = sec_decoder.init_history_cache()
|
||||
|
||||
for k in range(k_max):
|
||||
has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
|
||||
history_feat = prev_repr.unsqueeze(1) # (B, 1, CONT_SLOT_DIM + type_dim)
|
||||
remaining_frac = remaining.unsqueeze(1) # (B, 1)
|
||||
slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
|
||||
if active_idx.numel() == 0:
|
||||
break
|
||||
Bc = active_idx.numel()
|
||||
cc = cond_cont.index_select(0, active_idx)
|
||||
ck = cond_cat.index_select(0, active_idx)
|
||||
s1 = stage1_out.index_select(0, active_idx)
|
||||
has_prev = torch.full((Bc, 1), k >= 1, dtype=torch.bool, device=device)
|
||||
history_feat = prev_repr.unsqueeze(1) # (Bc, 1, CONT_SLOT_DIM + type_dim)
|
||||
remaining_frac = remaining.unsqueeze(1) # (Bc, 1)
|
||||
slot_idx = torch.full((Bc, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
|
||||
hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)
|
||||
|
||||
if use_stop_token:
|
||||
stop_logit = sec_decoder.predict_stop(
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
cc,
|
||||
ck,
|
||||
s1,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
@@ -346,21 +390,31 @@ def sample_secondaries_ar(
|
||||
hist=hist,
|
||||
).squeeze(1)
|
||||
if sec_decoder.n_sec_sampling == "sample":
|
||||
stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit)
|
||||
stop_now = torch.rand(Bc, device=device) < torch.sigmoid(stop_logit)
|
||||
else:
|
||||
stop_now = stop_logit >= 0.0
|
||||
derived_n_sec[stop_now & ~finished] = k
|
||||
finished = finished | stop_now
|
||||
if finished.all():
|
||||
break
|
||||
newly_stopped = stop_now & ~finished.index_select(0, active_idx)
|
||||
derived_n_sec[active_idx[newly_stopped]] = k
|
||||
finished[active_idx[stop_now]] = True
|
||||
if not full_length:
|
||||
keep = ~stop_now
|
||||
active_idx = active_idx[keep]
|
||||
cc, ck, s1 = cc[keep], ck[keep], s1[keep]
|
||||
has_prev, remaining_frac, slot_idx = has_prev[keep], remaining_frac[keep], slot_idx[keep]
|
||||
history_feat, hist = history_feat[keep], hist[keep]
|
||||
history_cache = sec_decoder.select_history_cache(history_cache, keep)
|
||||
prev_repr, remaining = prev_repr[keep], remaining[keep]
|
||||
if active_idx.numel() == 0:
|
||||
break
|
||||
Bc = active_idx.numel()
|
||||
|
||||
if objective.is_adversarial:
|
||||
z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
|
||||
z = torch.randn(Bc, 1, sec_decoder.noise_dim, device=device)
|
||||
token = sec_decoder(
|
||||
z,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
cc,
|
||||
ck,
|
||||
s1,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
@@ -368,15 +422,15 @@ def sample_secondaries_ar(
|
||||
hist=hist,
|
||||
)
|
||||
else:
|
||||
x = torch.randn(B, 1, token_dim, device=device)
|
||||
x = torch.randn(Bc, 1, token_dim, device=device)
|
||||
dt = 1.0 / steps
|
||||
for i in range(steps):
|
||||
t = torch.full((B, 1), i * dt, device=device)
|
||||
t = torch.full((Bc, 1), i * dt, device=device)
|
||||
v = sec_decoder(
|
||||
x,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
cc,
|
||||
ck,
|
||||
s1,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
@@ -387,15 +441,15 @@ def sample_secondaries_ar(
|
||||
x = x + v * dt
|
||||
token = x
|
||||
|
||||
token = token.squeeze(1) # (B, token_dim)
|
||||
token = token.squeeze(1) # (Bc, token_dim)
|
||||
cont_k = token[:, :CONT_SLOT_DIM]
|
||||
if type_folded:
|
||||
type_k = token[:, CONT_SLOT_DIM:]
|
||||
else:
|
||||
type_k = sec_decoder.predict_type(
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
cc,
|
||||
ck,
|
||||
s1,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
@@ -403,8 +457,8 @@ def sample_secondaries_ar(
|
||||
hist=hist,
|
||||
).squeeze(1)
|
||||
|
||||
sec_cont[:, k] = cont_k
|
||||
sec_type[:, k] = type_k
|
||||
sec_cont[active_idx, k] = cont_k
|
||||
sec_type[active_idx, k] = type_k
|
||||
|
||||
if target == "onehot":
|
||||
type_for_history = F.one_hot(type_k.argmax(dim=-1), num_classes=type_dim).float()
|
||||
@@ -415,6 +469,12 @@ def sample_secondaries_ar(
|
||||
prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
|
||||
remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)
|
||||
|
||||
if not full_length and not use_stop_token:
|
||||
keep2 = n_sec_pred.index_select(0, active_idx) > (k + 1)
|
||||
active_idx = active_idx[keep2]
|
||||
prev_repr, remaining = prev_repr[keep2], remaining[keep2]
|
||||
history_cache = sec_decoder.select_history_cache(history_cache, keep2)
|
||||
|
||||
resolved_n_sec = derived_n_sec if use_stop_token else n_sec_pred
|
||||
sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < resolved_n_sec.unsqueeze(1)
|
||||
return sec_cont, sec_type, sec_valid
|
||||
|
||||
@@ -291,9 +291,13 @@ def _assemble_stage2_ar_inputs_scheduled(
|
||||
skips self-sampling entirely), so callers can call this unconditionally.
|
||||
|
||||
The free-running estimate is a REAL autoregressive self-sample —
|
||||
`giant.sample.sample_secondaries_ar` under `torch.no_grad()` — not a
|
||||
cheap one-step proxy, so building it costs the same `k_max` (`* steps`
|
||||
for flow) sequential forwards `sample.py` pays at inference, EVERY batch
|
||||
`giant.sample.sample_secondaries_ar` under `torch.no_grad()`, called here
|
||||
with `full_length=True` — not a cheap one-step proxy, so building it
|
||||
costs the full `k_max` (`* steps` for flow) sequential forwards for every
|
||||
row regardless of that row's own secondary count (`full_length=True`
|
||||
disables `sample.py`'s inference-time row compaction — see that
|
||||
function's docstring for why: the mixing below needs a real prediction
|
||||
at every slot up to `k_max`, not just the valid ones). Paid EVERY batch
|
||||
this is called on (paid at train time too whenever teacher_forcing !=
|
||||
"always"). Fully detached: gradient only ever flows
|
||||
through the "real" target path each stage trainer already uses
|
||||
@@ -306,7 +310,7 @@ def _assemble_stage2_ar_inputs_scheduled(
|
||||
|
||||
was_training = model.training
|
||||
sec_cont_pred, sec_type_pred, _ = sample_secondaries_ar(
|
||||
model, cond_cont, cond_cat, stage1_ctx, n_sec, steps=sample_steps
|
||||
model, cond_cont, cond_cat, stage1_ctx, n_sec, steps=sample_steps, full_length=True
|
||||
)
|
||||
if was_training:
|
||||
model.train()
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "giant"
|
||||
version = "0.3.17"
|
||||
version = "0.3.18"
|
||||
description = "Geant4 step-function surrogate via conditional flow matching"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
|
||||
@@ -341,6 +341,87 @@ def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises():
|
||||
sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
|
||||
|
||||
|
||||
# ── Row compaction: full_length=False (default, inference) must agree with
|
||||
# full_length=True (the pre-compaction behaviour, still exercised by
|
||||
# _assemble_stage2_ar_inputs_scheduled's training-time self-sample) ────────
|
||||
|
||||
|
||||
def _zero_randn(*size, **kwargs):
|
||||
"""Drop-in replacement for `torch.randn` that returns zeros of the same
|
||||
shape — makes the ODE/WGAN noise deterministic so a compacted run and a
|
||||
full_length run can be compared row-for-row regardless of how many
|
||||
`torch.randn` calls each makes (compaction changes the batch size, and
|
||||
therefore the RNG stream position, at every slot)."""
|
||||
device = kwargs.get("device")
|
||||
dtype = kwargs.get("dtype")
|
||||
return torch.zeros(*size, device=device, dtype=dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("history", ["markov", "attention"])
|
||||
@pytest.mark.parametrize("generator", ["flow", "wgan"])
|
||||
def test_sample_secondaries_ar_compaction_matches_full_length_head_mode(generator, history, monkeypatch):
|
||||
B, k_max, emb_dim = 4, 5, 6
|
||||
decoder = _stage2_ar("physical", generator, emb_dim=emb_dim, k_max=k_max, history=history)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec_pred = torch.tensor([0, 1, 3, k_max])
|
||||
|
||||
monkeypatch.setattr(torch, "randn", _zero_randn)
|
||||
sec_cont_c, sec_type_c, sec_valid_c = sample_secondaries_ar(
|
||||
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=False
|
||||
)
|
||||
sec_cont_f, sec_type_f, sec_valid_f = sample_secondaries_ar(
|
||||
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=True
|
||||
)
|
||||
|
||||
assert torch.equal(sec_valid_c, sec_valid_f)
|
||||
assert torch.equal(sec_valid_c, torch.arange(k_max).unsqueeze(0) < n_sec_pred.unsqueeze(1))
|
||||
assert torch.allclose(sec_cont_c[sec_valid_c], sec_cont_f[sec_valid_f], atol=1e-4, rtol=1e-4)
|
||||
assert torch.allclose(sec_type_c[sec_valid_c], sec_type_f[sec_valid_f], atol=1e-4, rtol=1e-4)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("generator", ["flow", "wgan"])
|
||||
def test_sample_secondaries_ar_compaction_matches_full_length_stop_token(generator, monkeypatch):
|
||||
"""`n_sec_sampling="greedy"` keeps the stop decision itself deterministic
|
||||
(no `torch.rand` draw), so only `torch.randn` needs zeroing."""
|
||||
B, k_max = 6, 5
|
||||
decoder = _stage2_ar_stop_token("physical", generator, n_sec_sampling="greedy", k_max=k_max)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
|
||||
monkeypatch.setattr(torch, "randn", _zero_randn)
|
||||
sec_cont_c, sec_type_c, sec_valid_c = sample_secondaries_ar(
|
||||
decoder, cond_cont, cond_cat, stage1_out, None, steps=2, full_length=False
|
||||
)
|
||||
sec_cont_f, sec_type_f, sec_valid_f = sample_secondaries_ar(
|
||||
decoder, cond_cont, cond_cat, stage1_out, None, steps=2, full_length=True
|
||||
)
|
||||
|
||||
assert torch.equal(sec_valid_c, sec_valid_f)
|
||||
assert torch.allclose(sec_cont_c[sec_valid_c], sec_cont_f[sec_valid_f], atol=1e-4, rtol=1e-4)
|
||||
assert torch.allclose(sec_type_c[sec_valid_c], sec_type_f[sec_valid_f], atol=1e-4, rtol=1e-4)
|
||||
|
||||
|
||||
def test_sample_secondaries_ar_full_length_ignores_n_sec_pred_zero_rows():
|
||||
"""A row with n_sec_pred == 0 would be dropped from the active set at
|
||||
slot 0 under compaction (full_length=False) — full_length=True must
|
||||
still run the model for it at every slot (only masked by sec_valid at
|
||||
the end), matching _assemble_stage2_ar_inputs_scheduled's contract."""
|
||||
B, k_max, emb_dim = 3, 4, 6
|
||||
decoder = _stage2_ar("physical", "flow", emb_dim=emb_dim, k_max=k_max)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec_pred = torch.tensor([0, 0, 0])
|
||||
sec_cont, sec_type, sec_valid = sample_secondaries_ar(
|
||||
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=True
|
||||
)
|
||||
assert not sec_valid.any()
|
||||
# every slot still ran the model (not left at the zero-init default) —
|
||||
# a real flow ODE output from randn-initialized noise is essentially
|
||||
# never exactly zero.
|
||||
assert not torch.allclose(sec_cont, torch.zeros_like(sec_cont))
|
||||
|
||||
|
||||
# ── resolve_n_sec: n_sec.mode = "head" sampling policy (gitea #86) ──────────
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user