docs: bring README in line with current architecture and CLI
Adds the physical/embedding conditioning split, WGAN/MoE-routing modes, HTCondor remote-GPU submission, and the giant analyze pipeline, none of which were reflected in the previous version. Updates the project structure listing to match the current module layout.
This commit is contained in:
@@ -2,15 +2,19 @@
|
||||
|
||||
**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate — a play on *Geant4* and the step function being the computationally heaviest part of the simulation.
|
||||
|
||||
Proof-of-concept surrogate model for the Geant4 step function. Given a pre-step particle state, the model samples a physically plausible post-step outcome — the primary's continuation plus the secondary particles it produces — replacing the stochastic Geant4 physics engine with a trained conditional generative model.
|
||||
Conditional generative surrogate for the Geant4 step function. Given a pre-step particle state, the model samples a physically plausible post-step outcome — the primary's continuation plus the variable-length list of secondary particles it produces — replacing the stochastic Geant4 physics engine with a trained generative model. A trained checkpoint autoregressively rolls out full showers, stepping each primary and pushing secondaries as new tracks.
|
||||
|
||||
Training is driven entirely from parquet files of the miniCaloSim steps tree. No Geant4 runtime dependency.
|
||||
|
||||
## Architecture
|
||||
|
||||
A **two-stage conditional flow matching** model (Lipman et al. 2022): a small MLP learns a vector field mapping noise → step outcomes in ~10 ODE steps per sample. Falls back to DDPM for comparison.
|
||||
A **two-stage model**, both stages checkpointed together, with a choice of generative mode per stage (`--mode`):
|
||||
|
||||
**Stage 1 — primary (9D, diffused):**
|
||||
- **`flow`** (default) — conditional flow matching (Lipman et al. 2022): an MLP learns a vector field mapping noise → step outcomes, sampled via ODE integration in ~10 steps.
|
||||
- **`ddpm`** — a standard denoising diffusion baseline for comparison (`giant/model/schedule.py:CosineSchedule`).
|
||||
- **`wgan`** — a single-pass Wasserstein-GAN-GP generator/critic (`giant/model/wgan.py`), trading iterative sampling for one forward pass, evaluated against a ~10× native-Geant4 latency budget.
|
||||
|
||||
**Stage 1 — primary (9D, `giant/constants.py:LOCAL_TARGET_NAMES`):**
|
||||
|
||||
| Index | Variable | Encoding |
|
||||
|-------|----------|----------|
|
||||
@@ -19,13 +23,20 @@ A **two-stage conditional flow matching** model (Lipman et al. 2022): a small ML
|
||||
| 3–5 | `post_dir` in local frame | unit vector |
|
||||
| 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector |
|
||||
|
||||
The two energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — **energy conservation is built into the parametrization**, not left to the loss. Stage 1 also has a classifier head predicting the number of secondaries `n_sec ∈ {0..15}` from the conditioning alone.
|
||||
The two energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — **energy conservation is built into the parametrization**, not left to the loss (`energy_simplex_decode`). Stage 1 also has a classifier head (`predict_n_sec`) predicting the number of secondaries `n_sec ∈ {0..K_MAX}` (`K_MAX = 15`) from the conditioning alone, no diffusion noise involved.
|
||||
|
||||
Both `post_dir` and `travel_dir` are expressed in the coordinate frame where `pre_dir = ẑ`, making the scattering distribution nearly azimuthally symmetric. `post_pos` itself is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, so the two stay consistent by construction instead of being learned (and potentially diverging) independently.
|
||||
|
||||
**Stage 2 — secondaries (`SecondaryDecoder`):** conditioned on the pre-step state *and* the Stage-1 outcome, a second flow net generates all `K_MAX = 15` secondary slots at once. Each slot carries a stick-breaking energy fraction, a local-frame direction, and a continuous particle-type embedding (snapped to the nearest PDG at inference), ordered by descending energy; slots beyond the predicted `n_sec` are masked. The secondary energies are a stick-breaking partition of the `e_sec` budget from Stage 1, so the full chain conserves energy. Each secondary's momentum is reconstructed afterward from `(energy, direction, species)` rather than predicted.
|
||||
**Stage 2 — secondaries (`SecondaryDecoder`):** conditioned on the pre-step state *and* the Stage-1 outcome, a second net generates all `K_MAX` secondary slots at once — `(stick-breaking energy logit, local-frame direction, log-mass, charge)` per slot, ordered by descending energy; slots beyond the predicted `n_sec` are masked. Secondary energies are a stick-breaking partition of the `e_sec` budget from Stage 1, so the whole chain conserves energy. A secondary's mass/charge are regressed directly against its ground-truth PDG code's physical values (`giant.particles.particle_mass_charge`) and used as-is at inference — including for its own conditioning if it takes further steps in a rollout. No snapping to a known PDG code happens in the model path; `giant.particles.nearest_known_pdg` is a reporting-only lookup used to populate a nominal `pdg` label on output rows.
|
||||
|
||||
**Conditioning:** PDG code (embedding), pre-step position, log(pre-energy), pre-step direction, material (embedding), layer ID. (`n_sec` / `e_sec` are outputs now, not inputs.)
|
||||
**Conditioning (`--conditioning`, per-checkpoint):** pre-step position, log(pre-energy), pre-step direction, layer ID, plus particle/material physical properties — mass/charge (`giant/particles.py`) and Z_eff/A_eff/density/X0/λ_int (`giant/materials.py`). Two mutually exclusive modes:
|
||||
|
||||
- **`physical`** (default) — the physical-property columns are routed through small MLPs, computable for any PDG code / material, letting the surrogate generalize to species/materials outside the training menu.
|
||||
- **`embedding`** — the original design: a learned `nn.Embedding` per PDG code / material, kept as a generalization-comparison baseline (memorizes the training menu).
|
||||
|
||||
`n_sec` and `e_sec` are model outputs, not conditioning inputs — a rollout is self-contained and never injects ground truth.
|
||||
|
||||
**Mixture-of-experts routing (`--router`, opt-in):** `giant/model/network.py` also implements a pluggable `Router` contract (`ROUTER_REGISTRY`: `energy`, `pdg`, `process`, plus a `composed` router combining several axes) that splits `DenoisingMLP`/`SecondaryDecoder` into per-expert trunks (`RoutedDenoisingMLP`/`RoutedSecondaryDecoder`), soft-gated in training and hard-dispatched at eval. See `--router-type`/`--n-experts`/`--router-axis` on `giant train`/`giant new-run`.
|
||||
|
||||
## Roadmap
|
||||
|
||||
@@ -33,11 +44,13 @@ Both `post_dir` and `travel_dir` are expressed in the coordinate frame where `pr
|
||||
|
||||
**Phase 2 (implemented — baseline):** the two-stage model above predicts `n_sec` and each secondary's energy, direction, and species jointly with the primary, so a shower rollout is fully self-contained.
|
||||
|
||||
**Next directions:** faster-eval architectures against a ~10× native-Geant4 budget (Wasserstein-GAN, mixture-of-experts routing tree), a multi-material sampling-calorimeter dataset, and physical-property conditioning over learned embeddings.
|
||||
**Physical-property conditioning (implemented):** replaces learned PDG/material embeddings with physical-property MLPs (see above); Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. Not yet done: the held-out-material/species generalization comparison against the `embedding` baseline.
|
||||
|
||||
**In progress:** a WGAN-GP single-pass mode (`--mode wgan`) and mixture-of-experts routing (`--router`) are implemented and under evaluation as faster-eval alternatives to iterative flow/DDPM sampling, against a ~10× native-Geant4 latency budget. A multi-material sampling-calorimeter dataset is the target for both the routing/WGAN evaluation and the physical-property generalization comparison.
|
||||
|
||||
## Data
|
||||
|
||||
Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from a ROOT file via `uv run dwarf convert`. Each row is one Geant4 step. Train/val split is by `event_id` (not row shuffle) to avoid leaking correlated steps from the same shower.
|
||||
Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from a ROOT file via `dwarf convert`. Each row is one Geant4 step. Train/val split is by `event_id` (not row shuffle) to avoid leaking correlated steps from the same shower.
|
||||
|
||||
## Project structure
|
||||
|
||||
@@ -49,21 +62,33 @@ giant/
|
||||
│ │ ├── transforms.py # log transforms, local-frame rotation, energy simplex, secondary encode/decode
|
||||
│ │ └── dataset.py # StepsDataset / StreamingStepsDataset (PyTorch)
|
||||
│ ├── model/
|
||||
│ │ ├── network.py # SinusoidalEmbedding, ConditionEncoder, DenoisingMLP, SecondaryDecoder
|
||||
│ │ └── schedule.py # CosineSchedule (DDPM) and flow matching utilities
|
||||
│ │ ├── network.py # ConditionEncoder, DenoisingMLP, SecondaryDecoder, Router/MoE, WGAN generator/critic
|
||||
│ │ ├── schedule.py # CosineSchedule (DDPM) and flow matching utilities
|
||||
│ │ └── wgan.py # WGAN-GP gradient penalty / critic / generator losses
|
||||
│ ├── constants.py # output/conditioning dims, K_MAX, secondary slot layout, schema keys
|
||||
│ ├── particles.py # PDG → (mass, charge) decode, incl. nuclear/ion codes; nearest-known-PDG lookup
|
||||
│ ├── 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
|
||||
│ ├── train.py # two-stage training loop, checkpointing, graceful shutdown
|
||||
│ ├── sample.py # DDPM / DDIM / flow matching samplers + secondary sampling
|
||||
│ ├── sample.py # DDPM / DDIM / flow matching / WGAN samplers + secondary sampling
|
||||
│ ├── geometry.py # GeometryOracle: position → (material, layer_id, escaped) for rollout
|
||||
│ ├── rollout.py # autoregressive shower rollout driver
|
||||
│ ├── validate.py # step-level marginal + KL-divergence validation
|
||||
│ ├── analysis.py # step- and shower-level diagnostics: marginals, correlations, rollout observables
|
||||
│ └── cli.py # `giant train` / `predict` / `rollout` Typer app
|
||||
├── scripts/ # dataset/tooling logic, unified under the `dwarf` CLI (`uv run dwarf --help`)
|
||||
│ ├── condor.py # HTCondor submission for `train-submit`/`rollout-submit` (remote GPU workers)
|
||||
│ ├── analysis/ # rollout-vs-reference analysis pipeline (see `giant analyze` below)
|
||||
│ │ ├── sources.py # canonical LazyFrames + secondary view
|
||||
│ │ ├── 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
|
||||
│ │ ├── catalog.py # declarative PlotSpec registry
|
||||
│ │ ├── condor.py # prep / compute-one / submit-description plumbing
|
||||
│ │ └── render.py # PDFs + HTML gallery (only module importing plotstyle/LaTeX)
|
||||
│ └── cli.py # `giant train` / `predict` / `rollout` / `new-run` / `*-submit` / `analyze` Typer app
|
||||
├── scripts/ # 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, hparam-scan
|
||||
│ │ # update-manifest, create-manifest, make-root, build-geometry-oracle,
|
||||
│ │ # 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`
|
||||
@@ -80,27 +105,48 @@ giant/
|
||||
```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
|
||||
```
|
||||
|
||||
`cpu` and `cuda` are mutually exclusive extras selecting the torch build; plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
|
||||
`cpu` and `cuda` are mutually exclusive extras selecting the torch build (pinned to 2.3.x); plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
|
||||
|
||||
## Training, prediction, and rollout
|
||||
|
||||
```bash
|
||||
giant train path/to/steps.parquet --mode flow
|
||||
giant new-run --hidden-dim 512 --lr 3e-4 --comment "..." # scaffold a config.toml + run dir for a new run
|
||||
giant train path/to/steps.parquet --mode flow # train (flow matching; also --mode ddpm / wgan)
|
||||
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
|
||||
|
||||
# Full-shower rollout needs a geometry oracle (position → material/layer_id):
|
||||
uv sync --extra cpu --extra geometry
|
||||
dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl
|
||||
giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
|
||||
```
|
||||
|
||||
`train`/`predict` accept a TOML config file (`--config`) and CLI overrides for hyperparameters; see `--help` on any command for the full option list. `giant rollout` seeds showers from the highest-energy entry step per event, then autoregressively steps the two-stage model to completion — pushing secondaries as new tracks and looking up `material`/`layer_id` from the oracle at each step. Tracks terminate on energy cutoff, per-track max steps, detector escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
|
||||
|
||||
## Validation
|
||||
### Remote-GPU training and rollout (HTCondor)
|
||||
|
||||
`giant.validate.validate_marginals` runs step-level marginal and KL-divergence checks during training (`--validate-every`). For deeper diagnostics on a trained checkpoint — stratified marginals, correlation structure, physical-constraint violations, and shower-level rollout observables (longitudinal/transverse profiles, PDG energy shares) — see `giant.analysis`.
|
||||
`giant train-submit`/`giant rollout-submit` run a `giant train`/`giant rollout` invocation as a remote-GPU HTCondor job (ETP's TOpAS V100/A100 or NEMO2 L40S workers). GPU workers are remote-only and reach `/ceph` (not `/work`/`/home`) — the checkout submitting these jobs, and its venv, must live under `/ceph`.
|
||||
|
||||
```bash
|
||||
giant train-submit path/to/steps.parquet --config run/config.toml --accounting-group cms
|
||||
giant rollout-submit path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl --accounting-group cms
|
||||
```
|
||||
|
||||
Each submission writes a self-contained `condor/`/`condor_<tag>/` directory (`run.sh`, `job.sub`, `submission.json`). Training jobs re-check `out_dir/last.pt` on every invocation and resume automatically if preempted. `train-submit` requires `--config` (rather than mirroring `train`'s hyperparameter flags) since the TOML is the reproducible source of truth, and passing the same one on every retry is what keeps a resumed run's architecture from drifting.
|
||||
|
||||
## Validation and analysis
|
||||
|
||||
`giant.validate.validate_marginals` runs step-level marginal and KL-divergence checks during training (`--validate-every`).
|
||||
|
||||
For deeper rollout-vs-reference diagnostics — marginals stratified by energy/pdg/material, per-event totals, shower profiles, species share, leakage, and secondaries — `giant analyze` runs a streaming compute/render pipeline against a `giant rollout` YAML sidecar:
|
||||
|
||||
```bash
|
||||
giant analyze submit rollout.yaml --accounting-group cms # prep + one HTCondor job per plot (compute only)
|
||||
giant analyze render <run_dir> --gallery # local: styled PDFs + HTML gallery (needs LaTeX)
|
||||
```
|
||||
|
||||
`<run_dir>` is derived next to the rollout parquet (`analyze prep`/`submit` print it). Compute jobs are polars/numpy only; only `render` imports plotstyle/LaTeX, so it always runs locally.
|
||||
|
||||
## Development
|
||||
|
||||
|
||||
Reference in New Issue
Block a user