lars
9277d79dff
- Data pipeline: loader (parquet→numpy), transforms (log, local-frame Rodrigues rotation, Normalizer), StepsDataset with event-ID-based split - Model: SinusoidalEmbedding, ConditionEncoder, ResBlock, DenoisingMLP - Schedule: cosine DDPM and conditional flow matching loss (Lipman 2022) - Samplers: flow (Euler ODE), DDPM ancestral, DDIM deterministic - Training loop: AdamW + cosine LR, grad clipping, best-val checkpoint - Validation: per-dimension marginal summary (normalised space) - CLI: TOML config support with CLI-overrides; hyperparam-encoded output directory; config.toml with git hash saved into each run's checkpoint dir - 21 unit tests covering transforms, network, flow/DDPM losses, dataset splits Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
68 lines
3.0 KiB
Markdown
68 lines
3.0 KiB
Markdown
# giant
|
||
|
||
**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 — replacing the stochastic Geant4 physics engine with a trained conditional generative model.
|
||
|
||
Training is driven entirely from parquet files of the miniCaloSim steps tree. No Geant4 runtime dependency.
|
||
|
||
## Architecture
|
||
|
||
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.
|
||
|
||
**Output space (6D, diffused):**
|
||
|
||
| Index | Variable | Transform |
|
||
|-------|----------|-----------|
|
||
| 0 | `step_length` [mm] | log |
|
||
| 1 | `ΔE = pre_E − post_E` [MeV] | log |
|
||
| 2 | `edep` [MeV] | log |
|
||
| 3–5 | `post_dir` in local frame | unit vector |
|
||
|
||
The post-step direction is expressed in the coordinate frame where `pre_dir = ẑ`, making the scattering distribution nearly azimuthally symmetric.
|
||
|
||
**Conditioning:** PDG code (embedding), pre-step position, log(pre-energy), pre-step direction, material (embedding), layer ID, number of secondaries.
|
||
|
||
## Roadmap
|
||
|
||
The model is developed in two phases:
|
||
|
||
**Phase 1 (current):** The number of secondaries produced in each step is passed as a conditioning input. This makes training easier because the model has direct access to multiplicity information and can focus on learning the continuous post-step kinematics.
|
||
|
||
**Phase 2 (target):** The number of secondaries is not given — the model must predict it jointly with all secondary properties (energy, direction, species) for each step. This requires extending the output space and likely an autoregressive or set-based generative approach for the variable-length secondary list.
|
||
|
||
## Data
|
||
|
||
Input: parquet files produced by [miniCaloSim](../minicalosim). 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
|
||
|
||
```
|
||
giant/
|
||
├── giant/
|
||
│ ├── data/
|
||
│ │ ├── loader.py # parquet → numpy arrays
|
||
│ │ ├── transforms.py # log transforms, local-frame rotation, normaliser
|
||
│ │ └── dataset.py # StepsDataset (PyTorch)
|
||
│ ├── model/
|
||
│ │ ├── network.py # SinusoidalEmbedding, ConditionEncoder, DenoisingMLP
|
||
│ │ └── schedule.py # CosineSchedule (DDPM) and flow matching utilities
|
||
│ ├── train.py # training loop and evaluation
|
||
│ ├── sample.py # DDPM / DDIM / flow matching samplers
|
||
│ └── validate.py # step-level and shower-level validation
|
||
└── scripts/
|
||
└── train.py # CLI entry point
|
||
```
|
||
|
||
## Setup
|
||
|
||
```bash
|
||
uv sync
|
||
```
|
||
|
||
## Training
|
||
|
||
```bash
|
||
python scripts/train.py --data path/to/steps.parquet --mode flow
|
||
```
|