docs: add Hugging Face model card

README.md

@@ -1,282 +1,246 @@
-# 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.
-
-```text
-ROOT files -> EventSample -> shared collider features
-                              ├── GraphSample -> ROOT-GNN
-                              └── future SequenceSample -> ROOT-Transformer
-```
-
-The new implementation lives under [`src/gnn4colliders`](src/gnn4colliders/).
-[`legacy/`](legacy/) is a frozen behavioral reference for parity work and
-historical checkpoint investigation, not a supported runtime backend.
-
-## Installation
-
-The supported development environment is Python 3.12 (`>=3.12,<3.13`). Core
-development is supported on macOS and Linux:
+---
+library_name: gnn4colliders
+tags:
+- graph-neural-network
+- high-energy-physics
+- particle-physics
+- event-classification
+- pytorch
+- dgl
+datasets:
+- HWresearch/Delphes
+---
+
+# GNN4Colliders: Pretrained Event Classification Model
+
+This repository contains a pretrained Graph Neural Network (GNN) for
+high-energy-physics collider event classification. The model learns from
+reconstructed event-level objects and can be fine-tuned for downstream
+classification tasks. The accompanying paper source archive is available as
+[`paper.tar.gz`](paper.tar.gz). The supplied PDF was used as the source for
+this card but is not committed because the Hugging Face repository requires
+binary files to use Xet storage.
+
+## Model summary
+
+The model is an event-level foundation-model prototype trained on approximately
+120 million simulated proton-proton collision events from 12 Standard Model
+processes. It represents each event as a fully connected graph whose nodes are
+reconstructed jets, electrons, muons, photons, and missing transverse energy.
+The pretrained representation can be adapted to binary and multiclass physics
+classification tasks by replacing the output layer and fine-tuning the model.
+
+The paper studies two pretrained checkpoints:
+
+- `multiclass_pretrained_model_12`: 12-process multiclass pretraining; this is
+  the recommended starting point for downstream classification.
+- `multilabel_pretrained_model_41`: 41-label pretraining over particle
+  multiplicities and kinematic properties; results in the paper show less
+  consistent downstream transfer from this objective.
+
+The checkpoint directories and their training configurations are included in
+this repository under `legacy/root_gnn_dgl/Pretrained_GNN/`.
+
+## Model details
+
+| Property | Description |
+| --- | --- |
+| Architecture | Graph Network with MLP encoders, message-passing blocks, a decoder, and a task-specific linear head |
+| Framework | PyTorch with DGL |
+| Parameters | Approximately 400,000 trainable parameters |
+| Latent width | 64 |
+| MLP depth | 4 linear layers |
+| Message-passing steps | 4 |
+| Graph topology | Fully connected event graph |
+| Pretraining data | Approximately 120 million events across 12 processes |
+| Pretraining site | Perlmutter supercomputer at NERSC |
+| Authors | Joshua Ho, Ryan Roberts, Shuo Han, and Haichen Wang |
+
+### Inputs
+
+Each graph node contains:
+
+1. transverse momentum \(p_T\),
+2. pseudorapidity \(\eta\),
+3. azimuthal angle \(\phi\),
+4. energy \(E = p_T \cosh(\eta)\) under a massless assumption,
+5. b-tagging information,
+6. electric charge, and
+7. an integer object-type identifier.
+
+Undefined object-specific features use a placeholder value of zero. Each edge
+contains \(\Delta\eta\), wrapped \(\Delta\phi\), and
+\(\Delta R = \sqrt{(\Delta\eta)^2 + (\Delta\phi)^2}\). The global input is the
+number of nodes in the event.
+
+The GNN uses summation aggregation. Its edge, node, and global update functions
+are applied in sequence four times:
+
+$$
+\begin{aligned}
+\vec{y}'_{ij} &= f_e(\vec{x}_i, \vec{x}_j, \vec{y}_{ij}, \vec{z}),\\
+\vec{x}'_i &= f_n\left(\vec{x}_i, \sum_j \vec{y}'_{ij}, \vec{z}\right),\\
+\vec{z}' &= f_g\left(\sum_i \vec{x}'_i, \sum_{i,j} \vec{y}'_{ij}, \vec{z}\right).
+\end{aligned}
+$$
+
+The model returns task logits. A downstream task is responsible for applying
+the appropriate sigmoid or softmax transformation.
+
+## Intended use
+
+This model is intended for research in collider-physics event classification,
+including transfer-learning studies and rapid prototyping of classifiers for
+new physics processes. It is suitable for fine-tuning on simulated events or
+compatible reconstructed event collections.
+
+It is not intended to replace detector, object-reconstruction, calibration, or
+analysis validation. Predictions should not be interpreted as measurements of
+physical parameters without task-specific validation and uncertainty studies.
+
+## Training data
+
+The pretraining samples were generated at a proton-proton center-of-mass energy
+of 13 TeV. The paper describes MadGraph@NLO 2.7.3 at NLO in QCD, Pythia 8.235
+for showering and heavy-particle decays, and Delphes 3.4.2 configured to emulate
+the ATLAS detector.
+
+The 12 pretraining processes are:
+
+- Higgs production: \(ggF\), \(VBF\), \(WH\), \(ZH\), \(t\bar{t}H\), and
+  \(tHq\);
+- top production: single top, \(t\bar{t}\), \(t\bar{t}\gamma\gamma\),
+  \(t\bar{t}W\), \(t\bar{t}t\), and \(t\bar{t}t\bar{t}\).
+
+The paper also evaluates transfer to four additional Delphes processes and to
+ATLAS Open Data. The Open Data evaluation uses the `GamGam` and `1LMET30`
+collections, which are reconstructed with the full ATLAS simulation and
+reconstruction chain rather than Delphes.
+
+The Hugging Face [HWresearch/Delphes dataset](https://huggingface.co/datasets/HWresearch/Delphes)
+contains the public data interface used by the repository for example and
+smoke-test workflows. The complete training samples described in the paper are
+not reproduced by the small smoke-test fixture.
+
+## Pretraining
+
+The multiclass checkpoint predicts the 12 event-process classes using
+categorical cross entropy. The multilabel checkpoint predicts 41 labels using a
+weighted combination of binary cross entropy and mean-squared-error losses.
+Both pretraining runs use an initial learning rate of \(10^{-4}\) with an
+exponential decay factor of \(0.99\) per epoch and stop after five epochs
+without training-loss improvement. The checkpoint with the highest validation
+AUC is selected for fine-tuning.
+
+For downstream fine-tuning, the original output layer is replaced with a new
+linear layer. The new layer starts at learning rate \(10^{-4}\), while the
+pretrained layers start at \(10^{-5}\); both use the same decay schedule.
+
+## Evaluation
+
+The paper evaluates five Delphes binary tasks and two ATLAS Open Data
+multiclass tasks. Reported results are averages over five independently trained
+models, except where dataset size limited the randomization study.
+
+### Delphes downstream tasks
+
+The five binary tasks are:
+
+1. \(t\bar{t}H(\rightarrow\gamma\gamma)\) CP-even versus CP-odd;
+2. FCNC top decays versus \(tHq\);
+3. \(t\bar{t}W\) versus \(t\bar{t}t\);
+4. s-top pair production versus \(t\bar{t}H\); and
+5. \(WH\) versus \(ZH\).
+
+Training sizes range from \(10^3\) to \(10^7\) events per class. Relative to
+training from scratch, multiclass fine-tuning produces its largest gains in the
+low-data regime, with improvements of more than 4 percentage points in
+accuracy and 2.5 AUC points in some settings. The benefit generally decreases
+as the downstream training sample grows.
+
+### ATLAS Open Data
+
+| Task | Baseline accuracy | Multiclass change | Baseline AUC (×100) | Multiclass change |
+| --- | ---: | ---: | ---: | ---: |
+| Higgs production (5 classes) | \(71.8 \pm 0.2\)% | \(-0.1 \pm 0.2\) points | \(91.2 \pm 0.0\) | \(+0.1 \pm 0.0\) points |
+| Triboson (3 classes) | \(54.8 \pm 2.3\)% | \(+3.7 \pm 2.3\) points | \(72.9 \pm 3.3\) | \(+4.6 \pm 3.3\) points |
+
+Multilabel pretraining is less reliable: it can be neutral or harmful on
+several downstream tasks, especially at low statistics. This is a documented
+result of the paper and an important model-selection consideration.
+
+### Computational efficiency
+
+Multiclass and multilabel pretraining required approximately 45.5 and 60.0
+GPU hours, respectively. At full downstream statistics, the paper reports a
+mean fine-tuned-to-baseline full-training-time ratio of 64.6% across the seven
+tasks. When using a target-aware early-stopping metric, fine-tuning reaches the
+baseline AUC target in an average of 43.5% of the baseline time. The estimated
+pretraining break-even point is approximately 14 to 52 downstream tasks,
+depending on the stopping criterion.
+
+### Representation analysis
+
+The paper uses Centered Kernel Alignment (CKA) to compare internal
+representations. Encoder representations remain highly similar across
+pretrained, fine-tuned, and baseline models, while message-passing stages are
+more distinct. Fine-tuning primarily changes the global decoder, preserving
+much of the pretrained event representation while adapting the final decision
+function to the downstream task.
+
+## Limitations and risks
+
+- The model is trained and evaluated on simulated or curated ATLAS Open Data;
+  simulation-to-data mismodelling can affect performance.
+- The pretraining data cover a finite set of processes and detector/object
+  definitions. Results may not transfer to other collision energies,
+  experiments, reconstruction versions, or object selections.
+- The model uses placeholder zeros for features that are undefined for some
+  object types. The paper did not test an explicit feature mask.
+- Multilabel pretraining can produce negative transfer. The multiclass
+  checkpoint is the preferred starting point based on the reported results.
+- The reported metrics are task-specific and do not establish physics-analysis
+  sensitivity, calibration, uncertainty coverage, or discovery significance.
+- The paper studies supervised pretraining and fine-tuning; it does not claim
+  unsupervised or generative capabilities.
+
+## Reproducibility and repository use
+
+The repository contains the legacy DGL implementation used by the paper as
+well as an incremental rewrite under [`src/gnn4colliders`](src/gnn4colliders/).
+The rewrite provides shared ROOT ingestion, collider features, graph building,
+training, evaluation, prediction, and checkpoint/inference utilities. See
+[`README_PROJECT.md`](README_PROJECT.md) for the package development workflow
+and [`docs/architecture.md`](docs/architecture.md) for the target architecture.
+
+For the current rewrite, the supported development environment is:
 
 ```bash
-# 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 parity 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 legacy-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:
-
-```bash
-uv run gnn4colliders train environment=macos
 ```
 
-## Data samples
-
-ROOT inputs are available from the
-[HWresearch/Delphes dataset](https://huggingface.co/datasets/HWresearch/Delphes).
-Download the 64-event smoke-test sample with the Hugging Face CLI:
+The paper checkpoints preserve the historical model and configuration layout;
+use the compatibility documentation before loading them with the rewritten
+package.
 
-```bash
-hf download HWresearch/Delphes testing/ttH_NLO_64.root \
-  --repo-type dataset --local-dir data/raw
-```
+## Citation
 
-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.
+If you use this model or the associated study, please cite:
 
-## 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`](docs/configuration.md).
-
-```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]
+```bibtex
+@article{Ho2026PretrainedEventClassification,
+  title   = {Pretrained Event Classification Model for High Energy Physics Analysis},
+  author  = {Ho, Joshua and Roberts, Ryan and Han, Shuo and Wang, Haichen},
+  note    = {Code and model: https://huggingface.co/HWresearch/GNN4Colliders}
+}
 ```
 
-Train, evaluate, and predict from that cache:
-
-```bash
-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`.
-
-## 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:
-
-```bash
-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:
-
-```bash
-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
-```
+The paper source archive is included in this repository as `paper.tar.gz`.
 
-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:
-
-```bash
-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:
-
-```bash
-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`:
-
-```bash
-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`](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:
-
-```bash
-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`](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:
-
-```bash
-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`](docs/configuration.md) for the group reference and
-[`docs/perlmutter.md`](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
-
-```bash
-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
-```
+## Acknowledgments
 
-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`](docs/performance.md) and
-[`benchmarks/README.md`](benchmarks/README.md).
-See [`docs/testing.md`](docs/testing.md) for test layers, optional dependency
-markers, and package smoke validation.
-
-## Architecture and migration status
-
-See [`docs/architecture.md`](docs/architecture.md) for responsibility
-boundaries and the future sequence-model extension point. See
-[`docs/migration.md`](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.
+This work was supported by the U.S. National Science Foundation Award No.
+2046280 and the U.S. Department of Energy, Office of Science, under contract
+DE-AC02-05CH11231. Joshua Ho acknowledges support from UC Berkeley Summer
+Undergraduate Research Fellowships (SURF) and its donors. The authors
+acknowledge the ATLAS Collaboration for the Open Data and supporting software.

README_PROJECT.md

@@ -0,0 +1,282 @@
+# 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.
+
+```text
+ROOT files -> EventSample -> shared collider features
+                              ├── GraphSample -> ROOT-GNN
+                              └── future SequenceSample -> ROOT-Transformer
+```
+
+The new implementation lives under [`src/gnn4colliders`](src/gnn4colliders/).
+[`legacy/`](legacy/) is a frozen behavioral reference for parity work and
+historical checkpoint investigation, not a supported runtime backend.
+
+## Installation
+
+The supported development environment is Python 3.12 (`>=3.12,<3.13`). Core
+development is supported on macOS and Linux:
+
+```bash
+# 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 parity 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 legacy-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:
+
+```bash
+uv run gnn4colliders train environment=macos
+```
+
+## Data samples
+
+ROOT inputs are available from the
+[HWresearch/Delphes dataset](https://huggingface.co/datasets/HWresearch/Delphes).
+Download the 64-event smoke-test sample with the Hugging Face CLI:
+
+```bash
+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`](docs/configuration.md).
+
+```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]
+```
+
+Train, evaluate, and predict from that cache:
+
+```bash
+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`.
+
+## 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:
+
+```bash
+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:
+
+```bash
+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:
+
+```bash
+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:
+
+```bash
+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`:
+
+```bash
+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`](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:
+
+```bash
+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`](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:
+
+```bash
+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`](docs/configuration.md) for the group reference and
+[`docs/perlmutter.md`](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
+
+```bash
+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`](docs/performance.md) and
+[`benchmarks/README.md`](benchmarks/README.md).
+See [`docs/testing.md`](docs/testing.md) for test layers, optional dependency
+markers, and package smoke validation.
+
+## Architecture and migration status
+
+See [`docs/architecture.md`](docs/architecture.md) for responsibility
+boundaries and the future sequence-model extension point. See
+[`docs/migration.md`](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.

paper.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9587a5b67b73dfc0f88b2a641b98410623a592d47cedd8fc71b35d28e9e1dfbc
+size 142839