2.9 KiB
giant
Geant4 Inference via Autoregressive Neural sTep surrogate.
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. 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
uv sync
Training
python scripts/train.py --data path/to/steps.parquet --mode flow