lars
fc19934ba6
b2luigi's AnalysisComputeTask now submits the per-(plot, chunk) jobs, so the bespoke submit-file generator has nothing left to do: - giant/analysis/condor.py -> giant/analysis/run.py, dropping SubmitConfig, the wrapper/submit-description templates, _job_walltimes and _resolve_giant_executable. What stays is the actual logic — prep, RunMeta, the rollout-YAML loading, compute_reduced/compute_one and merge_one/merge_all — and the module no longer submits anything, hence the name. - `giant analyze submit` is gone; prep / compute-one / merge-one / list / render / metrics remain as the single-step primitives the workflow calls. - tests/test_condor.py -> tests/test_analysis_run.py, minus the submit-description cases. CLAUDE.md and README.md document the workflow package, the new `workflow` extra, and — for whenever condor-gpu-train-rollout is merged — that its train-submit/rollout-submit commands are deliberately superseded and must not be revived. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
222 lines
19 KiB
Markdown
222 lines
19 KiB
Markdown
# giant
|
||
|
||
**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate.
|
||
|
||
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.
|
||
|
||
## 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
|
||
|
||
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
|
||
```
|
||
|
||
Every command takes `--help` for the 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).
|
||
|
||
**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 |
|
||
|
||
- 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)`.
|
||
|
||
**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`):
|
||
|
||
- `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`
|
||
|
||
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).
|
||
|
||
**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.
|
||
|
||
**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.
|
||
|
||
## Data
|
||
|
||
- 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.
|
||
|
||
## Project structure
|
||
|
||
```
|
||
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 the job requests
|
||
│ │ ├── run.py # prep / compute-one / merge plumbing
|
||
│ │ └── render.py # PDFs + HTML gallery (only module importing plotstyle/LaTeX)
|
||
│ ├── workflow/ # b2luigi pipeline orchestration (`giant workflow run spec.toml`)
|
||
│ │ ├── spec.py # workflow TOML -> frozen dataclasses, validation, per-task spec hashes
|
||
│ │ ├── htcondor.py # CPU/GPU submit settings (docker image, +RemoteJob, GPU requirements)
|
||
│ │ ├── tasks.py # the task graph: cache-warm -> train (one job/epoch) -> rollout -> analysis
|
||
│ │ └── run.py # the script b2luigi re-executes on every worker
|
||
│ └── cli.py # `giant train` / `new-run` / `model summary` / `predict` / `rollout` / `analyze` / `workflow`
|
||
├── 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/
|
||
```
|
||
|
||
## Setup
|
||
|
||
```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 workflow # b2luigi, for `giant workflow run`
|
||
```
|
||
|
||
The `dev` extra pulls in `convert`, `analysis`, `geometry`, `wandb` and `workflow` as well.
|
||
|
||
`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.
|
||
|
||
## Training, prediction, rollout
|
||
|
||
```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 prep rollout.yaml --chunks 8 # lay out the run directory
|
||
giant analyze prep a.yaml b.yaml --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 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)
|
||
```
|
||
|
||
The per-(plot, chunk) compute jobs themselves are submitted by the workflow (below), not by `giant analyze` — these commands are the single-step primitives it calls. `<run_dir>` defaults to `<cwd>/analysis_runs/analysis_<id>` (`--run-dir` overrides it; `prep` prints 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.
|
||
|
||
## Workflow orchestration
|
||
|
||
Multi-step pipelines run through [b2luigi](https://github.com/belle2/b2luigi) — one spec file describes a whole experiment, and every step's outputs are files on `/ceph` that are only recomputed when their spec (or an upstream one) changes:
|
||
|
||
```bash
|
||
uv sync --extra cpu --extra workflow
|
||
giant workflow run configs/workflow_example.toml --mode dry-run # what would run
|
||
giant workflow run configs/workflow_example.toml --mode show-output # where every output goes
|
||
giant workflow run configs/workflow_example.toml --batch --workers 20 # submit to HTCondor and wait
|
||
```
|
||
|
||
The spec holds `[workflow]`/`[condor]`/`[dataset]`/`[geometry]` plus repeated `[[train]]`, `[[rollout]]` and `[[analysis]]` tables cross-referenced by name (see `configs/workflow_example.toml`). The task graph is `DatasetTask → WarmCacheTask/GeometryOracleTask → TrainEpochTask… → TrainTask → RolloutTask → AnalysisPrepTask → AnalysisComputeTask(plot, chunk) → AnalysisRenderTask`. Training is split into one short GPU job per epoch (chained by `--resume`), which schedules better on a busy farm and survives preemption; `TrainTask` then publishes one `best.pt`/`last.pt`/`metrics.csv` for everything downstream. Rendering always runs locally, since it is the only step that needs LaTeX.
|
||
|
||
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 ruff check . # lint
|
||
uv run ruff format . # format
|
||
uv run ty check . # type check
|
||
```
|