lars
a867fc4aae
Provides stratified marginal comparisons, joint-structure checks (correlation matrices, physically-coupled pairwise plots, direction alignment), and physical-constraint validation (unit-norm directions, non-negative raw targets) for a trained model's generated samples, building on the aggregate marginal/KL check already in giant.validate. Supports two entry points: live sampling against a checkpoint + val data (load_model_bundle/collect_samples), or loading a precomputed `giant predict --coord local` parquet directly (load_predicted_local) without needing the checkpoint at all. Predict output is now tagged with parquet schema metadata so the loader can verify a file's format and reject coord=global or untagged files with a clear error instead of guessing from column names. Also extends the config git-hash mismatch warning (added for --config loading) to checkpoint loading: both `giant predict` and analysis.load_model_bundle now look for a config.toml next to the checkpoint and warn (without failing) if it was generated from a different git commit. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
giant
Geant4 Inference via Autoregressive Neural sTep 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. 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