GNN4Colliders / docs /quickstart.md
ho22joshua's picture
feat: add ttH CP parity benchmark configuration
d2d4ccf
|
Raw
History Blame Contribute Delete
5.83 kB
# GNN4Colliders quickstart
GNN4Colliders turns collider ROOT events into shared event samples, applies
physics features, builds a representation for a selected model family, and
runs training or inference through one CLI. The current production model
family is `root_gnn`; the shared data and task layers are intended to support
future representations as well.
```text
ROOT -> EventSample -> features -> GraphSampleCache -> ROOT-GNN -> logits
```
## Install
The supported environment is Python 3.12. From the repository root:
```bash
uv sync --dev --extra root-gnn
```
The `root-gnn` extra is required for DGL-backed graph workflows and ROOT-GNN
parity tests. Add `--extra onnx` when using export validation.
## Run a local end-to-end smoke test
This is the fastest way to verify the package without downloading production
data:
```bash
uv run python scripts/dev/smoke_end_to_end.py
```
The script creates a tiny temporary ROOT file, prepares a cache, trains one
CPU epoch, evaluates, predicts, and cleans up. It is the recommended first
command after installation or a code change.
## Run the public Hugging Face smoke configuration
The small pinned fixture is selected by `config_hf_smoke`:
```bash
uv run gnn4colliders prepare --config-name config_hf_smoke
uv run gnn4colliders train --config-name config_hf_smoke
```
The first command resolves and caches the public ROOT fixture. The default
debug trainer runs one epoch. Outputs are written below `outputs/hf_smoke/`;
the resolved configuration and checkpoint path are recorded there.
For the full 12-process pretraining source, use this only when the larger
download is intentional:
```bash
uv run gnn4colliders prepare --config-name config_hf_delphes
```
The data source, revision, checksum, tree, feature branches, and split rules
are documented in [configuration.md](configuration.md).
The inclusive no-selection ttH CP benchmark has a dedicated configuration. It
uses `ttH_NLO.root` as label 0 (CP-even) and `ttH_CPodd.root` as label 1
(CP-odd), with all configured object branches passed through without event or
object selections:
```bash
uv run gnn4colliders prepare \
--config-name config_tth_cp_even_odd \
data.num_workers=8
```
This task configuration enables absolute event weights for optimization
because the NLO files contain signed weights. The original signed weights are
still retained in event metadata and used by the configured evaluation
semantics.
## Use your own ROOT file
Preparation accepts Hydra overrides. Supply the ROOT file, tree, cache path,
feature branch specification, object types, and scales:
```bash
uv run gnn4colliders prepare \
data.files=[data/events.root] \
data.tree_name=Events \
data.cache.path=cache/events.pt \
'data.feature_branches=[[jet_pt],[jet_eta],[jet_phi],CALC_E,[1.0],[0.0],NODE_TYPE]' \
data.object_types=[vector] \
data.scales=[1,1,1,1,1,1,1]
```
Then train, evaluate, or predict using the same cache:
```bash
uv run gnn4colliders train data.cache.path=cache/events.pt trainer.max_epochs=1
uv run gnn4colliders evaluate \
data.cache.path=cache/events.pt \
inference.checkpoint=outputs/default/checkpoints/epoch_0000.pt
uv run gnn4colliders predict \
data.cache.path=cache/events.pt \
inference.checkpoint=outputs/default/checkpoints/epoch_0000.pt \
inference.output=outputs/predictions.npz
```
The feature override above is illustrative: replace it with branches that
exist in the input tree. See [configuration.md](configuration.md) for the
full schema.
## Models, tasks, and checkpoints
Select behavior with semantic config groups rather than Python module paths:
```bash
uv run gnn4colliders train \
model=root_gnn/edge_network \
task=pretraining_multiclass \
data.cache.path=cache/events.pt \
trainer.max_epochs=20
```
To fine-tune a pretrained ROOT-GNN checkpoint, use the fine-tuning model and
provide `checkpoint.pretrained`:
```bash
uv run gnn4colliders train \
model=root_gnn/fine_tuned_edge_network \
task=binary_classification \
data.cache.path=cache/target.pt \
checkpoint.pretrained=/path/to/pretrained.pt \
model.freeze_backbone=true
```
`checkpoint.resume` continues an existing run, restoring its training state;
`checkpoint.pretrained` starts a new task from model weights and creates a new
head and optimizer. Evaluation uses the held-out test split by default.
The first transformer model reuses node features as ordered tokens and the
same task/trainer lifecycle. It can be trained from an existing prepared cache:
```bash
uv run gnn4colliders train \
model=root_transformer/transformer \
task=binary_classification \
data.cache.path=cache/target.pt \
model.hid_size=64 model.n_heads=4
```
This bridge intentionally ignores graph edges. Native sequence-cache storage
and distributed sequence loading are follow-up work.
## Validate changes
Run focused tests first, then the full suite and style checks:
```bash
uv run pytest tests/unit -q
GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 uv run pytest tests/parity -q
uv run pytest -q
uv run ruff check .
uv run ruff format --check .
```
For the portable public-data regression:
```bash
uv run python -m validation.run_public_validation --no-download
```
The parity gate requires DGL and compares the new implementation against
committed frozen reference fixtures. Historical runtime code is not imported
from main; compatibility adapters are reserved for old checkpoint, metadata,
and output boundaries.
## Where to look next
- [Architecture](architecture.md) — package boundaries and data flow.
- [Configuration](configuration.md) — Hydra groups and data sources.
- [Migration](migration.md) — rewrite status and compatibility policy.
- [Developer guide](../README_PROJECT.md) — transfer learning, distributed
training, ONNX export, and output contracts.