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.

ROOT -> EventSample -> features -> GraphSampleCache -> ROOT-GNN -> logits

Install

The supported environment is Python 3.12. From the repository root:

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:

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:

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:

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.

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:

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:

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:

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 for the full schema.

Models, tasks, and checkpoints

Select behavior with semantic config groups rather than Python module paths:

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:

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:

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:

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:

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 — package boundaries and data flow.
Configuration — Hydra groups and data sources.
Migration — rewrite status and compatibility policy.
Developer guide — transfer learning, distributed training, ONNX export, and output contracts.