README_PROJECT.md

GNN4Colliders

GNN4Colliders is a collider-machine-learning toolkit. The repository name reflects its first production model family, ROOT-GNN; the Python package is gnn4colliders, and the configuration identifier is root_gnn. Shared ROOT ingestion, collider features, metadata, tasks, training, inference, and distributed utilities are designed so that a future sequence model can reuse them without requiring every event to be a graph.

ROOT files -> EventSample -> shared collider features
                              ├── GraphSample -> ROOT-GNN
                              └── future SequenceSample -> ROOT-Transformer

The new implementation lives under src/gnn4colliders. Historical behavior is preserved by the root-gnn-parity-baseline tag and committed reference fixtures, not by a supported historical runtime backend.

Installation

The supported development environment is Python 3.12 (>=3.12,<3.13). Core development is supported on macOS and Linux:

# macOS (Apple Silicon): CPU ROOT-GNN development and tests
uv sync --dev --extra root-gnn

# Linux x86_64 with an NVIDIA GPU: validated ROOT-GNN development
uv sync --dev --extra root-gnn

The core package can be installed without DGL when only shared data or task code is needed. ROOT-GNN models, graph construction, and ROOT-GNN reference tests require the root-gnn extra. On Linux x86_64, it uses the validated CUDA 12.1 wheels configured in pyproject.toml; a compatible NVIDIA driver is still required. On Apple Silicon macOS, it installs the CPU DGL wheel, supporting local graph/cache development. The default ROOT-GNN backend performs training with native PyTorch graph tensors, so it runs on Apple MPS, NVIDIA CUDA, and CPU; DGL remains a cache and graph compatibility adapter. Do not add site-specific CUDA, Slurm, or filesystem paths to model or task configuration.

Use the MPS profile on an Apple Silicon Mac:

uv run gnn4colliders train environment=macos

Data samples

ROOT inputs are available from the HWresearch/Delphes dataset. Download the 64-event smoke-test sample with the Hugging Face CLI:

hf download HWresearch/Delphes testing/ttH_NLO_64.root \
  --repo-type dataset --local-dir data/raw

The sample is data/raw/testing/ttH_NLO_64.root, has tree name output, and is suitable for checking the prepare/train workflow. The dataset also provides larger process-specific ROOT samples under samples/, derived datasets under derived/, and analysis-specific ntuples under analyses/. These data are intentionally ignored by Git; inspect a selected ROOT file's tree and branches before writing its preparation configuration.

Quick start

Prepare a graph cache from a ROOT tree. The feature specifications below are illustrative placeholders; replace them with the branches in the input tree. The full preparation interface is documented in docs/configuration.md.

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]

Train, evaluate, and predict from that cache:

uv run gnn4colliders train \
  data.cache.path=cache/events.pt \
  trainer.max_epochs=1 \
  environment.output_root=outputs/pretraining_multiclass

uv run gnn4colliders evaluate \
  data.cache.path=cache/events.pt \
  inference.checkpoint=outputs/pretraining_multiclass/checkpoints/epoch_0000.pt

uv run gnn4colliders predict \
  data.cache.path=cache/events.pt \
  inference.checkpoint=outputs/pretraining_multiclass/checkpoints/epoch_0000.pt \
  inference.output=outputs/pretraining_multiclass/predictions.npz

For a dependency-complete, temporary-data version of this flow, run uv run python scripts/dev/smoke_end_to_end.py.

Preparation can use local worker processes for larger inputs. Workers write ordered temporary shards and the application merges them into one cache:

uv run gnn4colliders prepare --config-name config_hf_smoke data.num_workers=4

Benchmark worker counts on the target machine with uv run python benchmarks/benchmark_prepare.py --workers 4; small fixtures may be slower because process startup dominates.

Core concepts

EventSample is the architecture-neutral event boundary. It contains the selected objects, label, global_features, and named EventMetadata. Metadata includes fold, weight, and stable sample_id; callers should not interpret public tracking[:, N] columns. Legacy tracking mappings exist only at compatibility boundaries.

The ROOT-GNN adapter converts shared features to a directed, fully connected graph with no self-loops: an event with N nodes has N * (N - 1) edges. Node columns are, in order, pt, eta, phi, energy, btag, charge, and node_type. Edge columns are deta, wrapped dphi, and dR. Object collections are concatenated in configured object-type order. The compatibility energy is pt * cosh(eta) before per-column scaling.

GraphSampleCache stores processed graph samples and schema metadata. It is a Level-2 graph cache, not the universal event cache. Feature, graph, and cache schema versions are checked when loading; incompatible versions fail before training.

ROOT-GNN training and transfer

EdgeNetwork encodes node, edge, and global features, performs iterative edge/node/global message passing, decodes a graph representation, and applies the classifier. Its output is raw logits; sigmoid or softmax is task-owned.

Multiclass pretraining uses the semantic model=root_gnn/edge_network and task=pretraining_multiclass groups:

uv run gnn4colliders train \
  data.cache.path=cache/events.pt \
  model=root_gnn/edge_network task=pretraining_multiclass \
  trainer.max_epochs=20 data.batch_size=64 \
  environment.output_root=outputs/pretraining_multiclass

Fine-tuning is a separate workflow. It loads a pretrained backbone, replaces the classifier, and creates a new task/head optimizer:

uv run gnn4colliders train \
  data.cache.path=cache/target.pt \
  model=root_gnn/fine_tuned_edge_network \
  task=binary_classification \
  checkpoint.pretrained=/path/to/pretrained.pt \
  model.freeze_backbone=true \
  trainer.max_epochs=10

Set model.freeze_backbone=false to train the reused backbone as well. Transfer learning is not resume training:

Workflow	Meaning	Restored state
Resume	Continue the same task/run	model, optimizer, scheduler, trainer, early stopping, and RNG state when present
Transfer	Start a new task from a pretrained backbone	model weights only; new classifier and optimizer

Resume example:

uv run gnn4colliders train \
  data.cache.path=cache/events.pt \
  checkpoint.resume=outputs/pretraining_multiclass/checkpoints/epoch_0000.pt \
  trainer.max_epochs=20

Validation is evaluated each epoch and drives scheduling/early stopping; test remains held out. Evaluation computes task metrics over the complete selected split, including weighted ROC AUC where defined:

uv run gnn4colliders evaluate \
  data.cache.path=cache/events.pt \
  inference.split=test \
  inference.checkpoint=/path/to/checkpoint.pt

Prediction writes a named compressed NPZ. Labeled data includes labels; fold and weight are included when available. Every result includes sample_id, logits, scores, and predictions:

uv run gnn4colliders predict \
  data.cache.path=cache/events.pt \
  inference.checkpoint=/path/to/checkpoint.pt \
  inference.output=outputs/predictions.npz

Optional Python-level ROOT writing is provided by gnn4colliders.inference.write_root_scores. It clones the selected tree, adds score (or score_class_N), and writes selection_pass; IDs ending in :<entry> preserve alignment and unselected entries receive NaN scores. The CLI currently exposes NPZ output only.

The supported legacy checkpoint, metadata, and output boundary is documented in docs/compatibility.md. New code should use named metadata fields; positional tracking is accepted only by the explicit compatibility adapter.

ONNX export

Install the optional export dependencies and export a prepared graph-cache checkpoint with numerical ONNX validation:

uv sync --extra root-gnn --extra onnx
uv run gnn4colliders export \
  export.checkpoint=/path/to/checkpoint.pt \
  export.output=model.onnx \
  data.cache.path=/path/to/graph-cache.pt

The model accepts processed graph tensors and returns raw logits. See docs/export.md for the tensor contract and limitations.

Configuration and environments

Hydra groups are data, model, task, trainer, checkpoint, inference, environment, and distributed. Use configuration for a new experiment and Python for new behavior. Examples:

uv run gnn4colliders train trainer.max_epochs=50 data.batch_size=64
uv run gnn4colliders train environment=perlmutter environment.device=cuda
uv run gnn4colliders train distributed=ddp environment=perlmutter

Each run writes a resolved configuration to <environment.output_root>/resolved_config.yaml. See docs/configuration.md for the group reference and docs/perlmutter.md for launch examples.

Distributed execution and reproducibility

Launch DDP with torchrun or the provided Slurm wrappers. data.batch_size and data.num_workers are per process, so the ordinary effective batch size is batch_size * world_size. Training shards may be padded for equal steps; validation and prediction are unpadded. Rank 0 writes shared checkpoints, configs, and predictions, and metrics/results are gathered across ranks.

The configured seed controls initialization and deterministic local loader ordering; distributed process seeds are rank-offset and samplers use set_epoch. CPU runs are reproducible for fixed inputs and environment. GPU kernels, DGL, and distributed scheduling can remain nondeterministic, so the project does not promise bitwise GPU reproducibility.

Development and validation

uv run pytest
uv run pytest tests/unit
GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 uv run pytest tests/parity -v
uv run ruff check .
uv run ruff format --check .
uv run python benchmarks/benchmark_preprocessing.py
uv run python benchmarks/benchmark_training.py --device cpu

Unit tests cover isolated components, integration tests cover small workflows, and parity tests compare deterministic behavior with the frozen legacy reference. Performance guidance and measured caveats are in docs/performance.md and benchmarks/README.md. See docs/testing.md for test layers, optional dependency markers, and package smoke validation.

Architecture and migration status

See docs/architecture.md for responsibility boundaries and the future sequence-model extension point. See docs/migration.md for the migration matrix, intentional redesigns, compatibility limits, and deferred work.

ROOT-GNN v1 covers ROOT preparation, validated feature/graph/model/task behavior, training, fine-tuning, checkpoint resume, evaluation, prediction, single-process/DDP execution, and validated ONNX export. Streaming distributed output, legacy cleanup, and ROOT-Transformer remain follow-up work.