Update README to match current architecture and tooling

Output space is 9D (post_dir + travel_dir), not 6D; documents the
giant CLI, analysis/validate modules, ROOT-to-parquet conversion
script, cpu/cuda install extras, and the ruff/ty dev tooling.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-06-18 17:43:04 +02:00
parent 8cebc4809d
commit 0cdb5db947
+43 -11
View File
@@ -10,7 +10,7 @@ Training is driven entirely from parquet files of the miniCaloSim steps tree. No
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):**
**Output space (9D, diffused):**
| Index | Variable | Transform |
|-------|----------|-----------|
@@ -18,10 +18,11 @@ Conditional **flow matching** model (Lipman et al. 2022): a small MLP learns a v
| 1 | `ΔE = pre_E post_E` [MeV] | log |
| 2 | `edep` [MeV] | log |
| 35 | `post_dir` in local frame | unit vector |
| 68 | `travel_dir` (`post_pos pre_pos`) 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.
Both `post_dir` and `travel_dir` are expressed in the coordinate frame where `pre_dir = ẑ`, making the scattering distribution nearly azimuthally symmetric. `post_pos` itself is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, so the two stay consistent by construction instead of being learned (and potentially diverging) independently.
**Conditioning:** PDG code (embedding), pre-step position, log(pre-energy), pre-step direction, material (embedding), layer ID, number of secondaries.
**Conditioning:** PDG code (embedding), pre-step position, log(pre-energy), pre-step direction, material (embedding), layer ID, number of secondaries (Phase 1 only — see Roadmap below).
## Roadmap
@@ -33,7 +34,7 @@ The model is developed in two phases:
## 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.
Input: parquet files produced by [miniCaloSim](../minicalosim), or converted from a ROOT file via `scripts/steps_to_parquet.py`. 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
@@ -41,27 +42,58 @@ Input: parquet files produced by [miniCaloSim](../minicalosim). Each row is one
giant/
├── giant/
│ ├── data/
│ │ ├── loader.py # parquet → numpy arrays
│ │ ├── loader.py # parquet → numpy arrays (incl. streaming/chunked reads)
│ │ ├── transforms.py # log transforms, local-frame rotation, normaliser
│ │ └── dataset.py # StepsDataset (PyTorch)
│ │ └── dataset.py # StepsDataset / StreamingStepsDataset (PyTorch)
│ ├── model/
│ │ ├── network.py # SinusoidalEmbedding, ConditionEncoder, DenoisingMLP
│ │ └── schedule.py # CosineSchedule (DDPM) and flow matching utilities
│ ├── train.py # training loop and evaluation
│ ├── config.py # default hyperparameters, TOML config merging, device autodetect
│ ├── pipeline.py # builds datasets/normalizers and kicks off a training run
│ ├── train.py # training loop, checkpointing, graceful shutdown
│ ├── sample.py # DDPM / DDIM / flow matching samplers
── validate.py # step-level and shower-level validation
└── scripts/
└── train.py # CLI entry point
── validate.py # step-level marginal + KL-divergence validation
│ ├── analysis.py # notebook diagnostics: marginals, correlations, constraint checks
└── cli.py # `giant train` / `giant predict` Typer app
├── scripts/
│ ├── train.py # argparse training entry point
│ └── steps_to_parquet.py # ROOT → parquet conversion (uproot/awkward/polars)
└── tests/
```
## Setup
```bash
uv sync
uv sync --extra cpu # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
uv sync --extra cpu --extra dev # add dev tools (pytest, ruff, ty)
```
`cpu` and `cuda` are mutually exclusive extras selecting the torch build; plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
## Training
```bash
python scripts/train.py --data path/to/steps.parquet --mode flow
```
or via the installed CLI:
```bash
giant train path/to/steps.parquet --mode flow
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
```
Both accept a TOML config file (`--config`) and CLI overrides for hyperparameters; see `--help` on either for the full option list.
## Validation
`giant.validate.validate_marginals` runs step-level marginal and KL-divergence checks during training (`--validate-every`). For deeper, notebook-driven diagnostics on a trained checkpoint — stratified marginals, correlation structure, physical constraint violations — see `giant.analysis`.
## Development
```bash
uv run pytest # run tests
uv run ruff check . # lint
uv run ruff format . # format
uv run ty check . # type check
```