From 0cdb5db94745216ca765ae306ff84322753fbb81 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 18 Jun 2026 17:43:04 +0200 Subject: [PATCH] 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 --- README.md | 54 +++++++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 43 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 8a1357a..e0d0c19 100644 --- a/README.md +++ b/README.md @@ -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 | | 3–5 | `post_dir` in local frame | unit vector | +| 6–8 | `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 +```