commit 943a044608ef3301421ea055ad4f77d3b5d49aec Author: Lars Bogner Date: Wed Jun 17 09:29:19 2026 +0200 Initial commit: giant surrogate model with two-phase roadmap in README Co-Authored-By: Claude Sonnet 4.6 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..505a3b1 --- /dev/null +++ b/.gitignore @@ -0,0 +1,10 @@ +# Python-generated files +__pycache__/ +*.py[oc] +build/ +dist/ +wheels/ +*.egg-info + +# Virtual environments +.venv diff --git a/.python-version b/.python-version new file mode 100644 index 0000000..6324d40 --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.14 diff --git a/README.md b/README.md new file mode 100644 index 0000000..f965886 --- /dev/null +++ b/README.md @@ -0,0 +1,67 @@ +# 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 +``` diff --git a/giant/__init__.py b/giant/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/giant/data/__init__.py b/giant/data/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/giant/data/dataset.py b/giant/data/dataset.py new file mode 100644 index 0000000..8e06954 --- /dev/null +++ b/giant/data/dataset.py @@ -0,0 +1 @@ +# PyTorch Dataset wrapping the preprocessed steps arrays. diff --git a/giant/data/loader.py b/giant/data/loader.py new file mode 100644 index 0000000..2e6bf49 --- /dev/null +++ b/giant/data/loader.py @@ -0,0 +1 @@ +# Load steps parquet files into numpy arrays. diff --git a/giant/data/transforms.py b/giant/data/transforms.py new file mode 100644 index 0000000..818e313 --- /dev/null +++ b/giant/data/transforms.py @@ -0,0 +1 @@ +# Log transforms, local-frame direction rotation, and per-dimension normaliser. diff --git a/giant/model/__init__.py b/giant/model/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/giant/model/network.py b/giant/model/network.py new file mode 100644 index 0000000..3ac58df --- /dev/null +++ b/giant/model/network.py @@ -0,0 +1 @@ +# SinusoidalEmbedding, ConditionEncoder, ResBlock, DenoisingMLP. diff --git a/giant/model/schedule.py b/giant/model/schedule.py new file mode 100644 index 0000000..2eddceb --- /dev/null +++ b/giant/model/schedule.py @@ -0,0 +1 @@ +# CosineSchedule (DDPM forward process) and flow matching loss utilities. diff --git a/giant/sample.py b/giant/sample.py new file mode 100644 index 0000000..4edb8ab --- /dev/null +++ b/giant/sample.py @@ -0,0 +1 @@ +# DDPM, DDIM, and flow matching samplers. diff --git a/giant/train.py b/giant/train.py new file mode 100644 index 0000000..218eb73 --- /dev/null +++ b/giant/train.py @@ -0,0 +1 @@ +# Training loop and validation loss evaluation. diff --git a/giant/validate.py b/giant/validate.py new file mode 100644 index 0000000..2cb9afe --- /dev/null +++ b/giant/validate.py @@ -0,0 +1 @@ +# Step-level marginal comparisons and (later) shower-level rollout validation. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..8c4bd3a --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "giant" +version = "0.1.0" +description = "Geant4 step-function surrogate via conditional flow matching" +readme = "README.md" +requires-python = ">=3.12" +dependencies = [ + "torch>=2.3", + "numpy>=1.26", + "pandas>=2.2", + "pyarrow>=16", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["giant"] diff --git a/scripts/train.py b/scripts/train.py new file mode 100644 index 0000000..04c24bf --- /dev/null +++ b/scripts/train.py @@ -0,0 +1 @@ +# CLI entry point: parse args, build dataset, instantiate model, call train loop.