File size: 5,830 Bytes
70ee5cc d2d4ccf 70ee5cc 8344837 70ee5cc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | # 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.
|