# giant

**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep 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](../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
```