PR 7 — Pretraining Objective Experiment Framework

.github/workflows/ci.yml

@@ -0,0 +1,87 @@
+name: CI
+
+on:
+  push:
+  pull_request:
+
+concurrency:
+  group: ci-${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  lint:
+    name: lint
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          version: "0.8.x"
+          enable-cache: true
+      - run: uv sync --dev
+      - run: uv run ruff check .
+      - run: uv run ruff format --check .
+
+  test:
+    name: test
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          version: "0.8.x"
+          enable-cache: true
+      - run: uv sync --dev
+      - run: uv run pytest
+
+  test-macos:
+    name: test (macOS ROOT-GNN CPU)
+    runs-on: macos-14
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          version: "0.8.x"
+          enable-cache: true
+      - run: uv sync --dev --extra root-gnn
+      - run: uv run pytest tests/unit tests/integration
+
+  package:
+    name: package
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          version: "0.8.x"
+          enable-cache: true
+      - run: uv sync --dev
+      - run: uv build
+      - run: uv run python -m twine check dist/*
+      - name: Wheel install smoke test
+        shell: bash
+        run: |
+          set -euo pipefail
+          smoke_dir="$(mktemp -d)"
+          trap 'rm -rf "$smoke_dir"' EXIT
+          uv venv "$smoke_dir/venv"
+          uv pip install --python "$smoke_dir/venv/bin/python" dist/*.whl
+          (
+            cd "$smoke_dir"
+            "$smoke_dir/venv/bin/python" -c 'import gnn4colliders; print(gnn4colliders.__version__)'
+            "$smoke_dir/venv/bin/gnn4colliders" --help
+            "$smoke_dir/venv/bin/gnn4colliders" train --help
+            "$smoke_dir/venv/bin/gnn4colliders" export --help
+          )
+
+  onnx:
+    name: onnx
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          version: "0.8.x"
+          enable-cache: true
+      - run: uv sync --dev --extra root-gnn --extra onnx
+      - run: uv run pytest tests/unit/export -v

.gitignore

@@ -1,16 +1,32 @@
-# Python tooling
 __pycache__/
+trainings/
+scores/
+slurm/
+.onnx
+.png
+
+# Python tooling
 *.py[cod]
-*.egg-info/
-.venv/
 .pytest_cache/
 .ruff_cache/
+.mypy_cache/
+*.egg-info/
+.venv/
+venv/
+htmlcov/
 .coverage
+profiles/
+*.trace.json
 dist/
 build/
+benchmark-results/
+validation_output/
 
 # Local data and generated outputs
-data/
+data/raw/*
+!data/raw/.gitkeep
+data/processed/
 outputs/
-logs/
-*.onnx
+checkpoints/
+
+# Keep the legacy implementation immutable by convention; do not add generated files there.

.python-version

@@ -1 +1 @@
-3.14
+3.12

AGENTS.md

@@ -1,52 +1,477 @@
-# GNN4Colliders agent instructions
+# GNN4Colliders Agent Instructions
 
-GNN4Colliders trains event classifiers (the ROOT-GNN edge network of Ho et
-al., JINST 21 (2026) P08006, a transformer and Deep Sets) on reconstructed
-collider objects. Keep it small and readable: a physicist should be able to follow a
-run from `cli.py` to the model in a few files.
+## Project purpose
 
-## Layout
+GNN4Colliders is a collider machine-learning package intended to support multiple model architectures over shared collider-physics data infrastructure.
+
+The first model family being rewritten is:
+
+```text
+root_gnn
+```
+
+Future model families may include architectures such as:
+
+```text
+root_transformer
+```
+
+Do not design shared infrastructure around assumptions that only apply to GNNs.
+
+---
+
+## Source of truth
+
+Before making non-trivial changes, inspect:
+
+```text
+docs/architecture.md
+docs/migration.md
+```
+
+and the relevant existing source and tests.
+
+The `legacy/` tree is a behavioral reference for the rewrite.
+
+Unless explicitly instructed otherwise:
+
+* do not modify legacy code
+* do not reorganize legacy code
+* do not mechanically copy legacy architecture into the new package
+
+When legacy behavior and documentation disagree, identify the discrepancy rather than silently choosing one.
+
+---
+
+## Package architecture
+
+Production Python code belongs under:
+
+```text
+src/gnn4colliders/
+```
+
+The intended boundaries are:
+
+```text
+data/
+    Data access, datasets, caching, batching, folds, and generic
+    representation-independent data infrastructure.
+
+features/
+    Collider-domain feature transformations, selections, scaling,
+    and derived physics quantities.
+
+graphs/
+    Graph representation, topology, edge construction, and other
+    graph-specific transformations.
+
+models/
+    Architecture-specific neural network implementations.
+
+models/root_gnn/
+    ROOT-GNN-specific model components.
+
+training/
+    Training lifecycle, losses, metrics, checkpointing,
+    reproducibility, and distributed-training utilities where
+    architecture-independent.
+
+inference/
+    Prediction, evaluation, output writing, and model export.
+
+cli/
+    Thin command-line entry points only.
+```
+
+Do not place ROOT file-reading logic inside a specific model family unless it is genuinely architecture-specific.
+
+Do not place reusable physics-feature logic inside `root_gnn`.
+
+Do not place core implementation in notebooks or shell scripts.
+
+---
+
+## Shared infrastructure vs model-specific code
+
+Use this decision rule:
+
+**Shared collider or experiment behavior → shared package module.**
+
+Examples:
+
+```text
+ROOT reading
+event selections
+physics features
+dataset splits
+metrics
+checkpoint orchestration
+```
+
+**Representation-specific behavior → representation module.**
+
+Examples:
+
+```text
+graph topology
+edge construction
+sequence/token construction
+```
+
+**Architecture-specific behavior → models/<model_family>/.**
+
+Examples:
+
+```text
+message-passing network
+transformer encoder
+architecture-specific layers
+```
+
+The package should allow future architectures to reuse the same data and physics infrastructure where practical.
+
+---
+
+## Configuration
+
+Experiment configuration should describe intent rather than expose implementation details.
+
+Prefer:
+
+```yaml
+model:
+  type: root_gnn
+```
+
+over configuration that directly names Python module paths and class names.
+
+Configuration should eventually support composition across concerns such as:
 
 ```text
-src/gnn4colliders/   flat modules; see docs/design.md for what each holds
-configs/             one self-contained YAML file per experiment
-scripts/             Slurm wrappers and small helpers (no core logic)
-tests/               pytest, with small ROOT fixtures in tests/fixtures
-Pretrained_GNN/      the published pretrained checkpoints (Git LFS)
+data
+model
+task
+trainer
+environment
 ```
 
-* `README.md` is the paper's Hugging Face model card. Do not edit it; usage
-  goes in `docs/usage.md`, design and results in `docs/design.md`.
-* The main benchmark is `configs/gamgam.yaml` (ATLAS Open Data GamGam, 5
-  classes); compare with Table 3 of the paper (`docs/design.md`).
+Do not hardcode site-specific filesystem paths, CUDA settings, Slurm settings, or machine configuration into model or task definitions.
+
+Environment-specific configuration belongs in an environment configuration layer.
 
-## Rules
+---
 
-* Add a module or an abstraction only when existing code cannot hold it
-  clearly. No new frameworks (Lightning, Hydra, DGL, PyG) without agreement.
-* Configuration is plain dataclasses in `config.py`; every key has a default
-  or is required, and unknown keys are errors. Site settings (accounts,
-  paths, `UV_CACHE_DIR`) belong in scripts, not in configs or code.
-* Name domain quantities (`labels`, `weight`, `fold`, `event_index`) instead
-  of relying on tensor column positions, except for the documented seven node
-  columns.
-* Randomness is seeded explicitly (`train.seed`, `data.seed`); never mutate
-  global random state elsewhere. Do not claim bitwise reproducibility on GPUs.
-* The node schema, edge construction, MLP layout and parameter names must
-  stay compatible with `Pretrained_GNN/`; a test loads both pretrained models.
-* Changes to features, splits, loss or metrics change physics results: say so
-  explicitly and re-check the GamGam benchmark.
+## CLI design
 
-## Checks
+The intended user-facing interface is a single project CLI with subcommands conceptually similar to:
 
 ```bash
-uv sync --dev
-uv run pytest            # GPU tests run on a CUDA node
-uv run ruff check . && uv run ruff format --check .
+gnn4colliders prepare
+gnn4colliders train
+gnn4colliders evaluate
+gnn4colliders predict
+gnn4colliders export
+```
+
+CLI modules should be thin.
+
+They may:
+
+* parse/compose configuration
+* construct application objects
+* invoke library functions
+* handle user-facing errors
+
+They should not contain core training, data-processing, graph-building, or model logic.
+
+---
+
+## Rewrite strategy
+
+This repository is being rewritten incrementally.
+
+Do not attempt to rewrite the entire legacy repository in one task.
+
+For substantial migrations:
+
+1. inspect the relevant legacy implementation
+2. identify externally observable behavior
+3. inspect existing characterization/parity tests
+4. state or infer the intended new interface
+5. implement the smallest coherent unit
+6. add or update tests
+7. run relevant validation
+8. report intentional differences and unresolved ambiguity
+
+Prefer vertical, testable migration steps over large speculative refactors.
+
+---
+
+## Behavioral parity
+
+The legacy implementation defines important behavior that may need to remain compatible during migration.
+
+Important compatibility areas include, where applicable:
+
+* input feature definitions
+* feature ordering
+* tensor shapes
+* tensor dtypes
+* graph topology
+* edge feature definitions
+* label semantics
+* event weights
+* fold semantics
+* loss calculations
+* metric definitions
+* model outputs
+* checkpoint compatibility
+* inference outputs
+
+Do not alter behavior merely because the legacy implementation appears unusual.
+
+If behavior seems incorrect or ambiguous:
+
+1. document it
+2. characterize it with a test when possible
+3. separate compatibility from proposed improvement
+
+Improvements can be made deliberately after the behavior is understood.
+
+---
+
+## Testing
+
+Tests belong under:
+
+```text
+tests/unit/
+tests/integration/
+tests/parity/
 ```
 
-On Perlmutter, run GPU work on compute nodes (`salloc -C gpu -q
-shared_interactive ...`) with `UV_CACHE_DIR=$SCRATCH/.cache/uv`.
+### Unit tests
+
+Use unit tests for isolated transformations and components.
+
+They should be:
+
+* fast
+* deterministic
+* focused
+* independent of large external datasets
+
+### Integration tests
+
+Use integration tests for small end-to-end workflows such as:
+
+```text
+input fixture
+→ features
+→ representation
+→ model
+→ output
+```
+
+### Parity tests
+
+Use parity tests to compare the rewrite with the legacy implementation.
+
+Prefer small deterministic fixtures and reference outputs.
+
+Do not weaken or delete parity tests simply to make new code pass.
+
+If a parity difference is intentional, document the reason.
+
+---
+
+## Test data
+
+Large production datasets do not belong in the repository.
+
+Small deterministic fixtures may live under:
+
+```text
+tests/fixtures/
+```
+
+Fixtures should be only large enough to exercise relevant behavior.
+
+Where possible, create reference outputs for deterministic legacy behavior before replacing that behavior.
+
+---
+
+## Reproducibility
+
+Reproducibility is a project requirement.
+
+Randomness should be explicit and controllable.
+
+Where relevant, account for:
+
+* Python random state
+* NumPy random state
+* PyTorch random state
+* CUDA random state
+* DataLoader workers
+* shuffling
+* sampling
+* data augmentation
+* distributed execution
+
+Do not introduce hidden global random-state mutation.
+
+Seeds should be passed or configured explicitly.
+
+Do not claim bitwise deterministic training unless the execution environment actually guarantees it.
+
+---
+
+## Code quality
+
+Domain meaning must be represented by named fields or schemas, not only by
+positional column indices. New APIs must not require callers to know that a
+particular tensor column means fold, weight, or another domain field.
+
+Prefer:
+
+* small modules with clear responsibilities
+* typed function signatures
+* explicit inputs and outputs
+* dataclasses or typed configuration where appropriate
+* composition over hidden global state
+* readable names over abbreviations
+* dependency injection over implicit filesystem/environment assumptions
+
+Avoid:
+
+* repository-relative `sys.path` modifications
+* mutable global configuration
+* hidden singleton state
+* wildcard imports
+* giant utility modules
+* giant training scripts
+* model-specific behavior in generic data code
+* duplicated preprocessing logic
+* unnecessary abstraction introduced before a second use case exists
+
+Keep public interfaces intentionally small.
+
+---
+
+## Dependencies
+
+Do not introduce a dependency merely to simplify a small amount of code.
+
+Before adding a major framework or runtime dependency, explain why it is necessary.
+
+In particular, do not introduce architectural frameworks such as:
+
+```text
+PyTorch Lightning
+Kedro
+```
+
+unless the task explicitly calls for evaluating or adopting them.
+
+Scientific dependencies should be introduced intentionally and with environment compatibility in mind.
+
+PyTorch is shared infrastructure for the expected model families and belongs
+in the base package dependencies. Architecture-specific dependencies belong in
+named extras, for example `root-gnn` for DGL. The canonical development
+environment for the active ROOT-GNN rewrite is:
+
+```bash
+uv sync --extra root-gnn
+```
+
+---
+
+## Documentation
+
+Update documentation when a change alters:
+
+* architecture
+* configuration
+* public interfaces
+* expected workflow
+* compatibility guarantees
+
+Do not document speculative functionality as if it already exists.
+
+Use:
+
+```text
+docs/architecture.md
+```
+
+for the target system structure and important architectural decisions.
+
+Use:
+
+```text
+docs/migration.md
+```
+
+for rewrite progress and migration sequencing.
+
+---
+
+## Scope discipline
+
+Do not make unrelated changes.
+
+When implementing a task:
+
+* modify only the modules necessary for that task
+* do not opportunistically refactor unrelated code
+* do not rename public concepts without a clear reason
+* do not remove compatibility behavior without explicit instruction
+* do not modify generated artifacts unless required
+
+If a larger architectural issue is discovered, report it rather than expanding the task automatically.
+
+---
+
+## Validation before finishing
+
+After code changes, run the relevant available checks.
+
+As the project matures, the expected baseline should include:
+
+```bash
+pytest
+ruff check .
+```
+
+ROOT-GNN parity validation requires DGL. A package-wide environment where DGL
+is intentionally absent may skip DGL-dependent tests, but that is incomplete
+ROOT-GNN validation. The required parity gate is:
+
+```bash
+GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 pytest
+```
+
+The gate must fail when DGL cannot be imported.
+
+Run more focused tests first when appropriate.
+
+If the full suite is expensive, run the relevant subset and clearly report what was and was not run.
+
+Do not claim validation succeeded unless the commands actually succeeded.
+
+---
+
+## Completion report
+
+For non-trivial changes, finish with a concise report containing:
+
+* files changed
+* behavior implemented
+* tests/checks run
+* parity status where relevant
+* intentional deviations from legacy behavior
+* unresolved questions or ambiguities
 
-Report what changed, what was run and what was not, and any change to
-results.
+Do not hide failed checks or incomplete behavior.

CHANGELOG.md

@@ -0,0 +1,13 @@
+# Changelog
+
+## Unreleased
+
+### Added
+
+- Release validation and CI coverage for the rewritten package.
+- Installed-package Hydra configuration discovery and public import smoke tests.
+
+### Compatibility
+
+- The `0.1.0` package version remains a static, single-source version exposed
+  as `gnn4colliders.__version__`.

MANIFEST.in

@@ -0,0 +1,3 @@
+include README.md
+include CHANGELOG.md
+recursive-include docs *.md

README.md

@@ -6,6 +6,7 @@ tags:
 - particle-physics
 - event-classification
 - pytorch
+- dgl
 datasets:
 - HWresearch/Delphes
 ---

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.

benchmarks/README.md

@@ -0,0 +1,48 @@
+# Performance benchmarks
+
+These scripts use fixed synthetic inputs, explicit seeds, warmup iterations,
+and JSON-lines output. GPU timings synchronize after every measured operation;
+warmup and one-time setup costs are excluded from steady-state numbers.
+
+```bash
+uv run python benchmarks/benchmark_preprocessing.py
+uv run python benchmarks/benchmark_dataloader.py
+uv run python benchmarks/benchmark_training.py --device cpu
+uv run python benchmarks/benchmark_inference.py --device cpu
+uv run python benchmarks/benchmark_training.py --device cuda --profile
+```
+
+DGL-dependent scripts report a structured `skipped` result when the optional
+`root-gnn` extra is absent. Profiler traces go under `profiles/` and are not
+committed.
+
+## Baseline measurements
+
+The portable baseline in this checkout uses Python 3.12.13, PyTorch 2.2.2+cu121,
+DGL 2.4.0+cu121, CPU, `nodes=32`, `iterations=10`, and `warmup=3`. Exact timings are machine
+dependent; the JSON output from a local run is authoritative.
+
+| Measurement | Mean | Median |
+| --- | ---: | ---: |
+| feature construction | 0.47 ms/event | 0.46 ms/event |
+| edge features | 0.11 ms/graph | 0.11 ms/graph |
+| graph-sample loader | 1.65 ms/iteration | 1.65 ms/iteration |
+| ROOT-GNN training step | 21.19 ms/step | 5.32 ms/step |
+| ROOT-GNN inference | 1.48 ms/graph | 1.49 ms/graph |
+
+On the available NVIDIA A100-PCIE-40GB with CUDA 12.1 runtime, the same small
+synthetic ROOT-GNN benchmark measured 6.41 ms/step (17.3 MiB peak allocated)
+and 2.00 ms/graph. These are microbenchmarks, not production-workload claims.
+
+DGL is available in this environment. The training mean is skewed by one CPU
+warmup-adjacent outlier; median is the more useful steady-state indicator. DDP
+measurements require a multi-process run and are not inferred from CPU numbers.
+
+The topology cache is bounded to 32 node-count/device/policy entries and only
+stores reusable index tensors, never event-specific edge features.
+
+On the same CPU, a direct topology microbenchmark for 64 nodes measured
+201.6 microseconds per cold construction versus 10.5 microseconds for a warm
+cache lookup (about 19x for this isolated operation). This is a targeted
+index-construction result, not an end-to-end training speedup; graph feature
+construction and DGL message passing remain separate costs.

benchmarks/__init__.py

@@ -0,0 +1 @@
+"""Standalone, reproducible performance measurements for GNN4Colliders."""

benchmarks/_common.py

@@ -0,0 +1,70 @@
+"""Small helpers shared by benchmark entry points."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import platform
+import statistics
+import time
+from collections.abc import Callable
+from typing import Any
+
+import torch
+
+
+def common_parser(description: str) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=description)
+    parser.add_argument("--iterations", type=int, default=20)
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--seed", type=int, default=1234)
+    return parser
+
+
+def synchronize(device: torch.device) -> None:
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+
+
+def measure(
+    operation: Callable[[], Any], *, iterations: int, warmup: int, device: torch.device
+) -> dict[str, float]:
+    for _ in range(warmup):
+        operation()
+    synchronize(device)
+    durations = []
+    for _ in range(iterations):
+        start = time.perf_counter()
+        operation()
+        synchronize(device)
+        durations.append(time.perf_counter() - start)
+    return {
+        "mean_ms": statistics.mean(durations) * 1000,
+        "median_ms": statistics.median(durations) * 1000,
+        "min_ms": min(durations) * 1000,
+        "max_ms": max(durations) * 1000,
+        "stdev_ms": statistics.stdev(durations) * 1000 if len(durations) > 1 else 0.0,
+    }
+
+
+def metadata(device: torch.device) -> dict[str, Any]:
+    result: dict[str, Any] = {
+        "python": platform.python_version(),
+        "torch": torch.__version__,
+        "device": str(device),
+        "cuda_available": torch.cuda.is_available(),
+    }
+    if device.type == "cuda" and torch.cuda.is_available():
+        result["gpu"] = torch.cuda.get_device_name(device)
+    try:
+        import dgl
+
+        result["dgl"] = dgl.__version__
+    except ImportError:
+        result["dgl"] = None
+    return result
+
+
+def report(name: str, values: dict[str, Any]) -> None:
+    print(json.dumps({"benchmark": name, **values}, sort_keys=True))

benchmarks/benchmark_dataloader.py

@@ -0,0 +1,63 @@
+"""Measure deterministic loader/batch construction when ROOT-GNN is installed."""
+
+from __future__ import annotations
+
+import torch
+
+try:
+    from ._common import common_parser, measure, metadata, report
+except ImportError:
+    from _common import common_parser, measure, metadata, report
+
+
+def main() -> None:
+    parser = common_parser(__doc__)
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--samples", type=int, default=64)
+    args = parser.parse_args()
+    device = torch.device(args.device)
+    try:
+        import dgl
+
+        from gnn4colliders.data import (
+            EventMetadata,
+            GraphDataLoader,
+            GraphDataset,
+            GraphSample,
+        )
+    except ImportError:
+        report(
+            "dataloader",
+            {**metadata(device), "status": "skipped: install root-gnn extra"},
+        )
+        return
+    samples = []
+    for index in range(args.samples):
+        graph = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 0])), num_nodes=2)
+        graph.ndata["features"] = torch.ones(2, 7)
+        graph.edata["features"] = torch.ones(2, 3)
+        samples.append(
+            GraphSample(
+                graph, torch.tensor(index % 2), None, EventMetadata(0, 1.0, str(index))
+            )
+        )
+    loader = GraphDataLoader(GraphDataset(samples), args.batch_size)
+    result = measure(
+        lambda: list(loader),
+        iterations=args.iterations,
+        warmup=args.warmup,
+        device=device,
+    )
+    report(
+        "dataloader",
+        {
+            **metadata(device),
+            **result,
+            "batch_size": args.batch_size,
+            "batches_per_second": 1000 / result["mean_ms"],
+        },
+    )
+
+
+if __name__ == "__main__":
+    main()

benchmarks/benchmark_inference.py

@@ -0,0 +1,52 @@
+"""Measure ROOT-GNN inference with fixed synthetic graph input."""
+
+from __future__ import annotations
+
+import torch
+
+try:
+    from ._common import common_parser, measure, metadata, report
+except ImportError:
+    from _common import common_parser, measure, metadata, report
+
+
+def main() -> None:
+    parser = common_parser(__doc__)
+    args = parser.parse_args()
+    device = torch.device(args.device)
+    try:
+        import dgl
+
+        from gnn4colliders.models.root_gnn import EdgeNetwork
+    except ImportError:
+        report(
+            "inference",
+            {**metadata(device), "status": "skipped: install root-gnn extra"},
+        )
+        return
+    graph = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 0])), num_nodes=2).to(
+        device
+    )
+    graph.ndata["features"] = torch.ones(2, 7, device=device)
+    graph.edata["features"] = torch.ones(2, 3, device=device)
+    model = (
+        EdgeNetwork(graph, None, hid_size=16, out_size=2, n_layers=1, n_proc_steps=1)
+        .to(device)
+        .eval()
+    )
+
+    def infer() -> None:
+        with torch.inference_mode():
+            model(graph)
+
+    result = measure(
+        infer, iterations=args.iterations, warmup=args.warmup, device=device
+    )
+    report(
+        "inference",
+        {**metadata(device), **result, "graphs_per_second": 1000 / result["mean_ms"]},
+    )
+
+
+if __name__ == "__main__":
+    main()

benchmarks/benchmark_preprocessing.py

@@ -0,0 +1,61 @@
+"""Measure shared feature and graph preprocessing on fixed synthetic events."""
+
+from __future__ import annotations
+
+import torch
+
+from gnn4colliders.features import build_node_features
+from gnn4colliders.graphs import build_edge_features, fully_connected_edges
+
+try:
+    from ._common import common_parser, measure, metadata, report
+except ImportError:
+    from _common import common_parser, measure, metadata, report
+
+
+def main() -> None:
+    parser = common_parser(__doc__)
+    parser.add_argument("--nodes", type=int, default=32)
+    args = parser.parse_args()
+    torch.manual_seed(args.seed)
+    event = {
+        "pt": torch.arange(args.nodes, dtype=torch.float32) + 1,
+        "eta": torch.linspace(-2, 2, args.nodes),
+        "phi": torch.linspace(-3.0, 3.0, args.nodes),
+    }
+    branches = [["pt"], ["eta"], ["phi"], "CALC_E", [1.0], [0.0], "NODE_TYPE"]
+    object_types = ["vector"]
+    scales = [1.0] * 7
+    device = torch.device(args.device)
+
+    feature_result = measure(
+        lambda: build_node_features(event, branches, object_types, scales),
+        iterations=args.iterations,
+        warmup=args.warmup,
+        device=device,
+    )
+    nodes = build_node_features(event, branches, object_types, scales)[0]
+    src, dst = fully_connected_edges(args.nodes)
+    graph_result = measure(
+        lambda: build_edge_features(nodes, src, dst, eta_index=1, phi_index=2),
+        iterations=args.iterations,
+        warmup=args.warmup,
+        device=device,
+    )
+    base = {**metadata(device), "nodes": args.nodes, "iterations": args.iterations}
+    report(
+        "feature_construction",
+        {
+            **base,
+            **feature_result,
+            "events_per_second": 1000 / feature_result["mean_ms"],
+        },
+    )
+    report(
+        "edge_features",
+        {**base, **graph_result, "graphs_per_second": 1000 / graph_result["mean_ms"]},
+    )
+
+
+if __name__ == "__main__":
+    main()

benchmarks/benchmark_training.py

@@ -0,0 +1,73 @@
+"""Measure a short ROOT-GNN training section and optionally emit a profiler trace."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+try:
+    from ._common import common_parser, measure, metadata, report
+except ImportError:
+    from _common import common_parser, measure, metadata, report
+
+
+def main() -> None:
+    parser = common_parser(__doc__)
+    parser.add_argument("--batch-size", type=int, default=4)
+    parser.add_argument("--profile", action="store_true")
+    parser.add_argument("--profile-dir", default="profiles")
+    args = parser.parse_args()
+    device = torch.device(args.device)
+    try:
+        import dgl
+
+        from gnn4colliders.models.root_gnn import EdgeNetwork
+    except ImportError:
+        report(
+            "training",
+            {**metadata(device), "status": "skipped: install root-gnn extra"},
+        )
+        return
+    torch.manual_seed(args.seed)
+    graph = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 0])), num_nodes=2).to(
+        device
+    )
+    graph.ndata["features"] = torch.randn(2, 7, device=device)
+    graph.edata["features"] = torch.randn(2, 3, device=device)
+    model = EdgeNetwork(
+        graph, None, hid_size=16, out_size=2, n_layers=1, n_proc_steps=1
+    ).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+
+    def step() -> None:
+        optimizer.zero_grad(set_to_none=True)
+        loss = model(graph).square().mean()
+        loss.backward()
+        optimizer.step()
+
+    if args.profile:
+        profile_dir = Path(args.profile_dir)
+        profile_dir.mkdir(parents=True, exist_ok=True)
+        with torch.profiler.profile(record_shapes=True, profile_memory=True) as prof:
+            for _ in range(args.warmup + args.iterations):
+                step()
+        prof.export_chrome_trace(str(profile_dir / "training.trace.json"))
+    result = measure(
+        step, iterations=args.iterations, warmup=args.warmup, device=device
+    )
+    report(
+        "training_step",
+        {
+            **metadata(device),
+            **result,
+            "steps_per_second": 1000 / result["mean_ms"],
+            "peak_memory_mb": torch.cuda.max_memory_allocated(device) / 2**20
+            if device.type == "cuda"
+            else 0.0,
+        },
+    )
+
+
+if __name__ == "__main__":
+    main()

configs/.gitkeep

@@ -0,0 +1 @@
+

configs/data/.gitkeep

@@ -0,0 +1 @@
+

configs/delphes_contrastive.yaml

@@ -1,81 +0,0 @@
-# Contrastive (SimCLR) pretraining of the object transformer on the 12
-# Delphes samples of configs/delphes_jepa.yaml, without their labels (same
-# event cache): two random views of every event (rotation, eta flip, energy
-# scale, object dropout) are pulled together (src/gnn4colliders/contrastive.py).
-# Fine-tune with configs/gamgam_jepa_finetune.yaml
-# model.pretrained=outputs/delphes_contrastive/last.pt, or compare encoders
-# without fine-tuning: gnn4colliders probe configs/gamgam_jepa_finetune.yaml
-# data.train_events=1000 split=validation model.pretrained=...
-#
-#   gnn4colliders train configs/delphes_contrastive.yaml
-name: delphes_contrastive
-
-data:
-  hf_repo: HWresearch/Delphes
-  hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e
-  files:
-    - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0}
-    - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1}
-    - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2}
-    - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3}
-    - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4}
-    - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5}
-    - {path: samples/photon/ttyy/ttyy.root, label: 6}
-    - {path: samples/top/multitop/tttt.root, label: 7}
-    - {path: samples/top/single-top/SingleT_schan.root, label: 8}
-    - {path: samples/top/ttbar/ttbar.root, label: 9}
-    - {path: samples/top/ttv/ttW.root, label: 10}
-    - {path: samples/top/multitop/ttt.root, label: 11}
-  tree_name: output
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-    - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  scales: [0.1, 1, 1, 0.1, 1, 1, 1]
-  # Number is the event index; Number % 10 is the fold. Pretraining trained
-  # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests.
-  fold_var: Number
-  fold_modulus: 10
-  train_folds: [2, 3, 4, 5, 6, 7, 8, 9]
-  validation_folds: [0]
-  test_folds: [1]
-  # JEPA ignores labels and weights; the cache is that of
-  # configs/delphes_pretraining.yaml.
-  weight_var: weight
-  cache: outputs/delphes_pretraining/events.pt
-  batch_size: 1024
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: contrastive
-  backbone: transformer
-  # The encoder; configs/gamgam_jepa_finetune.yaml must match it.
-  hid_size: 128
-  n_layers: 4
-  n_heads: 8
-  ffn_size: 512
-  dropout: 0.0
-  pair_bias: true
-  temperature: 0.1
-  scale_jitter: 0.1
-  object_dropout: 0.1
-
-task:
-  name: contrastive
-
-train:
-  max_epochs: 30
-  optimizer: adamw
-  learning_rate: 5.0e-4
-  weight_decay: 0.05
-  lr_decay: 0.95
-  early_stopping: null
-  # Fine-tune from last.pt.
-  best_metric: loss

configs/delphes_jepa.yaml

@@ -1,77 +0,0 @@
-# JEPA pretraining of the object transformer on the 12 Delphes samples of
-# the supervised pretraining (configs/delphes_pretraining.yaml), without
-# their labels. Fine-tune with configs/gamgam_jepa_finetune.yaml.
-#
-#   gnn4colliders prepare configs/delphes_jepa.yaml
-#   gnn4colliders train configs/delphes_jepa.yaml
-name: delphes_jepa
-
-data:
-  hf_repo: HWresearch/Delphes
-  hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e
-  files:
-    - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0}
-    - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1}
-    - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2}
-    - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3}
-    - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4}
-    - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5}
-    - {path: samples/photon/ttyy/ttyy.root, label: 6}
-    - {path: samples/top/multitop/tttt.root, label: 7}
-    - {path: samples/top/single-top/SingleT_schan.root, label: 8}
-    - {path: samples/top/ttbar/ttbar.root, label: 9}
-    - {path: samples/top/ttv/ttW.root, label: 10}
-    - {path: samples/top/multitop/ttt.root, label: 11}
-  tree_name: output
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-    - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  scales: [0.1, 1, 1, 0.1, 1, 1, 1]
-  # Number is the event index; Number % 10 is the fold. Pretraining trained
-  # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests.
-  fold_var: Number
-  fold_modulus: 10
-  train_folds: [2, 3, 4, 5, 6, 7, 8, 9]
-  validation_folds: [0]
-  test_folds: [1]
-  # JEPA ignores labels and weights; the cache is that of
-  # configs/delphes_pretraining.yaml.
-  weight_var: weight
-  cache: outputs/delphes_pretraining/events.pt
-  batch_size: 1024
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: jepa
-  # The encoder; configs/gamgam_jepa_finetune.yaml must match it.
-  hid_size: 128
-  n_layers: 4
-  n_heads: 8
-  ffn_size: 512
-  dropout: 0.0
-  pair_bias: true
-  predictor_layers: 2
-  mask_fraction: 0.4
-  ema_momentum: 0.996
-
-task:
-  name: jepa
-
-train:
-  max_epochs: 30
-  optimizer: adamw
-  learning_rate: 5.0e-4
-  weight_decay: 0.05
-  lr_decay: 0.95
-  early_stopping: null
-  # The target drifts with the EMA, so the loss does not compare epochs;
-  # fine-tune from last.pt.
-  best_metric: loss

configs/delphes_jepa_edge.yaml

@@ -1,82 +0,0 @@
-# JEPA pretraining of the edge network, of the size of the supervised
-# pretrained models (Pretrained_GNN/), on the same 12 Delphes samples without
-# their labels. Fine-tune with configs/gamgam_finetune.yaml and
-# model.pretrained=outputs/delphes_jepa_edge/last.pt.
-#
-#   gnn4colliders prepare configs/delphes_jepa_edge.yaml
-#   gnn4colliders train configs/delphes_jepa_edge.yaml
-name: delphes_jepa_edge
-
-data:
-  hf_repo: HWresearch/Delphes
-  hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e
-  files:
-    - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0}
-    - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1}
-    - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2}
-    - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3}
-    - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4}
-    - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5}
-    - {path: samples/photon/ttyy/ttyy.root, label: 6}
-    - {path: samples/top/multitop/tttt.root, label: 7}
-    - {path: samples/top/single-top/SingleT_schan.root, label: 8}
-    - {path: samples/top/ttbar/ttbar.root, label: 9}
-    - {path: samples/top/ttv/ttW.root, label: 10}
-    - {path: samples/top/multitop/ttt.root, label: 11}
-  tree_name: output
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-    - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  scales: [0.1, 1, 1, 0.1, 1, 1, 1]
-  # Number is the event index; Number % 10 is the fold. Pretraining trained
-  # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests.
-  fold_var: Number
-  fold_modulus: 10
-  train_folds: [2, 3, 4, 5, 6, 7, 8, 9]
-  validation_folds: [0]
-  test_folds: [1]
-  # JEPA ignores labels and weights; the cache is that of
-  # configs/delphes_pretraining.yaml.
-  weight_var: weight
-  cache: outputs/delphes_pretraining/events.pt
-  batch_size: 1024
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: jepa
-  backbone: edge_network
-  # The encoder: as in Pretrained_GNN/*/config.yaml.
-  hid_size: 64
-  n_layers: 4
-  n_proc_steps: 4
-  dropout: 0.0
-  # The predictor: transformer blocks.
-  predictor_layers: 2
-  n_heads: 4
-  ffn_size: 256
-  mask_fraction: 0.4
-  ema_momentum: 0.996
-  # Without it the encodings collapse (docs/design.md, JEPA pretraining).
-  variance_weight: 0.1
-
-task:
-  name: jepa
-
-train:
-  max_epochs: 30
-  optimizer: adamw
-  learning_rate: 5.0e-4
-  weight_decay: 0.05
-  lr_decay: 0.95
-  early_stopping: null
-  # The target drifts with the EMA, so the loss does not compare epochs;
-  # fine-tune from last.pt.
-  best_metric: loss

configs/delphes_masked.yaml

@@ -1,78 +0,0 @@
-# Masked object modelling of the object transformer on the 12 Delphes
-# samples of configs/delphes_jepa.yaml, without their labels (same event
-# cache): JEPA's masking and predictor, but the predictor predicts the hidden
-# objects' kinematics instead of their target encodings (model.jepa_target:
-# features; src/gnn4colliders/jepa.py). Fine-tune or probe with
-# configs/gamgam_jepa_finetune.yaml model.pretrained=outputs/delphes_masked/last.pt.
-#
-#   gnn4colliders train configs/delphes_masked.yaml
-name: delphes_masked
-
-data:
-  hf_repo: HWresearch/Delphes
-  hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e
-  files:
-    - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0}
-    - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1}
-    - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2}
-    - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3}
-    - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4}
-    - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5}
-    - {path: samples/photon/ttyy/ttyy.root, label: 6}
-    - {path: samples/top/multitop/tttt.root, label: 7}
-    - {path: samples/top/single-top/SingleT_schan.root, label: 8}
-    - {path: samples/top/ttbar/ttbar.root, label: 9}
-    - {path: samples/top/ttv/ttW.root, label: 10}
-    - {path: samples/top/multitop/ttt.root, label: 11}
-  tree_name: output
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-    - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  scales: [0.1, 1, 1, 0.1, 1, 1, 1]
-  # Number is the event index; Number % 10 is the fold. Pretraining trained
-  # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests.
-  fold_var: Number
-  fold_modulus: 10
-  train_folds: [2, 3, 4, 5, 6, 7, 8, 9]
-  validation_folds: [0]
-  test_folds: [1]
-  # JEPA ignores labels and weights; the cache is that of
-  # configs/delphes_pretraining.yaml.
-  weight_var: weight
-  cache: outputs/delphes_pretraining/events.pt
-  batch_size: 1024
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: jepa
-  jepa_target: features
-  # The encoder; configs/gamgam_jepa_finetune.yaml must match it.
-  hid_size: 128
-  n_layers: 4
-  n_heads: 8
-  ffn_size: 512
-  dropout: 0.0
-  pair_bias: true
-  predictor_layers: 2
-  mask_fraction: 0.4
-
-task:
-  name: jepa
-
-train:
-  max_epochs: 30
-  optimizer: adamw
-  learning_rate: 5.0e-4
-  weight_decay: 0.05
-  lr_decay: 0.95
-  early_stopping: null
-  # Fine-tune from last.pt.
-  best_metric: loss

configs/delphes_pfn.yaml

@@ -1,86 +0,0 @@
-# Meta-training of a prior-fitted network (src/gnn4colliders/pfn.py) on
-# classification tasks drawn from the 12 Delphes samples: it learns to
-# classify events given labelled context events, and is then evaluated in
-# context on new data (configs/gamgam_pfn.yaml). The encoder has the sizes of
-# configs/delphes_jepa.yaml, so model.pretrained=outputs/delphes_jepa/last.pt
-# starts it from the JEPA encoder. task.prior_classes=[1,6,7,8,9,10,11]
-# leaves out the five processes of the GamGam task (ttH, ggF, VBF, WH, ZH).
-#
-#   gnn4colliders train configs/delphes_pfn.yaml
-name: delphes_pfn
-
-data:
-  hf_repo: HWresearch/Delphes
-  hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e
-  files:
-    - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0}
-    - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1}
-    - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2}
-    - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3}
-    - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4}
-    - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5}
-    - {path: samples/photon/ttyy/ttyy.root, label: 6}
-    - {path: samples/top/multitop/tttt.root, label: 7}
-    - {path: samples/top/single-top/SingleT_schan.root, label: 8}
-    - {path: samples/top/ttbar/ttbar.root, label: 9}
-    - {path: samples/top/ttv/ttW.root, label: 10}
-    - {path: samples/top/multitop/ttt.root, label: 11}
-  tree_name: output
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-    - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  scales: [0.1, 1, 1, 0.1, 1, 1, 1]
-  # Number is the event index; Number % 10 is the fold. Pretraining trained
-  # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests.
-  fold_var: Number
-  fold_modulus: 10
-  train_folds: [2, 3, 4, 5, 6, 7, 8, 9]
-  validation_folds: [0]
-  test_folds: [1]
-  # Tasks use the labels, not the weights; the cache is that of
-  # configs/delphes_pretraining.yaml.
-  weight_var: weight
-  cache: outputs/delphes_pretraining/events.pt
-  # Events per step (tasks of context and query events).
-  batch_size: 4096
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: pfn
-  hid_size: 128
-  n_layers: 4
-  n_heads: 8
-  ffn_size: 512
-  dropout: 0.0
-  pair_bias: true
-  pretrained: null
-  max_classes: 10
-  pfn_layers: 6
-
-task:
-  name: pfn
-  # The prior: classes per task, context events (log-uniform) and queries.
-  prior_classes: null
-  task_classes: [2, 10]
-  context_events: [16, 2048]
-  query_events: 256
-  energy_jitter: 0.2
-  steps_per_epoch: 1000
-  validation_steps: 50
-
-train:
-  max_epochs: 50
-  optimizer: adamw
-  learning_rate: 1.0e-4
-  weight_decay: 0.01
-  lr_decay: 0.95
-  early_stopping: null
-  best_metric: loss

configs/delphes_pretraining.yaml

@@ -1,61 +0,0 @@
-# 12-class pretraining on the Delphes samples of Ho et al. (2026), read from
-# the HWresearch/Delphes Hugging Face dataset.
-name: delphes_pretraining
-
-data:
-  hf_repo: HWresearch/Delphes
-  hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e
-  files:
-    - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0}
-    - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1}
-    - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2}
-    - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3}
-    - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4}
-    - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5}
-    - {path: samples/photon/ttyy/ttyy.root, label: 6}
-    - {path: samples/top/multitop/tttt.root, label: 7}
-    - {path: samples/top/single-top/SingleT_schan.root, label: 8}
-    - {path: samples/top/ttbar/ttbar.root, label: 9}
-    - {path: samples/top/ttv/ttW.root, label: 10}
-    - {path: samples/top/multitop/ttt.root, label: 11}
-  tree_name: output
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-    - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  scales: [0.1, 1, 1, 0.1, 1, 1, 1]
-  # Number is the event index; Number % 10 is the fold. Pretraining trained
-  # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests.
-  fold_var: Number
-  fold_modulus: 10
-  train_folds: [2, 3, 4, 5, 6, 7, 8, 9]
-  validation_folds: [0]
-  test_folds: [1]
-  weight_var: weight
-  cache: outputs/delphes_pretraining/events.pt
-  batch_size: 1024
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: edge_network
-  hid_size: 64
-  out_size: 12
-  n_layers: 4
-  n_proc_steps: 4
-
-task:
-  name: multiclass
-
-train:
-  max_epochs: 100
-  learning_rate: 1.0e-4
-  lr_decay: 0.99
-  early_stopping: {monitor: train_loss, patience: 5, min_delta: 0.0}
-  best_metric: auc

configs/delphes_tabicl.yaml

@@ -1,89 +0,0 @@
-# Meta-training of the TabICL network (src/gnn4colliders/tabular.py; BSD-3,
-# so meta-trained weights can be shared freely) on the classification tasks
-# of configs/delphes_pfn.yaml, on the fixed-width event rows of
-# configs/gamgam_tabicl.yaml. By default TabICL is frozen and its in-context
-# blocks learn LoRA updates (model.lora_rank; 0 trains all 28M weights).
-# task.prior_classes=[1,6,7,8,9,10,11] leaves out the five processes of the
-# GamGam task. Evaluate in context with
-#   gnn4colliders evaluate configs/gamgam_tabicl.yaml split=validation \
-#     checkpoint=outputs/delphes_tabicl/best.pt
-# Needs `uv sync --extra tabicl`.
-#
-#   gnn4colliders train configs/delphes_tabicl.yaml
-name: delphes_tabicl
-
-data:
-  hf_repo: HWresearch/Delphes
-  hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e
-  files:
-    - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0}
-    - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1}
-    - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2}
-    - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3}
-    - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4}
-    - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5}
-    - {path: samples/photon/ttyy/ttyy.root, label: 6}
-    - {path: samples/top/multitop/tttt.root, label: 7}
-    - {path: samples/top/single-top/SingleT_schan.root, label: 8}
-    - {path: samples/top/ttbar/ttbar.root, label: 9}
-    - {path: samples/top/ttv/ttW.root, label: 10}
-    - {path: samples/top/multitop/ttt.root, label: 11}
-  tree_name: output
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-    - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  scales: [0.1, 1, 1, 0.1, 1, 1, 1]
-  # Number is the event index; Number % 10 is the fold. Pretraining trained
-  # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests.
-  fold_var: Number
-  fold_modulus: 10
-  train_folds: [2, 3, 4, 5, 6, 7, 8, 9]
-  validation_folds: [0]
-  test_folds: [1]
-  # Tasks use the labels, not the weights; the cache is that of
-  # configs/delphes_pretraining.yaml.
-  weight_var: weight
-  cache: outputs/delphes_pretraining/events.pt
-  # Events per step (tasks of context and query events). TabICL (28M
-  # weights) recomputes activations in the backward pass.
-  batch_size: 4096
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: tabicl
-  tabicl_checkpoint: tabicl-classifier-v2-20260212.ckpt
-  # Leading jets, electrons, muons, photons and MET in a table row.
-  slots: [6, 2, 2, 3, 1]
-  # As in configs/gamgam_tabicl.yaml, which evaluates this model.
-  relative_phi: false
-  pair_objects: 0
-  max_classes: 10
-  lora_rank: 8
-
-task:
-  name: pfn
-  # The prior: classes per task, context events (log-uniform) and queries.
-  prior_classes: null
-  task_classes: [2, 10]
-  context_events: [16, 2048]
-  query_events: 256
-  energy_jitter: 0.2
-  steps_per_epoch: 500
-  validation_steps: 50
-
-train:
-  max_epochs: 20
-  optimizer: adamw
-  learning_rate: 1.0e-4
-  weight_decay: 0.01
-  lr_decay: 0.95
-  early_stopping: null
-  best_metric: loss

configs/delphes_tabpfn.yaml

@@ -1,89 +0,0 @@
-# Meta-training of the TabPFN-3.5 network (src/gnn4colliders/tabular.py) on
-# the classification tasks of configs/delphes_pfn.yaml, on the fixed-width
-# event rows of configs/gamgam_tabpfn.yaml. By default TabPFN is frozen and
-# its in-context blocks learn LoRA updates (model.lora_rank; 0 trains all
-# weights). task.prior_classes=[1,6,7,8,9,10,11] leaves out the five
-# processes of the GamGam task. Evaluate in context with
-#   gnn4colliders evaluate configs/gamgam_tabpfn.yaml split=validation \
-#     checkpoint=outputs/delphes_tabpfn/best.pt
-# Needs `uv sync --extra tabpfn` and the TabPFN-3.5 weights (docs/usage.md).
-#
-#   gnn4colliders train configs/delphes_tabpfn.yaml
-name: delphes_tabpfn
-
-data:
-  hf_repo: HWresearch/Delphes
-  hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e
-  files:
-    - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0}
-    - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1}
-    - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2}
-    - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3}
-    - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4}
-    - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5}
-    - {path: samples/photon/ttyy/ttyy.root, label: 6}
-    - {path: samples/top/multitop/tttt.root, label: 7}
-    - {path: samples/top/single-top/SingleT_schan.root, label: 8}
-    - {path: samples/top/ttbar/ttbar.root, label: 9}
-    - {path: samples/top/ttv/ttW.root, label: 10}
-    - {path: samples/top/multitop/ttt.root, label: 11}
-  tree_name: output
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-    - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  scales: [0.1, 1, 1, 0.1, 1, 1, 1]
-  # Number is the event index; Number % 10 is the fold. Pretraining trained
-  # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests.
-  fold_var: Number
-  fold_modulus: 10
-  train_folds: [2, 3, 4, 5, 6, 7, 8, 9]
-  validation_folds: [0]
-  test_folds: [1]
-  # Tasks use the labels, not the weights; the cache is that of
-  # configs/delphes_pretraining.yaml.
-  weight_var: weight
-  cache: outputs/delphes_pretraining/events.pt
-  # Events per step (tasks of context and query events): one task of the
-  # largest context. TabPFN recomputes activations in the backward pass.
-  batch_size: 2304
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: tabpfn
-  tabpfn_version: v3.5
-  tabpfn_checkpoint: null
-  # Leading jets, electrons, muons, photons and MET in a table row.
-  slots: [6, 2, 2, 3, 1]
-  # As in configs/gamgam_tabpfn.yaml, which evaluates this model.
-  relative_phi: false
-  pair_objects: 0
-  max_classes: 10
-  lora_rank: 8
-
-task:
-  name: pfn
-  # The prior: classes per task, context events (log-uniform) and queries.
-  prior_classes: null
-  task_classes: [2, 10]
-  context_events: [16, 2048]
-  query_events: 256
-  energy_jitter: 0.2
-  steps_per_epoch: 500
-  validation_steps: 50
-
-train:
-  max_epochs: 20
-  optimizer: adamw
-  learning_rate: 1.0e-4
-  weight_decay: 0.01
-  lr_decay: 0.95
-  early_stopping: null
-  best_metric: loss

configs/environment/.gitkeep

@@ -0,0 +1 @@
+

configs/gamgam.yaml

@@ -1,97 +0,0 @@
-# Higgs production-mode classification on ATLAS Open Data (13 TeV, 2020
-# release) GamGam MC: the 5-class from-scratch baseline of Ho et al.,
-# JINST 21 (2026) P08006, Tables 1 and 3. 2,473,335 events.
-#
-#   gnn4colliders prepare configs/gamgam.yaml
-#   gnn4colliders train configs/gamgam.yaml [train.fast=true] [train.seed=2]
-#   gnn4colliders evaluate configs/gamgam.yaml
-name: gamgam
-
-data:
-  hf_repo: ho22joshua/atlas_opendata
-  hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e
-  files:
-    - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root
-      label: 0
-      size_bytes: 217819076
-      sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7
-    - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root
-      label: 1
-      size_bytes: 218418347
-      sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb
-    - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root
-      label: 2
-      size_bytes: 108616148
-      sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2
-    - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root
-      label: 3
-      size_bytes: 29757932
-      sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248
-    - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root
-      label: 4
-      size_bytes: 58701804
-      sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748
-  class_names: [ttH, ggF, VBF, WH, ZH]
-  tree_name: mini
-  # Electrons and muons share the lep_* collection; b-tagging uses the
-  # MV2c10 70% efficiency working point of the Open Data documentation.
-  derived_branches:
-    ele_pt: {source: lep_pt, where: {lep_type: 11}}
-    ele_eta: {source: lep_eta, where: {lep_type: 11}}
-    ele_phi: {source: lep_phi, where: {lep_type: 11}}
-    ele_charge: {source: lep_charge, where: {lep_type: 11}}
-    mu_pt: {source: lep_pt, where: {lep_type: 13}}
-    mu_eta: {source: lep_eta, where: {lep_type: 13}}
-    mu_phi: {source: lep_phi, where: {lep_type: 13}}
-    mu_charge: {source: lep_charge, where: {lep_type: 13}}
-    jet_btag: {source: jet_MV2c10, above: 0.8244273}
-  # Object types (node types) as in pretraining: jet, electron, muon, photon, MET.
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, photon_pt, met_et]
-    - [jet_eta, ele_eta, mu_eta, photon_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV.
-  scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1]
-  # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split.
-  fold_var: eventNumber
-  fold_modulus: 10
-  train_folds: [0, 1, 2, 3, 4, 5, 6, 7]
-  validation_folds: [8]
-  test_folds: [9]
-  # Unweighted: each event counts once and classes are balanced by the
-  # per-class mean loss (23% of ttH events have mcWeight = -1).
-  weight_var: null
-  cache: outputs/gamgam/events.pt
-  batch_size: 1024
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-# Baseline GNN of Ho et al. (2026), Table 1.
-model:
-  name: edge_network
-  hid_size: 64
-  out_size: 5
-  n_layers: 4
-  n_proc_steps: 4
-  dropout: 0.0
-
-task:
-  name: multiclass
-
-# Sections 3.2 and 3.5: Adam at 1e-4 decayed by 0.99 per epoch, stopped
-# after 5 epochs without training-loss improvement; the epoch with the best
-# validation AUC is kept (best.pt). train.fast=true adds torch.compile, CUDA
-# graphs, bfloat16 and TF32 (about 3x faster on an A100; results agree to
-# rounding, not bitwise).
-train:
-  max_epochs: 200
-  learning_rate: 1.0e-4
-  lr_decay: 0.99
-  early_stopping: {monitor: train_loss, patience: 5, min_delta: 0.0}
-  best_metric: auc

configs/gamgam_deep_sets.yaml

@@ -1,88 +0,0 @@
-# The GamGam 5-class task with Deep Sets (a per-object MLP, summed): the
-# baseline without object interactions (no paper reference). Uses the cache
-# of configs/gamgam.yaml.
-#
-#   gnn4colliders train configs/gamgam_deep_sets.yaml
-name: gamgam_deep_sets
-
-data:
-  hf_repo: ho22joshua/atlas_opendata
-  hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e
-  files:
-    - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root
-      label: 0
-      size_bytes: 217819076
-      sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7
-    - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root
-      label: 1
-      size_bytes: 218418347
-      sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb
-    - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root
-      label: 2
-      size_bytes: 108616148
-      sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2
-    - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root
-      label: 3
-      size_bytes: 29757932
-      sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248
-    - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root
-      label: 4
-      size_bytes: 58701804
-      sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748
-  class_names: [ttH, ggF, VBF, WH, ZH]
-  tree_name: mini
-  # Electrons and muons share the lep_* collection; b-tagging uses the
-  # MV2c10 70% efficiency working point of the Open Data documentation.
-  derived_branches:
-    ele_pt: {source: lep_pt, where: {lep_type: 11}}
-    ele_eta: {source: lep_eta, where: {lep_type: 11}}
-    ele_phi: {source: lep_phi, where: {lep_type: 11}}
-    ele_charge: {source: lep_charge, where: {lep_type: 11}}
-    mu_pt: {source: lep_pt, where: {lep_type: 13}}
-    mu_eta: {source: lep_eta, where: {lep_type: 13}}
-    mu_phi: {source: lep_phi, where: {lep_type: 13}}
-    mu_charge: {source: lep_charge, where: {lep_type: 13}}
-    jet_btag: {source: jet_MV2c10, above: 0.8244273}
-  # Object types (node types) as in pretraining: jet, electron, muon, photon, MET.
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, photon_pt, met_et]
-    - [jet_eta, ele_eta, mu_eta, photon_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV.
-  scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1]
-  # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split.
-  fold_var: eventNumber
-  fold_modulus: 10
-  train_folds: [0, 1, 2, 3, 4, 5, 6, 7]
-  validation_folds: [8]
-  test_folds: [9]
-  # Unweighted: each event counts once and classes are balanced by the
-  # per-class mean loss (23% of ttH events have mcWeight = -1).
-  weight_var: null
-  cache: outputs/gamgam/events.pt
-  batch_size: 1024
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: deep_sets
-  hid_size: 64
-  out_size: 5
-  n_layers: 3
-  dropout: 0.0
-
-task:
-  name: multiclass
-
-train:
-  max_epochs: 100
-  learning_rate: 1.0e-4
-  lr_decay: 0.99
-  early_stopping: {monitor: train_loss, patience: 5, min_delta: 0.0}
-  best_metric: auc

configs/gamgam_finetune.yaml

@@ -1,97 +0,0 @@
-# The GamGam 5-class task fine-tuned from the multiclass pretrained model of
-# Ho et al. (2026): its edge network with a new 5-class classifier, all
-# weights trained. Uses the cache of configs/gamgam.yaml.
-#
-#   gnn4colliders train configs/gamgam_finetune.yaml
-name: gamgam_finetune
-
-data:
-  hf_repo: ho22joshua/atlas_opendata
-  hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e
-  files:
-    - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root
-      label: 0
-      size_bytes: 217819076
-      sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7
-    - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root
-      label: 1
-      size_bytes: 218418347
-      sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb
-    - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root
-      label: 2
-      size_bytes: 108616148
-      sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2
-    - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root
-      label: 3
-      size_bytes: 29757932
-      sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248
-    - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root
-      label: 4
-      size_bytes: 58701804
-      sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748
-  class_names: [ttH, ggF, VBF, WH, ZH]
-  tree_name: mini
-  # Electrons and muons share the lep_* collection; b-tagging uses the
-  # MV2c10 70% efficiency working point of the Open Data documentation.
-  derived_branches:
-    ele_pt: {source: lep_pt, where: {lep_type: 11}}
-    ele_eta: {source: lep_eta, where: {lep_type: 11}}
-    ele_phi: {source: lep_phi, where: {lep_type: 11}}
-    ele_charge: {source: lep_charge, where: {lep_type: 11}}
-    mu_pt: {source: lep_pt, where: {lep_type: 13}}
-    mu_eta: {source: lep_eta, where: {lep_type: 13}}
-    mu_phi: {source: lep_phi, where: {lep_type: 13}}
-    mu_charge: {source: lep_charge, where: {lep_type: 13}}
-    jet_btag: {source: jet_MV2c10, above: 0.8244273}
-  # Object types (node types) as in pretraining: jet, electron, muon, photon, MET.
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, photon_pt, met_et]
-    - [jet_eta, ele_eta, mu_eta, photon_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV.
-  scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1]
-  # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split.
-  fold_var: eventNumber
-  fold_modulus: 10
-  train_folds: [0, 1, 2, 3, 4, 5, 6, 7]
-  validation_folds: [8]
-  test_folds: [9]
-  # Unweighted: each event counts once and classes are balanced by the
-  # per-class mean loss (23% of ttH events have mcWeight = -1).
-  weight_var: null
-  cache: outputs/gamgam/events.pt
-  batch_size: 1024
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-# The pretrained network's size (Pretrained_GNN/*/config.yaml).
-model:
-  name: edge_network
-  pretrained: Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_71.pt
-  freeze_backbone: false
-  hid_size: 64
-  out_size: 5
-  n_layers: 4
-  n_proc_steps: 4
-  dropout: 0.0
-
-task:
-  name: multiclass
-
-# Sections 3.2 and 3.5: Adam at 1e-4 decayed by 0.99 per epoch, stopped
-# after 5 epochs without training-loss improvement; the epoch with the best
-# validation AUC is kept (best.pt). train.fast=true adds torch.compile, CUDA
-# graphs, bfloat16 and TF32 (about 3x faster on an A100; results agree to
-# rounding, not bitwise).
-train:
-  max_epochs: 200
-  learning_rate: 1.0e-4
-  lr_decay: 0.99
-  early_stopping: {monitor: train_loss, patience: 5, min_delta: 0.0}
-  best_metric: auc

configs/gamgam_jepa_finetune.yaml

@@ -1,95 +0,0 @@
-# The GamGam 5-class task with the transformer encoder pretrained by JEPA on
-# the Delphes samples (configs/delphes_jepa.yaml). model.pretrained=null
-# trains the same architecture from scratch. Uses the cache of
-# configs/gamgam.yaml.
-#
-#   gnn4colliders train configs/gamgam_jepa_finetune.yaml
-name: gamgam_jepa_finetune
-
-data:
-  hf_repo: ho22joshua/atlas_opendata
-  hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e
-  files:
-    - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root
-      label: 0
-      size_bytes: 217819076
-      sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7
-    - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root
-      label: 1
-      size_bytes: 218418347
-      sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb
-    - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root
-      label: 2
-      size_bytes: 108616148
-      sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2
-    - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root
-      label: 3
-      size_bytes: 29757932
-      sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248
-    - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root
-      label: 4
-      size_bytes: 58701804
-      sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748
-  class_names: [ttH, ggF, VBF, WH, ZH]
-  tree_name: mini
-  # Electrons and muons share the lep_* collection; b-tagging uses the
-  # MV2c10 70% efficiency working point of the Open Data documentation.
-  derived_branches:
-    ele_pt: {source: lep_pt, where: {lep_type: 11}}
-    ele_eta: {source: lep_eta, where: {lep_type: 11}}
-    ele_phi: {source: lep_phi, where: {lep_type: 11}}
-    ele_charge: {source: lep_charge, where: {lep_type: 11}}
-    mu_pt: {source: lep_pt, where: {lep_type: 13}}
-    mu_eta: {source: lep_eta, where: {lep_type: 13}}
-    mu_phi: {source: lep_phi, where: {lep_type: 13}}
-    mu_charge: {source: lep_charge, where: {lep_type: 13}}
-    jet_btag: {source: jet_MV2c10, above: 0.8244273}
-  # Object types (node types) as in pretraining: jet, electron, muon, photon, MET.
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, photon_pt, met_et]
-    - [jet_eta, ele_eta, mu_eta, photon_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV.
-  scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1]
-  # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split.
-  fold_var: eventNumber
-  fold_modulus: 10
-  train_folds: [0, 1, 2, 3, 4, 5, 6, 7]
-  validation_folds: [8]
-  test_folds: [9]
-  # Unweighted: each event counts once and classes are balanced by the
-  # per-class mean loss (23% of ttH events have mcWeight = -1).
-  weight_var: null
-  cache: outputs/gamgam/events.pt
-  batch_size: 1024
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: transformer
-  out_size: 5
-  hid_size: 128
-  n_layers: 4
-  n_heads: 8
-  ffn_size: 512
-  dropout: 0.0
-  pair_bias: true
-  pretrained: outputs/delphes_jepa/last.pt
-  # true trains only the classifier (linear probe).
-  freeze_backbone: false
-
-task:
-  name: multiclass
-
-train:
-  max_epochs: 100
-  learning_rate: 1.0e-4
-  lr_decay: 0.99
-  early_stopping: {monitor: train_loss, patience: 5, min_delta: 0.0}
-  best_metric: auc

configs/gamgam_pfn.yaml

@@ -1,90 +0,0 @@
-# The GamGam 5-class task classified in context by the PFN meta-trained on the
-# Delphes samples (configs/delphes_pfn.yaml), with no training on GamGam:
-# data.train_events labelled training events, an equal number per class, are
-# the context. The model section must match the meta-training. Uses the
-# cache of configs/gamgam.yaml.
-#
-#   gnn4colliders evaluate configs/gamgam_pfn.yaml split=validation
-name: gamgam_pfn
-checkpoint: outputs/delphes_pfn/best.pt
-
-data:
-  hf_repo: ho22joshua/atlas_opendata
-  hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e
-  files:
-    - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root
-      label: 0
-      size_bytes: 217819076
-      sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7
-    - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root
-      label: 1
-      size_bytes: 218418347
-      sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb
-    - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root
-      label: 2
-      size_bytes: 108616148
-      sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2
-    - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root
-      label: 3
-      size_bytes: 29757932
-      sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248
-    - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root
-      label: 4
-      size_bytes: 58701804
-      sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748
-  class_names: [ttH, ggF, VBF, WH, ZH]
-  tree_name: mini
-  # Electrons and muons share the lep_* collection; b-tagging uses the
-  # MV2c10 70% efficiency working point of the Open Data documentation.
-  derived_branches:
-    ele_pt: {source: lep_pt, where: {lep_type: 11}}
-    ele_eta: {source: lep_eta, where: {lep_type: 11}}
-    ele_phi: {source: lep_phi, where: {lep_type: 11}}
-    ele_charge: {source: lep_charge, where: {lep_type: 11}}
-    mu_pt: {source: lep_pt, where: {lep_type: 13}}
-    mu_eta: {source: lep_eta, where: {lep_type: 13}}
-    mu_phi: {source: lep_phi, where: {lep_type: 13}}
-    mu_charge: {source: lep_charge, where: {lep_type: 13}}
-    jet_btag: {source: jet_MV2c10, above: 0.8244273}
-  # Object types (node types) as in pretraining: jet, electron, muon, photon, MET.
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, photon_pt, met_et]
-    - [jet_eta, ele_eta, mu_eta, photon_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV.
-  scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1]
-  # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split.
-  fold_var: eventNumber
-  fold_modulus: 10
-  train_folds: [0, 1, 2, 3, 4, 5, 6, 7]
-  validation_folds: [8]
-  test_folds: [9]
-  # Labelled events (equal per class) that form the context.
-  train_events: 1000
-  # Unweighted: each event counts once and classes are balanced by the
-  # per-class mean loss (23% of ttH events have mcWeight = -1).
-  weight_var: null
-  cache: outputs/gamgam/events.pt
-  batch_size: 1024
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: pfn
-  out_size: 5
-  hid_size: 128
-  n_layers: 4
-  n_heads: 8
-  ffn_size: 512
-  pair_bias: true
-  max_classes: 10
-  pfn_layers: 6
-
-task:
-  name: multiclass

configs/gamgam_tabdpt.yaml

@@ -1,91 +0,0 @@
-# The GamGam 5-class task classified in context by TabDPT (Layer 6;
-# Apache-2.0 weights from Hugging Face), with no training: data.train_events
-# labelled training events, an equal number per class, are the context of
-# every prediction (src/gnn4colliders/tabular.py). TabDPT is trained on real
-# tables; with model.tabdpt_context_size each event's context is its nearest
-# context events (retrieval), which allows contexts beyond what fits at once.
-# Needs `uv sync --extra tabdpt`. Uses the cache of configs/gamgam.yaml.
-#
-#   gnn4colliders evaluate configs/gamgam_tabdpt.yaml split=validation
-name: gamgam_tabdpt
-
-data:
-  hf_repo: ho22joshua/atlas_opendata
-  hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e
-  files:
-    - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root
-      label: 0
-      size_bytes: 217819076
-      sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7
-    - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root
-      label: 1
-      size_bytes: 218418347
-      sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb
-    - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root
-      label: 2
-      size_bytes: 108616148
-      sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2
-    - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root
-      label: 3
-      size_bytes: 29757932
-      sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248
-    - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root
-      label: 4
-      size_bytes: 58701804
-      sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748
-  class_names: [ttH, ggF, VBF, WH, ZH]
-  tree_name: mini
-  # Electrons and muons share the lep_* collection; b-tagging uses the
-  # MV2c10 70% efficiency working point of the Open Data documentation.
-  derived_branches:
-    ele_pt: {source: lep_pt, where: {lep_type: 11}}
-    ele_eta: {source: lep_eta, where: {lep_type: 11}}
-    ele_phi: {source: lep_phi, where: {lep_type: 11}}
-    ele_charge: {source: lep_charge, where: {lep_type: 11}}
-    mu_pt: {source: lep_pt, where: {lep_type: 13}}
-    mu_eta: {source: lep_eta, where: {lep_type: 13}}
-    mu_phi: {source: lep_phi, where: {lep_type: 13}}
-    mu_charge: {source: lep_charge, where: {lep_type: 13}}
-    jet_btag: {source: jet_MV2c10, above: 0.8244273}
-  # Object types (node types) as in pretraining: jet, electron, muon, photon, MET.
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, photon_pt, met_et]
-    - [jet_eta, ele_eta, mu_eta, photon_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV.
-  scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1]
-  # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split.
-  fold_var: eventNumber
-  fold_modulus: 10
-  train_folds: [0, 1, 2, 3, 4, 5, 6, 7]
-  validation_folds: [8]
-  test_folds: [9]
-  # Labelled events (equal per class) that form the context.
-  train_events: 1000
-  # Unweighted: each event counts once and classes are balanced by the
-  # per-class mean loss (23% of ttH events have mcWeight = -1).
-  weight_var: null
-  cache: outputs/gamgam/events.pt
-  # Events classified per TabDPT call.
-  batch_size: 8192
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: tabdpt
-  out_size: 5
-  tabdpt_estimators: 8
-  tabdpt_context_size: null
-  # Leading jets, electrons, muons, photons and MET in a table row.
-  slots: [6, 2, 2, 3, 1]
-  relative_phi: false
-  pair_objects: 0
-
-task:
-  name: multiclass

configs/gamgam_tabicl.yaml

@@ -1,95 +0,0 @@
-# The GamGam 5-class task classified in context by TabICLv2 (Inria; BSD-3
-# weights from Hugging Face), with no training: data.train_events labelled
-# training events, an equal number per class, are the context of every
-# prediction (src/gnn4colliders/tabular.py). Needs `uv sync --extra tabicl`.
-# Uses the cache of configs/gamgam.yaml. The rows are those of
-# configs/gamgam_tabpfn.yaml; model.relative_phi=true and
-# model.pair_objects=4 add the physics features.
-#
-#   gnn4colliders evaluate configs/gamgam_tabicl.yaml split=validation
-#
-# With checkpoint=outputs/delphes_tabicl/best.pt, the TabICL network
-# meta-trained by configs/delphes_tabicl.yaml (model.lora_rank and
-# model.max_classes must match it) classifies with its own forward pass.
-name: gamgam_tabicl
-
-data:
-  hf_repo: ho22joshua/atlas_opendata
-  hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e
-  files:
-    - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root
-      label: 0
-      size_bytes: 217819076
-      sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7
-    - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root
-      label: 1
-      size_bytes: 218418347
-      sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb
-    - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root
-      label: 2
-      size_bytes: 108616148
-      sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2
-    - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root
-      label: 3
-      size_bytes: 29757932
-      sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248
-    - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root
-      label: 4
-      size_bytes: 58701804
-      sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748
-  class_names: [ttH, ggF, VBF, WH, ZH]
-  tree_name: mini
-  # Electrons and muons share the lep_* collection; b-tagging uses the
-  # MV2c10 70% efficiency working point of the Open Data documentation.
-  derived_branches:
-    ele_pt: {source: lep_pt, where: {lep_type: 11}}
-    ele_eta: {source: lep_eta, where: {lep_type: 11}}
-    ele_phi: {source: lep_phi, where: {lep_type: 11}}
-    ele_charge: {source: lep_charge, where: {lep_type: 11}}
-    mu_pt: {source: lep_pt, where: {lep_type: 13}}
-    mu_eta: {source: lep_eta, where: {lep_type: 13}}
-    mu_phi: {source: lep_phi, where: {lep_type: 13}}
-    mu_charge: {source: lep_charge, where: {lep_type: 13}}
-    jet_btag: {source: jet_MV2c10, above: 0.8244273}
-  # Object types (node types) as in pretraining: jet, electron, muon, photon, MET.
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, photon_pt, met_et]
-    - [jet_eta, ele_eta, mu_eta, photon_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV.
-  scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1]
-  # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split.
-  fold_var: eventNumber
-  fold_modulus: 10
-  train_folds: [0, 1, 2, 3, 4, 5, 6, 7]
-  validation_folds: [8]
-  test_folds: [9]
-  # Labelled events (equal per class) that form the context.
-  train_events: 1000
-  # Unweighted: each event counts once and classes are balanced by the
-  # per-class mean loss (23% of ttH events have mcWeight = -1).
-  weight_var: null
-  cache: outputs/gamgam/events.pt
-  # Events classified per TabICL call.
-  batch_size: 8192
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: tabicl
-  out_size: 5
-  tabicl_checkpoint: tabicl-classifier-v2-20260212.ckpt
-  tabicl_estimators: 8
-  # Leading jets, electrons, muons, photons and MET in a table row.
-  slots: [6, 2, 2, 3, 1]
-  relative_phi: false
-  pair_objects: 0
-
-task:
-  name: multiclass

configs/gamgam_tabpfn.yaml

@@ -1,96 +0,0 @@
-# The GamGam 5-class task classified in context by TabPFN (Prior Labs), with
-# no training: data.train_events labelled training events, an equal number
-# per class, are the context of every prediction (src/gnn4colliders/tabular.py).
-# Needs `uv sync --extra tabpfn` and, for TabPFN-2.5 and later, a one-time
-# login (a token in ~/.cache/tabpfn or TABPFN_TOKEN). Uses the cache of
-# configs/gamgam.yaml.
-#
-#   gnn4colliders evaluate configs/gamgam_tabpfn.yaml split=validation
-#
-# With checkpoint=outputs/delphes_tabpfn/best.pt, the TabPFN network
-# meta-trained by configs/delphes_tabpfn.yaml (model.lora_rank and
-# model.max_classes must match it) classifies with its own forward pass.
-name: gamgam_tabpfn
-
-data:
-  hf_repo: ho22joshua/atlas_opendata
-  hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e
-  files:
-    - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root
-      label: 0
-      size_bytes: 217819076
-      sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7
-    - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root
-      label: 1
-      size_bytes: 218418347
-      sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb
-    - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root
-      label: 2
-      size_bytes: 108616148
-      sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2
-    - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root
-      label: 3
-      size_bytes: 29757932
-      sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248
-    - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root
-      label: 4
-      size_bytes: 58701804
-      sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748
-  class_names: [ttH, ggF, VBF, WH, ZH]
-  tree_name: mini
-  # Electrons and muons share the lep_* collection; b-tagging uses the
-  # MV2c10 70% efficiency working point of the Open Data documentation.
-  derived_branches:
-    ele_pt: {source: lep_pt, where: {lep_type: 11}}
-    ele_eta: {source: lep_eta, where: {lep_type: 11}}
-    ele_phi: {source: lep_phi, where: {lep_type: 11}}
-    ele_charge: {source: lep_charge, where: {lep_type: 11}}
-    mu_pt: {source: lep_pt, where: {lep_type: 13}}
-    mu_eta: {source: lep_eta, where: {lep_type: 13}}
-    mu_phi: {source: lep_phi, where: {lep_type: 13}}
-    mu_charge: {source: lep_charge, where: {lep_type: 13}}
-    jet_btag: {source: jet_MV2c10, above: 0.8244273}
-  # Object types (node types) as in pretraining: jet, electron, muon, photon, MET.
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, photon_pt, met_et]
-    - [jet_eta, ele_eta, mu_eta, photon_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV.
-  scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1]
-  # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split.
-  fold_var: eventNumber
-  fold_modulus: 10
-  train_folds: [0, 1, 2, 3, 4, 5, 6, 7]
-  validation_folds: [8]
-  test_folds: [9]
-  # Labelled events (equal per class) that form the context.
-  train_events: 1000
-  # Unweighted: each event counts once and classes are balanced by the
-  # per-class mean loss (23% of ttH events have mcWeight = -1).
-  weight_var: null
-  cache: outputs/gamgam/events.pt
-  # Events classified per TabPFN call.
-  batch_size: 8192
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: tabpfn
-  out_size: 5
-  tabpfn_version: v3.5
-  tabpfn_estimators: 4
-  # Leading jets, electrons, muons, photons and MET in a table row.
-  slots: [6, 2, 2, 3, 1]
-  # Physics features: azimuths from the leading object's, and the mass and
-  # Delta R of every pair of the leading pair_objects objects.
-  relative_phi: false
-  pair_objects: 0
-
-task:
-  name: multiclass

configs/gamgam_transformer.yaml

@@ -1,92 +0,0 @@
-# The GamGam 5-class task with a transformer over the same objects (no
-# paper reference). Uses the cache of configs/gamgam.yaml.
-#
-#   gnn4colliders train configs/gamgam_transformer.yaml
-name: gamgam_transformer
-
-data:
-  hf_repo: ho22joshua/atlas_opendata
-  hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e
-  files:
-    - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root
-      label: 0
-      size_bytes: 217819076
-      sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7
-    - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root
-      label: 1
-      size_bytes: 218418347
-      sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb
-    - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root
-      label: 2
-      size_bytes: 108616148
-      sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2
-    - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root
-      label: 3
-      size_bytes: 29757932
-      sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248
-    - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root
-      label: 4
-      size_bytes: 58701804
-      sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748
-  class_names: [ttH, ggF, VBF, WH, ZH]
-  tree_name: mini
-  # Electrons and muons share the lep_* collection; b-tagging uses the
-  # MV2c10 70% efficiency working point of the Open Data documentation.
-  derived_branches:
-    ele_pt: {source: lep_pt, where: {lep_type: 11}}
-    ele_eta: {source: lep_eta, where: {lep_type: 11}}
-    ele_phi: {source: lep_phi, where: {lep_type: 11}}
-    ele_charge: {source: lep_charge, where: {lep_type: 11}}
-    mu_pt: {source: lep_pt, where: {lep_type: 13}}
-    mu_eta: {source: lep_eta, where: {lep_type: 13}}
-    mu_phi: {source: lep_phi, where: {lep_type: 13}}
-    mu_charge: {source: lep_charge, where: {lep_type: 13}}
-    jet_btag: {source: jet_MV2c10, above: 0.8244273}
-  # Object types (node types) as in pretraining: jet, electron, muon, photon, MET.
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, photon_pt, met_et]
-    - [jet_eta, ele_eta, mu_eta, photon_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV.
-  scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1]
-  # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split.
-  fold_var: eventNumber
-  fold_modulus: 10
-  train_folds: [0, 1, 2, 3, 4, 5, 6, 7]
-  validation_folds: [8]
-  test_folds: [9]
-  # Unweighted: each event counts once and classes are balanced by the
-  # per-class mean loss (23% of ttH events have mcWeight = -1).
-  weight_var: null
-  cache: outputs/gamgam/events.pt
-  batch_size: 1024
-  prepare_workers: 8
-  loader_workers: 4
-  pin_memory: true
-
-model:
-  name: transformer
-  hid_size: 64
-  out_size: 5
-  n_layers: 2
-  n_heads: 4
-  ffn_size: 256
-  dropout: 0.1
-  # Attention bias from object-pair (deta, dphi, dR, mass) features, as in the
-  # Particle Transformer; false gives a plain transformer.
-  pair_bias: true
-
-task:
-  name: multiclass
-
-train:
-  max_epochs: 100
-  learning_rate: 1.0e-4
-  lr_decay: 0.99
-  early_stopping: {monitor: train_loss, patience: 5, min_delta: 0.0}
-  best_metric: auc

configs/model/.gitkeep

@@ -0,0 +1 @@
+

configs/model/root_gnn/.gitkeep

@@ -0,0 +1 @@
+

configs/task/.gitkeep

@@ -0,0 +1 @@
+

configs/trainer/.gitkeep

@@ -0,0 +1 @@
+

configs/tth_cp_even_odd.yaml

@@ -1,49 +0,0 @@
-# CP-even versus CP-odd top-Yukawa ttH (H -> gamma gamma), Delphes samples.
-# Point data.files at local copies of the two samples.
-name: tth_cp_even_odd
-
-data:
-  files:
-    - {path: data/raw/ttH_NLO_download/samples/higgs/top-associated/tth/ttH_NLO.root, label: 0}
-    - {path: data/raw/ttH_cp_even_odd/samples/higgs/top-associated/tth/ttH_CPodd.root, label: 1}
-  tree_name: output
-  feature_branches:
-    - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-    - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-    - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-    - CALC_E
-    - [jet_btag, 0, 0, 0, 0]
-    - [0, ele_charge, mu_charge, 0, 0]
-    - NODE_TYPE
-  object_types: [vector, vector, vector, vector, single]
-  # Delphes momenta are in GeV; 0.1 gives the pretraining unit of 10 GeV.
-  scales: [0.1, 1, 1, 0.1, 1, 1, 1]
-  # Number is the event index; Number % 10 is the fold.
-  fold_var: Number
-  fold_modulus: 10
-  train_folds: [0, 1, 2, 3, 4, 5, 6, 7]
-  validation_folds: [8]
-  test_folds: [9]
-  weight_var: weight
-  cache: outputs/tth_cp_even_odd/events.pt
-  batch_size: 1024
-
-model:
-  name: edge_network
-  hid_size: 64
-  out_size: 1
-  n_layers: 4
-  n_proc_steps: 4
-
-# NLO samples have signed weights; |weight| keeps the loss denominator
-# stable (metrics use events with positive weight).
-task:
-  name: binary
-  absolute_weights: true
-
-train:
-  max_epochs: 100
-  learning_rate: 1.0e-4
-  lr_decay: 0.99
-  early_stopping: {monitor: train_loss, patience: 5, min_delta: 0.0}
-  best_metric: auc

data/fixtures/.gitkeep

@@ -0,0 +1 @@
+

data/raw/.gitkeep

@@ -0,0 +1 @@
+

docs/architecture.md

@@ -0,0 +1,378 @@
+# GNN4Colliders architecture
+
+## Current v1 architecture
+
+The supported rewrite is layered around an architecture-neutral event boundary:
+
+```text
+ROOT/Awkward
+    ↓
+EventSample + EventMetadata
+    ↓
+shared collider feature construction
+    ├── GraphSample -> versioned graph cache -> GraphBatch
+    └── future SequenceSample -> ROOT-Transformer (not implemented)
+                                      ↓
+                               ROOT-GNN EdgeNetwork
+                                      ↓
+                               raw logits -> Task
+                                  ├── loss
+                                  ├── predictions
+                                  └── full-split metrics
+                                      ↓
+                          Trainer / Predictor / outputs
+```
+
+| Layer | Responsibility |
+| --- | --- |
+| `data` | ROOT/Awkward ingestion, event samples, metadata, graph caches, folds, and batching |
+| `features` | Shared collider-object features and derived physics quantities |
+| `graphs` | Topology, edge features, and the DGL representation adapter |
+| `models/root_gnn` | ROOT-GNN encoders, portable tensor message passing, classifier, and transfer boundary |
+| `tasks` | Loss, score/prediction, labels, weights, and metrics |
+| `training` | Optimizer lifecycle, validation, early stopping, checkpointing, and reproducibility |
+| `inference` | Ordered prediction/evaluation and NPZ/ROOT output adapters |
+| `distributed` | Rank-local devices, sharding, DDP, and cross-rank collection |
+| `config` / `cli` | Semantic Hydra composition and thin user-facing commands |
+
+`EventSample` is shared infrastructure, not a ROOT-GNN object. `GraphSample`
+is the current representation-specific adapter. At batching, the default
+ROOT-GNN backend converts cached DGL graphs into `TensorGraph` values with
+explicit node/edge indices and graph membership. Native message passing then
+uses only PyTorch tensor operations, allowing the same model to train on CPU,
+CUDA, and Apple MPS. This separation is the extension point for a future
+sequence/token representation.
+
+The new public metadata contract is named `EventMetadata(fold, weight,
+sample_id, extra)`. The legacy positional tracking tensor is accepted only by
+compatibility-facing ingestion code. A `GraphSampleCache` is deliberately a
+Level-2 graph cache; replacing it with a universal cache would couple future
+model families to DGL.
+
+For deployment, a prepared `GraphBatch` can pass through the isolated
+`RootGNNExportAdapter` into an ONNX model. ONNX is still an inference boundary:
+it does not read ROOT or construct collider features.
+
+## Current public workflow
+
+`prepare` reads ROOT through `RootEventDataset`, builds shared features and
+DGL graphs, and saves a schema-checked cache. `train` creates a model/task and
+`Trainer`; validation is the model-selection split and test is held out.
+`evaluate` computes metrics after collecting the complete split. `predict`
+returns detached CPU tensors in loader order and writes named NPZ fields.
+`write_root_scores` is an optional Python adapter with explicit entry alignment;
+the CLI currently exposes NPZ output.
+
+Compatibility responsibilities are isolated in `gnn4colliders.compat`.
+Supported historical checkpoint prefixes, classifier names, and the two-column
+tracking conversion are listed in [`compatibility.md`](compatibility.md).
+The modern pipeline does not propagate positional tracking or historical NPZ
+fields.
+
+Checkpoints are independent of the model implementation: they carry model and
+task metadata, lifecycle state, schema versions, and optional RNG state.
+Prefix normalization supports DDP `module.` and compiled `_orig_mod.` weights,
+plus the active ROOT-GNN historical classifier-name compatibility path.
+
+---
+
+## Historical behavioral reference
+
+This document covers the target system in `legacy/root_gnn_dgl/`. The sibling
+`legacy/physicsnemo/` tree is a prior rewrite attempt and is not a behavioral
+target.
+
+## 1. High-level system description
+
+The active rewrite exposes `gnn4colliders.models.root_gnn.EdgeNetwork`. Its
+encoders and message-passing blocks form a reusable backbone whose decoded
+graph representation is passed to an explicit classifier. `FineTunedEdgeNetwork`
+reuses that backbone and replaces only the task-specific classifier, with
+explicit frozen or trainable-backbone control.
+
+`root_gnn_dgl` is a ROOT-to-DGL graph classification system. YAML selects
+dataset, model, loss, and finish-function classes by import path. The dataset
+reads ROOT trees, converts collider objects to fully connected DGL graphs, and
+saves graph chunks. Training loads those chunks, applies fold selection and
+optional pre-batching/padding, trains a graph network, writes one PyTorch
+checkpoint per epoch, and reports weighted loss, accuracy, and ROC AUC.
+
+The primary model is `models.GCN.Edge_Network` ([`GCN.py:182-251`](../legacy/root_gnn_dgl/models/GCN.py)).
+It encodes node, edge, and global features, repeats edge -> node -> global
+message passing `n_proc_steps` times, decodes the global state, and applies
+`classify`. Fine-tuning uses `models.GCN.Transferred_Learning_Finetuning`
+([`GCN.py:884-997`](../legacy/root_gnn_dgl/models/GCN.py)), which loads a
+pretrained `Edge_Network`, removes its final classifier, and applies a new one.
+The active configs use output size 12 for multiclass pretraining and output
+size 1 for binary tasks ([`configs/stats_100K/pretraining_multiclass.yaml:1-45`](../legacy/root_gnn_dgl/configs/stats_100K/pretraining_multiclass.yaml),
+[`configs/stats_100K/finetuning_ttH_CP_even_vs_odd.yaml:1-45`](../legacy/root_gnn_dgl/configs/stats_100K/finetuning_ttH_CP_even_vs_odd.yaml)).
+
+### Entry points and flows
+
+- `scripts/training_script.py:main` and its CLI parser load YAML, create
+  loaders, construct the model, and call `train`; `--evaluate` calls
+  `evaluate` ([`training_script.py:638-843`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+- `scripts/prep_data.py:main` creates configured graph caches
+  ([`prep_data.py:68-110`](../legacy/root_gnn_dgl/scripts/prep_data.py)).
+- `scripts/inference.py:main` reconstructs an unlazy dataset, loads one or
+  more checkpoints, and writes `.npz` or ROOT scores
+  ([`inference.py:163-387`](../legacy/root_gnn_dgl/scripts/inference.py)).
+- `scripts/export_onnx.py:main` exports an ONNX-friendly model
+  ([`export_onnx.py:979-1035`](../legacy/root_gnn_dgl/scripts/export_onnx.py)).
+- `selections.py:main`, `check_dataset_files.py:main`, and
+  `plot_config_distributions.py:main` are diagnostic entry points. `run_demo.sh`
+  sequences pretraining, binary training, fine-tuning, and inference
+  ([`run_demo.sh:3-59`](../legacy/root_gnn_dgl/run_demo.sh)).
+
+The training flow is:
+
+```text
+YAML -> load_config/buildFromConfig -> RootDataset/LazyDataset
+     -> ROOT/Awkward -> DGL graph + labels/tracking/globals -> .bin cache
+     -> fold_selection -> prebatch/padding -> GraphDataLoader
+     -> Edge_Network or transfer model -> weighted loss/metrics
+     -> model_epoch_N.pt, logs, evaluation/inference output
+```
+
+`training_script.train` is the lifecycle implementation
+([`training_script.py:143-614`](../legacy/root_gnn_dgl/scripts/training_script.py));
+distributed paths use NCCL/DDP ([`training_script.py:616-839`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+Inference uses `CustomPreBatchedDataset`, applies a configured finish function,
+and collects `scores`, `labels`, and `tracking_info`
+([`inference.py:20-76`](../legacy/root_gnn_dgl/scripts/inference.py),
+[`inference.py:223-325`](../legacy/root_gnn_dgl/scripts/inference.py)).
+
+## 2. Dependency and data-flow map
+
+```text
+ROOT files (raw_dir/file_names, tree_name)
+  -> gnn4colliders.data.RootEventDataset (Uproot/Awkward)
+  -> node feature construction (dataset.py:15-50)
+  -> full_connected_graph (dataset.py:52-59)
+  -> EdgeDataset.make_graph: [deta, dphi, dR] (dataset.py:471-482)
+  -> DGL .bin cache / LazyDataset / PreBatchedDataset
+  -> GraphDataLoader -> models.GCN -> loss/metrics -> outputs
+```
+
+`RootDataset` provides `process`, `save`, `load`, `__getitem__`, and `__len__`
+([`dataset.py:160-469`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py));
+`LazyDataset` loads one chunk through a ring buffer
+([`dataset.py:525-578`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py));
+`PreBatchedDataset.process` selects, shuffles, batches, pads, and caches
+([`batched_dataset.py:34-146`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
+
+`load_config` uses PyYAML `FullLoader` and shallow `include` merging, while
+`buildFromConfig` dynamically imports `module`, resolves `class`, merges extra
+keys into `args`, converts list-valued weights to tensors, and injects runtime
+arguments ([`utils.py:10-43`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)).
+This reflection shape is a de facto interface for configured components.
+
+### Data and preprocessing
+
+The active node schema is seven columns: `pt`, `eta`, `phi`, `energy`, `btag`,
+`charge`, and `node_type`. `CALC_E` is `pt*cosh(eta)`, constants are broadcast
+per object type, `NODE_TYPE` is an integer type code, and feature scales are
+applied columnwise ([`dataset.py:15-50`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)).
+`full_connected_graph` makes directed all-pairs edges; `EdgeDataset` requests
+no self-loops and stores `[deta, dphi, dR]`
+([`dataset.py:52-59`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
+[`dataset.py:471-482`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)).
+
+Selections are strings evaluated with builtins disabled or
+`(variable, cut, operator)` triples (`check_selection`, `selection_mask`;
+[`dataset.py:75-145`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)). Fold
+selection uses `tracking[:,0] % n_folds`; tracking column 0 is fold and column 1
+is weight ([`utils.py:121-143`](../legacy/root_gnn_dgl/root_gnn_base/utils.py),
+[`dataset.py:176-182`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)).
+`hash_partition` and a seeded Torch generator control pre-batch order
+([`batched_dataset.py:27-99`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
+Padding modes are `NONE`, `STEPS`, `FIXED`, and `NODE`; `FIXED` is hardcoded to
+16,000 nodes and 104,000 edges ([`batched_dataset.py:100-125`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
+
+### Model, losses, and metrics
+
+`Make_MLP` builds linear/ReLU/dropout blocks followed by LayerNorm
+([`GCN.py:18-35`](../legacy/root_gnn_dgl/models/GCN.py)). Each `Edge_Network`
+step encodes inputs, copies source/destination states to edges, updates edges,
+sums edge messages into nodes, updates nodes, then mean-pools nodes/edges to
+update globals ([`GCN.py:195-249`](../legacy/root_gnn_dgl/models/GCN.py)). It
+returns logits `[graphs, out_size]` without sigmoid/softmax.
+
+The default objective is elementwise `BCEWithLogitsLoss`, multiplied by
+`tracking[:,1]`, averaged separately per unique label, then averaged across
+labels ([`training_script.py:143-185`](../legacy/root_gnn_dgl/scripts/training_script.py),
+[`training_script.py:320-359`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+`--abs` makes weights positive. Binary metrics use sigmoid threshold 0.5 and
+weighted ROC AUC; multiclass metrics use argmax and one-vs-rest ROC AUC
+([`training_script.py:438-510`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+Additional configurable losses/finishers live in `models/loss.py`
+([`loss.py:6-310`](../legacy/root_gnn_dgl/models/loss.py)).
+
+### Checkpoints and outputs
+
+Training writes `Training_Directory/model_epoch_<epoch>.pt` containing `epoch`,
+`model_state_dict`, `optimizer_state_dict`, and serialized `early_stop`
+([`training_script.py:565-604`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+Keys strip `module.` and compiled models save the underlying `_orig_mod` state;
+`get_last_epoch`, `get_specific_epoch`, and `get_best_epoch` load the files
+([`utils.py:145-248`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)).
+`evaluate` writes `evaluation_<epoch>.npz`; inference writes `.npz` fields
+`scores`, `labels`, `tracking_info`, or adds score branches and `selection_pass`
+to a cloned ROOT tree ([`training_script.py:57-140`](../legacy/root_gnn_dgl/scripts/training_script.py),
+[`inference.py:328-385`](../legacy/root_gnn_dgl/scripts/inference.py)).
+
+### Training lifecycle boundary
+
+The active rewrite keeps lifecycle orchestration architecture-independent:
+
+```text
+GraphDataLoader -> GraphBatch -> Model -> Task -> Trainer
+                                      loss/metrics
+```
+
+`gnn4colliders.training.Trainer` owns device placement, train/evaluation mode,
+gradient and optimizer steps, epoch aggregation, optional scheduler stepping,
+early stopping, and in-memory history. Tasks own loss and metric semantics;
+the trainer does not inspect positional tracking columns or collider-specific
+features. Evaluation concatenates detached outputs across the complete split
+before calling task metrics, so ROC AUC is not computed per mini-batch.
+
+Checkpoint persistence and the Python inference/output layer are implemented
+as separate adapters. The semantic CLI and distributed application boundary
+are implemented in the current stack. `gnn4colliders.inference.Predictor` accumulates detached CPU
+logits, task-defined scores/predictions, labels, and named event metadata in
+loader order; `write_npz` is the primary named-field format and ROOT score
+writing is an optional alignment-aware adapter.
+
+The new lifecycle uses conventional split semantics. `train` updates model
+parameters, `validation` is evaluated after every epoch and drives scheduler,
+early-stopping, and later model selection, and `test` is held out. The trainer
+does not accept a test loader in `fit`; callers evaluate the held-out test set
+separately after training. This deliberately corrects the legacy convention
+where a loader named `test` was used for model selection and `val` represented
+held-out testing.
+
+## Distributed execution
+
+`gnn4colliders.distributed` contains the small DDP boundary used by the
+application layer. `DistributedContext` reads the standard `torchrun`
+environment (`RANK`, `LOCAL_RANK`, and `WORLD_SIZE`), selects the rank-local
+device, and owns process-group cleanup. Graph samples are sharded before
+batching; training may pad rank shards for equal step counts, while validation
+and prediction use unpadded shards so events are not counted twice.
+
+The configured graph `batch_size` is per process. DDP wraps an otherwise
+ordinary model after device placement, and checkpoint state is normalized to
+the underlying model keys. Loss gradients are synchronized by DDP; epoch
+metrics and evaluation outputs are gathered across ranks. Rank 0 writes
+resolved configuration, checkpoints, and NPZ predictions. Moderate-size
+prediction gathering is in-memory; streaming/sharded output is a future
+extension.
+
+
+## Randomness, external services, and coupling
+
+The CLI exposes `--seed`, but `main` passes it to model construction rather
+than globally seeding Python, NumPy, or Torch
+([`training_script.py:638-753`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+Pre-batching has explicit seeds, but `AugmentedDataset` mutates the process-wide
+NumPy seed ([`dataset.py:716-827`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)),
+clustering uses unseeded `torch.randperm`/`randint` (`loss.py:259-295`), and
+model reset/fine-tuning hardcodes `torch.manual_seed(2)`
+([`GCN.py:58-65`](../legacy/root_gnn_dgl/models/GCN.py),
+[`GCN.py:900-915`](../legacy/root_gnn_dgl/models/GCN.py)). CUDA kernels, DDP,
+and DataLoader behavior are not made deterministic.
+
+Implicit coupling includes repository-relative `sys.path` insertion
+([`training_script.py:14-20`](../legacy/root_gnn_dgl/scripts/training_script.py)),
+dynamic imports, mutable default lists/dicts, global `FEATURE_DTYPE`
+([`dataset.py:13`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)),
+in-place `tracking_info` mutation ([`dataset.py:176-181`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)),
+and in-place DGL graph mutation during forward.
+
+The legacy environment assumes Python 3.8, PyTorch 2.0.1, CUDA 11.8, DGL
+1.1.1, ROOT, Awkward, Uproot, PyYAML, and scikit-learn
+([`setup/environment.yml:1-10`](../legacy/root_gnn_dgl/setup/environment.yml),
+[`setup/environment.yml:240-295`](../legacy/root_gnn_dgl/setup/environment.yml)).
+The active Linux development environment is intentionally separate: Python
+3.12, PyTorch 2.2.2/CUDA 12.1, and DGL 2.4.0 from the official DGL wheel
+repository. CUDA runtime wheels do not replace the compatible host NVIDIA
+driver and do not encode Perlmutter module settings.
+Standard configs assume `/global/cfs/` and `/pscratch/` paths, CUDA/NCCL,
+Slurm, and optionally Podman-HPC. `setup/download_data.sh` downloads the
+external Hugging Face dataset `HWresearch/Delphes`
+([`download_data.sh:13-67`](../legacy/root_gnn_dgl/setup/download_data.sh)).
+
+## Apparent unused or secondary code
+
+Not selected by the standard stats/Delphes configs, or only reachable from
+optional workflows, are `GCN_global`, `GCN_global_2way`, most transfer variants,
+attention models, `MultiModel`, and `Clustering`
+([`GCN.py:122-1933`](../legacy/root_gnn_dgl/models/GCN.py)); `UprootDataset`,
+`tHbbEdgeDataset`, `AugmentedDataset`, and photon-ID paths
+([`uproot_dataset.py:10-31`](../legacy/root_gnn_dgl/root_gnn_base/uproot_dataset.py),
+[`dataset.py:484-827`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
+[`photon_ID_dataset.py:1-33`](../legacy/root_gnn_dgl/root_gnn_base/photon_ID_dataset.py));
+optional loss and similarity utilities; and the no-op
+`root_gnn_base.utils.graph_augmentation` ([`utils.py:393-395`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)).
+The main path evaluates `test_loaders`; validation loaders are only assembled
+when a config has a validation fold ([`training_script.py:682-747`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+
+## 3. De facto interfaces to preserve
+
+### Metadata-aware dataset boundary
+
+The rewrite uses named `EventMetadata` (`fold`, `weight`, and stable
+`sample_id`) instead of exposing the legacy positional tracking tensor.
+`GraphSample`, `GraphBatch`, `SplitDefinition`, and `GraphDataLoader` form the
+ROOT-GNN orchestration boundary. Graph caches carry feature, graph, and cache
+schema versions and reject incompatible artifacts before loading.
+
+The current cache implementation stores processed `GraphSample` values (the
+Level-2 cache). The separation from `RootEventDataset` is intentional: a
+future Level-1 cache can store normalized `EventSample`/feature data for
+sequence or transformer representations without requiring DGL graph caches.
+
+The rewrite's shared data boundary is `gnn4colliders.data`: it reads selected
+ROOT/Awkward branches and returns architecture-neutral event samples. Feature
+construction and graph building remain separate downstream boundaries, so the
+same samples can be reused by non-graph model families.
+
+### Configuration and CLI boundary
+
+Hydra composes semantic YAML groups under `configs/` and passes the resolved
+configuration to explicit application factories in `gnn4colliders.config`.
+Those factories allow-list supported models, tasks, and trainer components;
+YAML is never treated as an arbitrary Python import specification. The thin
+`gnn4colliders` CLI selects `prepare`, `train`, `evaluate`, or `predict` and
+delegates to the stable data, training, checkpoint, and inference APIs. A new
+experiment should generally be a YAML change; new behavior belongs in Python.
+
+1. YAML `module`, `class`, `args`, plus runtime `sample_graph` and
+   `sample_global` injection.
+2. Dataset items `(DGLGraph, label, tracking, global_features)`.
+3. `ndata['features']`, `edata['features']`, seven node columns, and three edge
+   columns in `[deta, dphi, dR]` order.
+4. Tracking column 0 fold and column 1 weight semantics.
+5. `model(graph, global_feats)`, logits shape `[batch, out_size]`, and
+   `representation` where used.
+6. Weighted per-label loss, metric thresholds, checkpoint keys/prefix cleanup,
+   epoch filenames, and `.npz`/ROOT output fields.
+
+## 4. Ambiguous behavior
+
+- Historical edge order/self-loop expectations; empty and padding graph inputs.
+- Whether negative weights are meaningful or should always be absolute.
+- Whether “validation” is intended to differ from the active test-loader path.
+- Shape semantics of multi-label finishers and experimental transfer classes.
+- Whether chunk IDs must match historical `np.array_split` boundaries.
+- Required behavior for missing branches and dynamic selection expressions.
+
+## 5. Recommended rewrite boundaries
+
+Separate typed configuration; ROOT/Awkward I/O; selections/folds/features/
+edges; DGL cache/lazy loading/batching; active models and checkpoint adapters;
+objectives/metrics; training lifecycle; and inference/ONNX applications.
+Establish parity for the active `LazyDataset -> PreBatchedDataset ->
+Edge_Network` binary/multiclass path first. Add experimental classes only when
+a config or consumer proves they are required.

docs/compatibility.md

@@ -0,0 +1,22 @@
+# Compatibility boundary
+
+The package's canonical APIs use named `EventMetadata` fields and the version
+1 checkpoint schema. Compatibility is explicit and one-way: old artifacts are
+adapted into the new representation and are never rewritten implicitly.
+
+| Artifact or behavior | Supported | Boundary | Notes |
+| --- | :---: | --- | --- |
+| New checkpoint (`schema_version: 1`) | yes | `CheckpointManager` | Full model/lifecycle resume when state is present |
+| Legacy `model_epoch_N.pt` checkpoint | yes | `gnn4colliders.compat.load_legacy_checkpoint` | Model state and active early-stop fields are adapted |
+| Legacy DDP/compiled prefixes | yes | `normalize_legacy_state_dict_keys` | Supports `module.` and `_orig_mod.` |
+| Legacy ROOT-GNN classifier name | yes | `map_legacy_edge_network_state_dict` | Maps `classify` to `classifier` |
+| Legacy optimizer state | partial | legacy checkpoint adapter | Loaded when present; scheduler state is not available in the historical format |
+| Legacy scheduler state | no | — | Historical checkpoints do not carry a supported scheduler state |
+| Positional tracking rows | yes, at ingestion boundary | `EventMetadata.from_legacy_tracking` | Exactly `tracking[0] = fold`, `tracking[1] = weight`; shorter rows fail |
+| Generic/unknown tracking layouts | no | — | The package does not guess historical column meanings |
+| Modern NPZ output | yes | `inference.write_npz` | Named fields: `sample_id`, `logits`, `scores`, `predictions`, and available metadata |
+| Historical `tracking_info` NPZ output | no | — | No active consumer remains; positional output is intentionally unsupported |
+| Legacy YAML `module`/`class`/`args` | compatibility only | configuration boundary | Accepted only where the semantic factory can safely interpret it; new configs use semantic model names |
+
+The frozen `legacy/` tree remains available to parity tests and historical
+investigation. Production modules do not import executable code from it.

docs/configuration.md

@@ -0,0 +1,96 @@
+# Configuration guide
+
+The CLI composes the installed `gnn4colliders.configs/config.yaml` with semantic
+Hydra groups. A config
+describes experiment intent; it does not contain arbitrary Python module or
+class import paths.
+
+## Groups
+
+| Group | Purpose |
+| --- | --- |
+| `data` | ROOT source or graph-cache path, batch settings, and fold splits |
+| `model` | Model family and supported ROOT-GNN constructor settings |
+| `task` | Binary or multiclass loss, score, and metric semantics |
+| `trainer` | Device, seed, epochs, optimizer, scheduler, and early stopping |
+| `checkpoint` | Output directory, resume checkpoint, or pretrained checkpoint |
+| `inference` | Checkpoint, split, output path, and output format |
+| `environment` | Device and experiment output root |
+| `distributed` | Single-process/DDP selection and process-group backend |
+
+The active model groups are `root_gnn/edge_network` and
+`root_gnn/fine_tuned_edge_network`. The active task groups are
+`pretraining_multiclass`, `binary_classification`, and `tth_cp_finetune`.
+
+## Preparation
+
+`prepare` requires `data.files`, `data.cache.path`, `data.feature_branches`,
+`data.object_types`, and `data.scales`. `feature_branches` follows the shared
+seven-column feature contract: one branch/constant specification per output
+column and one entry per configured object type. `CALC_E` and `NODE_TYPE` are
+reserved derived specifications. `object_types` entries are `vector` or
+`single`. Preparation reads the configured tree in file order and writes a
+versioned `GraphSampleCache`.
+
+Example overrides are easiest to maintain in a YAML file for real datasets:
+
+```yaml
+# project-local example: data/my_events.yaml
+files: [data/events.root]
+tree_name: Events
+cache:
+  path: cache/events.pt
+feature_branches:
+  - [jet_pt]
+  - [jet_eta]
+  - [jet_phi]
+  - CALC_E
+  - [1.0]
+  - [0.0]
+  - NODE_TYPE
+object_types: [vector]
+scales: [1, 1, 1, 1, 1, 1, 1]
+fold_var: eventNumber
+weight_var: weight
+```
+
+Then compose it with `data=my_events`. The cache stores processed graph
+samples, labels, globals, named metadata, and feature/graph/cache schema
+versions. Changing the feature or graph schema requires a new compatible cache;
+loading a mismatched schema raises an error.
+
+## Common overrides
+
+```bash
+uv run gnn4colliders train \
+  data.cache.path=cache/events.pt \
+  data.batch_size=64 \
+  trainer.max_epochs=50 \
+  trainer.seed=123 \
+  environment.output_root=outputs/my_run
+```
+
+`data.batch_size` is per process. `data.splits.train_folds`,
+`validation_folds`, and `test_folds` define conventional train/validation/test
+selection and must be disjoint. Model/task mismatches are rejected during
+config validation; binary tasks require `model.out_size=1`, while multiclass
+tasks require `model.out_size=task.num_classes`.
+
+## Checkpoints and resolved configuration
+
+Training writes `epoch_####.pt` and the fully resolved configuration at
+`<environment.output_root>/resolved_config.yaml`. A checkpoint includes schema
+version, model weights/config, task config, trainer/optimizer/scheduler state,
+early stopping state, metadata, and optional RNG state. Set
+`checkpoint.resume=/path/to/epoch_####.pt` to continue a run. Set
+`checkpoint.pretrained=/path/to/epoch_####.pt` with the fine-tuned model group
+to load weights into a new task head; these options are mutually exclusive.
+
+## Environment profiles
+
+`environment=local` selects CPU by default. `environment=macos` selects the
+Apple MPS device for the portable PyTorch ROOT-GNN backend.
+`environment=perlmutter` selects
+the CUDA device and a conventional output-root pattern. Profiles should hold
+device/output policy only; site-specific module loads and filesystem paths
+belong in a launcher or shell environment.

docs/design.md

@@ -1,540 +0,0 @@
-# Design
-
-## Modules
-
-| Module | Contents |
-| --- | --- |
-| `config.py` | the experiment dataclasses, YAML loading and overrides |
-| `data.py` | ROOT reading, Hugging Face download, derived branches, object features, `PackedEvents` and its cache |
-| `loader.py` | `EventLoader`: deterministic shuffling, DDP sharding, collation workers |
-| `graphs.py` | fully connected `TensorGraph` batches, edge features, CUDA-graph padding buffers |
-| `sequences.py` | padded token batches for the transformer and Deep Sets, event augmentations |
-| `models.py` | `EdgeNetwork`, `FineTunedEdgeNetwork`, `ObjectTransformer`, `TransformerClassifier`, `DeepSets`, checkpoint loading |
-| `jepa.py` | JEPA pretraining of the transformer or edge network encoder (encoding or kinematic targets), event embeddings |
-| `contrastive.py` | contrastive (SimCLR) pretraining of the same encoders |
-| `probe.py` | k-NN and linear probes of frozen encoders, linear CKA between them |
-| `pfn.py` | our prior-fitted network over events and the task prior of meta-training |
-| `tabular.py` | event rows and encoder embeddings for TabPFN, TabICL and TabDPT; LoRA meta-training of TabPFN and TabICL |
-| `tasks.py` | multiclass and binary losses and metrics, the pretraining and PFN losses |
-| `train.py` | `Trainer`, early stopping, checkpoints |
-| `cuda_graphs.py` | captured training steps |
-| `distributed.py` | `torchrun` process group and collectives |
-| `inference.py` | model outputs over a loader, prediction files |
-| `export.py` | ONNX export |
-| `cli.py` | the six commands |
-
-## Data flow
-
-`prepare` reads each ROOT file in chunks of 100,000 entries (in parallel with
-`data.prepare_workers`) and stores every event's objects as rows of one
-`[objects, 7]` tensor with offsets, plus per-event labels, folds, weights and
-source file and entry. Nothing in the cache depends on the model: graphs and
-token sequences are built per batch, in the loader workers.
-
-A batch of the edge network is a set of fully connected graphs in flat
-edge-index form (`edge_src`, `edge_dst`, `node_batch`, `edge_batch`). Edges
-are directed, source-major and without self-loops except for one-node graphs,
-with features `[deta, dphi, dR]`. These conventions, the seven node columns,
-the MLP layout and the parameter names are those of the original ROOT-GNN, so
-the pretrained checkpoints in `Pretrained_GNN/` load without conversion.
-
-The transformer and Deep Sets take each event's objects as tokens
-`[events, objects, 7]`, zero-padded to the largest event of the batch, with
-a mask of real objects. A token is a linear map of the kinematic columns
-(`phi` as its cosine and sine) plus a learned embedding of `node_type`. The
-objects are stored by type (jets, electrons, muons, photons, MET), so a
-position embedding would mostly relearn the type; without one, neither
-model depends on object order, as with the edge network.
-
-* The transformer is a stack of pre-norm attention blocks, a mean over real
-  objects and a linear classifier. With `model.pair_bias` (as in the
-  Particle Transformer, Qu et al. 2022), an MLP of the pair features
-  `[deta, dphi, dR, log(1 + m^2)]` (the edge features plus the invariant
-  mass of the pair, objects taken as massless) adds a bias per head to the
-  attention logits, so the transformer sees the same pairwise angles as the
-  edge network.
-* Deep Sets (Zaheer et al. 2017; the particle flow network of Komiske et
-  al. 2019) applies an MLP to each object, sums over real objects and
-  applies an MLP to the sum: the baseline without object interactions.
-
-### JEPA pretraining
-
-`model.name: jepa` pretrains an object encoder without labels, as in
-I-JEPA (Assran et al. 2023), on the 12 Delphes samples of the supervised
-pretraining. `model.backbone` selects the encoder: the transformer's
-(`configs/delphes_jepa.yaml`) or the edge network, of the size of the
-supervised pretrained models (`configs/delphes_jepa_edge.yaml`), whose
-object encodings are its node states after the last processing step and
-whose context is the fully connected graph of the visible objects only.
-
-* a random `model.mask_fraction` of each event's objects (rounded, at least
-  one, and at least one left visible) is hidden;
-* the context encoder sees the visible objects; the target encoder, an
-  exponential moving average of it (momentum rising linearly from
-  `model.ema_momentum` to 1 over training) without gradients, sees all
-  objects;
-* a predictor (`model.predictor_layers` attention blocks over the visible
-  objects' encodings and one query per hidden object) predicts the hidden
-  objects' layer-normalized target encodings, with a mean squared error.
-  A query is told only the hidden object's type and its index among the
-  objects of that type (their pt order), not its kinematics; pair features
-  are used only between visible objects. For the edge network, its decoded
-  global state (what its classifier reads) is one more predictor input, so
-  the whole network is pretrained.
-
-The history reports `target_spread`, the standard deviation of the target
-encodings over a batch's objects averaged over dimensions, which falls
-towards zero if the representation collapses. The validation loss compares
-predictions with a moving target, so it does not select epochs: fine-tuning
-starts from `last.pt` (`configs/gamgam_jepa_finetune.yaml` for the
-transformer, `configs/gamgam_finetune.yaml` with `model.pretrained` for the
-edge network; encoder sizes must match). With `model.freeze_backbone=true` only the classifier is
-trained (a linear probe); `model.pretrained=null` trains the same
-architecture from scratch.
-
-`model.variance_weight` adds the VICReg variance term (Bardes et al. 2022):
-the mean shortfall below 1 of each dimension's standard deviation over the
-batch's context encodings (and the edge network's global states).
-
-First checks (one A100, `train.fast=true`, 2 pretraining epochs on Delphes
-fold 2, then 2 GamGam epochs, validation AUC, seed 1; not converged):
-
-| GamGam model | Validation AUC |
-| --- | --- |
-| transformer from scratch / from JEPA | 0.898 / 0.904 |
-| edge network from scratch | 0.895 |
-| edge network from the supervised pretrained model | 0.907 |
-| edge network from JEPA, last node states only | 0.845 |
-| edge network from JEPA, mean node states | 0.869 |
-| the same with `variance_weight` 0.1 / 1 | 0.874 / 0.810 |
-
-The edge network's JEPA encodings collapse without the variance term
-(`target_spread` 0.09 against about 1 with it), and none of the edge
-network variants yet beats training from scratch.
-
-### Contrastive and masked-object pretraining, and probes
-
-Two more label-free objectives pretrain the same encoders on the same
-samples, and their encoders load the same way (`context_encoder`):
-
-* `contrastive.py` (`configs/delphes_contrastive.yaml`): SimCLR, as JetCLR
-  and RS3L for jets. Two views of each event (random phi rotation, eta
-  flip, momentum scale within `model.scale_jitter`, each object but the
-  first dropped with `model.object_dropout`) are embedded and projected;
-  the NT-Xent loss at `model.temperature` contrasts them with the other
-  events of the batch. `embedding_spread` watches for collapse.
-* `model.jepa_target: features` (`configs/delphes_masked.yaml`): masked
-  object modelling (Golling et al. 2024). JEPA's masking and predictor
-  predict the hidden objects' kinematics instead of their target
-  encodings, with no target encoder: the comparison with JEPA isolates
-  predicting in embedding space.
-
-`gnn4colliders probe` compares encoders without fine-tuning (as jBOT):
-k-NN (cosine, `model.probe_neighbours`) and linear (L-BFGS logistic
-regression, `model.probe_l2`) classifiers of frozen event embeddings,
-trained on `data.train_events` events. A randomly initialized encoder
-(`model.pretrained=null`) is the reference; supervised checkpoints probe
-too. Results on GamGam: "Label efficiency on GamGam" below.
-
-`scripts/compare_encoders.py` compares encoders of one architecture by
-linear CKA (Kornblith et al. 2019), as the paper does: of their event
-embeddings (a matrix), and layer by layer against the first one (the mean
-over objects after each transformer block, or the edge network's mean node
-states after each step, then the event embedding). With a random encoder
-first, it shows where each objective moves the representation.
-
-### In-context models: TabPFN, TabICL, TabDPT and a collider PFN
-
-Both classify a dataset's events given `data.train_events` labelled
-training events (an equal number per class) as context, without training on
-that dataset; `cli.in_context_forward` builds them for `evaluate` and
-`predict`.
-
-* `tabular.py`: TabPFN (Prior Labs; `model.tabpfn_version`, default 3.5)
-  on fixed-width rows: the leading `model.slots` objects of each type (6
-  jets, 2 electrons, 2 muons, 3 photons, MET) with their six kinematic
-  columns, NaN when absent, and the number of objects of each type.
-* `pfn.py`: our own prior-fitted network. The object transformer (which
-  can start from the JEPA encoder) embeds each event; the context events get
-  a label embedding; `model.pfn_layers` attention blocks run over a task's
-  events, every event attending to the context events only, so queries are
-  independent and the context's keys and values are computed once. It is
-  meta-trained (`task.name: pfn`) on tasks drawn from the Delphes samples:
-  2-10 classes of `task.prior_classes` in random order, Dirichlet class
-  proportions, 16-2048 context events (log-uniform), 256 queries, and per
-  event a random phi rotation and eta flip, per task a momentum and energy
-  scale within `task.energy_jitter`. The GamGam processes are among the
-  Delphes samples; `task.prior_classes=[1,6,7,8,9,10,11]` leaves them out.
-  `configs/opendata_pfn.yaml` (written by `scripts/opendata_prior.py`)
-  meta-trains it on the ATLAS Open Data 2025 MC instead: 925 classes of 331
-  processes, H -> gamma gamma left out, a tenth of the processes held out
-  for validation, classes mixing 1 to 3 processes.
-* Options that bring the PFN closer to TabPFN, all off by default (the
-  default model's outputs match the earlier code to 2e-7):
-  `model.pfn_summary_tokens` learned tokens carry each event through the
-  encoder, with a context event's label added at the input instead of after
-  the encoder; `model.pfn_alternating` adds attention across the task's
-  events (summary tokens, context only) after every encoder block, as
-  TabPFN alternates attention over features and over rows, so an event's
-  encoding depends on the task; `model.pfn_normalize` scales each task's
-  momenta so that the context's geometric-mean object pT is 50 GeV (as
-  TabPFN standardizes features by the context); `model.pfn_ensemble`
-  averages the class probabilities of passes with the events rotated in
-  phi, flipped in eta and the class labels permuted (as TabPFN ensembles
-  over class orders and preprocessings); and `task.synthetic_tasks` makes
-  a fraction of the training tasks synthetic, as TabPFN's prior draws
-  random functions: events of 1 to 3 processes ranked by a random network
-  (16 tanh units) of 1 to 4 of 21 event observables (per object type the
-  count, scalar sum pT, leading pT and mass; the b-tag count), cut into
-  classes of Dirichlet fractions. A large context is encoded in chunks, one
-  block at a time, but the alternating PFN keeps every context event's
-  tokens and the keys and values of every block on the GPU: 100k events
-  fit, all 1.98M GamGam training events do not (out of memory on a 40 GB
-  A100).
-* Options from TabPFN-3.5 (219M parameters, 79 times ours; read from
-  its checkpoint and code), all off by default: `model.pfn_retrieval` makes
-  the logits the log of an attention-weighted vote of the context's labels
-  (TabPFN's many-class decoder) instead of a linear classifier;
-  `model.pfn_qk_norm` RMS-norms queries and keys per head and scales the
-  queries by a learned factor times the log of the number of keys
-  (TabPFN-3's scalable softmax), in all attention across events;
-  `model.pfn_query_kv` lets the queries read the context through one
-  key/value head of their own, so the memory of a large context is
-  `n_heads` times smaller; `model.pfn_concat` joins the summary tokens into
-  one token `pfn_summary_tokens * hid_size` wide for the blocks after the
-  encoder (TabPFN-3.5 joins 8 x 128 to 1024), with the label embedded
-  again; `model.pfn_fourier` learned frequencies of each object's log pT,
-  eta and phi and `model.pfn_ranks` (each object's pT and energy ranked
-  among the context's objects of its type, as TabPFN adds ECDF ranks) are
-  added to the object embeddings, through maps that start at zero.
-  `train.lr_schedule: cosine` warms up over `train.warmup_epochs`.
-* Meta-training speed: training batches are encoded in 4 parts of events
-  with similar numbers of objects (only 37% of the padded object tokens
-  were real), tasks are sampled in the loader workers rather than all at
-  the start of an epoch, and under `torchrun` all ranks draw the same
-  context size for a step, so that they wait less for each other. Steps are
-  drawn from a generator per seed, epoch, rank and step, so the tasks
-  differ from those of runs before this change. In-context inference sorts
-  the context by number of objects and attends in chunks of
-  `data.batch_size` events.
-
-* `tabular.py` also runs TabICLv2 (Qu et al., ICML 2026; BSD-3, open
-  pretraining code) on the same rows, as a second tabular in-context model.
-  Rows can measure azimuths from the leading object's (`model.relative_phi`)
-  and end with the invariant mass and Delta R of every pair of the
-  `model.pair_objects` leading objects by pt (energies are `pt cosh(eta)`,
-  so objects are massless; a pair with MET gives a transverse-like mass).
-  Both are off by default, so the default rows are unchanged.
-* `tabular.py` also runs TabDPT (Ma et al. 2024; Apache-2.0), trained on
-  real tables rather than a synthetic prior; with
-  `model.tabdpt_context_size` each event's context is its nearest context
-  events (faiss retrieval), beyond what fits in one pass. It computes in
-  float64 in places (on a Mac it runs on the CPU), and on macOS faiss and
-  torch crash together unless `OMP_NUM_THREADS=1`.
-* Tabular models can see, instead of or beside the rows, the event
-  embeddings of a pretrained object encoder (`model.tabular_inputs:
-  embeddings` or `both`, the encoder of `model.backbone` and
-  `model.pretrained`): physics-aware embeddings under a generic tabular
-  prior. In meta-training the encoder trains with the LoRA updates unless
-  `model.freeze_backbone`.
-* `tabular.py` also meta-trains the TabPFN-3.5 network itself on the
-  PFN's tasks (`model.name: tabpfn`, `task.name: pfn`), on the rows above
-  (built after the augmentation), starting from Prior Labs' weights. Calling
-  the network directly skips TabPFN's preprocessing and ensembling, which
-  on the GamGam fixtures cost nothing (validation AUC 0.852 against 0.843
-  with `predict_proba`, 61 events). By default the network is frozen and
-  the four attention projections and two MLP layers of each of its 24
-  in-context blocks get rank-8 LoRA updates (2.8M parameters, zero at the
-  start, so training starts from TabPFN as released); `model.lora_rank=0`
-  trains all weights. In context, the context's keys and values are cached
-  (the same logits as one pass over context and queries, to 2e-5). It
-  separates what Prior Labs' prior and what collider tasks contribute:
-  TabPFN as released, TabPFN meta-trained here, and our PFN. TabICL
-  (`configs/delphes_tabicl.yaml`; BSD-3, so meta-trained weights can be
-  shared) meta-trains the same way. Its attention reads fused projection
-  weights directly, so LoRA is a parametrization of the weights (12 blocks,
-  0.59M of 28M parameters at rank 8), for TabPFN too. TabICL's network takes
-  no NaN: rows are standardized by the context's statistics, absent values
-  at the mean, clipped at 8 sigma. On the fixtures this gives AUC 0.827
-  against 0.820 with its classifier, and the cached in-context forward
-  equals the training forward only in float32 (with mixed precision on the
-  CPU, AUC 0.59).
-
-### Label efficiency on GamGam (first study, 2026-10-22)
-
-GamGam validation AUC (%) against the number of labelled training events N
-(`data.train_events`, equal per class), one seed, on one 4-GPU A100 node in
-an interactive allocation (3.5 h). Each entry is evaluated on 20,000 random
-validation events (`data.eval_events=20000`, about +-0.3); on all 246,922,
-evaluated for the supervised runs and the leading TabPFN runs, each is 0.05
-to 0.41 lower, in the same order (TabPFN-3.5: 80.28, 86.92, 89.01, 0.08 to
-0.13 lower, matching the M2 Max's 80.26 and 86.93).
-
-| model | N = 100 | 1000 | 10000 | 50000 |
-|---|---|---|---|---|
-| edge network, from scratch | 78.00 | 82.74 | 87.66 | |
-| edge network, from `Pretrained_GNN` | 84.95 | 87.48 | 89.98 | |
-| transformer, from scratch | 72.33 | 83.47 | 87.03 | |
-| transformer, from JEPA | 73.25 | 85.10 | 88.60 | |
-| transformer, from masked objects | 75.92 | 86.12 | 88.73 | |
-| TabPFN-3.5 (rows) | 80.41 | 87.05 | 89.09 | 89.56 |
-| TabPFN-3.5 (rows, `relative_phi`, `pair_objects: 4`) | 80.61 | 87.78 | 89.56 | |
-| TabPFN-3.5 (rows + JEPA embeddings) | 80.97 | 87.43 | 89.44 | |
-| TabPFN-3.5 (rows + masked-object embeddings) | 82.99 | 88.22 | 89.94 | |
-| TabPFN-3.5 (rows + `Pretrained_GNN` embeddings) | 87.32 | 88.96 | 90.07 | |
-| TabPFN-3.5, released network, direct forward | 80.74 | 86.64 | 89.03 | |
-| TabPFN-3.5, meta-trained on Delphes (LoRA) | 80.52 | 84.77 | 88.34 | |
-| TabICLv2 | 80.53 | 85.68 | 88.41 | 89.12 |
-| TabDPT | 77.48 | 84.71 | 88.42 | |
-| our PFN, meta-trained on Delphes | 78.13 | 76.18 | 76.51 | |
-
-Probes (k-NN / linear) of the frozen encoders, the same N:
-
-| encoder | N = 100 | 1000 | 10000 |
-|---|---|---|---|
-| random | 73.4 / 72.6 | 76.6 / 83.1 | 79.5 / 85.4 |
-| JEPA | 73.0 / 73.4 | 78.9 / 83.4 | 81.8 / 87.9 |
-| masked objects | 75.0 / 77.5 | 81.7 / 85.7 | 84.4 / 88.5 |
-| contrastive | 61.6 / 67.9 | 70.0 / 80.3 | 76.4 / 85.9 |
-
-Setup (every run and its overrides: `scripts/label_efficiency_study.sh`):
-the pretraining runs saw the same 4M class-balanced Delphes events for 8
-epochs (`data.train_events=4000000`, 50 to 70 min each); the PFN
-meta-trained for 25 epochs of 1000 steps (25,000 steps, 95 min) and TabPFN
-for 6 epochs of 500 (3,000 steps, 48 min). Supervised runs train for a
-fixed number of steps with early stopping off (`train.early_stopping=null`;
-batch size, epochs: 100, 1000 for N = 100; 128, 400 for N = 1000; 256, 100
-for N = 10000; 20 to 35 min at N = 100, 2 to 4 min at N = 10000) and keep
-the epoch of best validation AUC, chosen on the same 20,000 events that are
-reported. That choice uses validation labels the in-context models do not
-get, and from N = 1000 most supervised runs overfit: their last epoch is far
-below the kept one. On all validation events (92% of them not used for the
-choice) the kept epochs lose 0.05 to 0.41, no more than the in-context runs
-lose, so picking the maximum over noisy evaluations inflates them little:
-
-| supervised run (N = 100 / 1000 / 10000) | kept epoch, 20k | kept epoch, all | last epoch, 20k |
-|---|---|---|---|
-| edge network, from scratch | 78.00 / 82.74 / 87.66 | 77.80 / 82.69 / 87.42 | 68.88 / 78.88 / 87.65 |
-| edge network, from `Pretrained_GNN` | 84.95 / 87.48 / 89.98 | 84.54 / 87.09 / 89.74 | 80.75 / 81.58 / 86.85 |
-| transformer, from scratch | 72.33 / 83.47 / 87.03 | 71.99 / 83.35 / 86.95 | 72.33 / 83.44 / 87.03 |
-| transformer, from JEPA | 73.25 / 85.10 / 88.60 | 73.07 / 84.96 / 88.45 | 72.34 / 78.82 / 81.46 |
-| transformer, from masked objects | 75.92 / 86.12 / 88.73 | 75.73 / 85.84 / 88.60 | 75.92 / 79.73 / 85.62 |
-
-With only N labelled events there is no large labelled set to choose the
-epoch on; without the choice (last epoch) every supervised model is below
-TabPFN-3.5 as released at every N except `Pretrained_GNN` at N = 100.
-
-With the configs' settings instead (batch size 1024, at most 100 epochs,
-early stopping after 5 epochs without training-loss improvement) N = 1000
-is one step per epoch, at most 100 steps: the transformer then reaches
-70.04 at N = 1000 instead of 83.47. The in-context models take 10 to 30 s
-for the 20,000 queries at N = 1000 and 10000 (TabPFN 109 s at N = 50000;
-the first run of each, N = 100, also loads the weights).
-
-Findings:
-
-* TabPFN-3.5 as released, without training, comes within 0.4 (N = 1000)
-  and 0.9 (N = 10000) of fine-tuning the supervised `Pretrained_GNN` (0.2
-  and 0.7 on all validation events), with its epoch chosen on validation
-  labels, and above it from N = 1000 without that choice. It beats every
-  model trained from scratch at every N. TabICL is close behind; TabDPT is
-  behind at small N. The physics row features add up to 0.7.
-* Encoder embeddings help TabPFN most at small N. With the label-free
-  masked-object encoder: +2.6 at N = 100. With `Pretrained_GNN` (trained
-  with labels, on Delphes samples that include the GamGam processes): the
-  best result at every N run (100 to 10000; not run at 50000), +6.9 at
-  N = 100.
-* Meta-training on Delphes tasks does not help. TabPFN with LoRA loses 1.9
-  at N = 1000 against the same network with zero updates (its direct
-  forward costs at most 0.4 against the classifier's preprocessing and
-  ensembling). Our PFN is below TabPFN and gets worse with more context
-  (76 at N = 1000, inside the meta-training context sizes 16 to 2048): to
-  investigate before the held-out-processes run, which was not done.
-* Of the label-free pretraining objectives, masked-object modelling is the
-  best, in probes and fine-tuned; JEPA's prediction in embedding space adds
-  nothing over it here. Both improve fine-tuning the transformer over
-  training it from scratch (+0.9 to +3.6) but stay below fine-tuning the
-  edge network from `Pretrained_GNN`. That gap mixes the architecture (from
-  scratch, the edge network is 5.7 above the transformer at N = 100) with
-  the objective (label-free or supervised); the like-for-like comparison,
-  JEPA of the edge network (`configs/delphes_jepa_edge.yaml`), was not run.
-* Contrastive pretraining is below a random encoder in the k-NN probe at
-  every N and in the linear probe at N = 100 and 1000 (67.9 against 72.6,
-  80.3 against 83.1); at N = 10000 its linear probe is 0.5 above (85.9
-  against 85.4). Without the momentum-scale augmentation
-  (`model.scale_jitter=0`, 4 epochs) it is no better (72.4 / 84.3 at
-  N = 10000).
-
-Not done:
-
-* the held-out-processes meta-training (`task.prior_classes` without the
-  five GamGam processes) of the PFN, gated on it beating TabPFN, and of
-  TabPFN and TabICL;
-* TabICL meta-training and `model.lora_rank=0`;
-* fine-tuning from contrastive pretraining (only probed), and JEPA of the
-  edge network;
-* more seeds, and the test split.
-
-The loss is a weighted mean per class present in the batch, averaged over
-those classes, which balances the classes without explicit class weights.
-Metrics (accuracy, one-vs-rest AUC averaged over classes, macro precision,
-recall and F1) use events with positive weight.
-
-### The TabPFN-like PFN on GamGam (2026-10-27)
-
-GamGam test AUC (%) in context, against the number of labelled context
-events N (equal per class), on 50,000 test events, one seed. Every PFN is
-meta-trained on the Open Data prior (`configs/opendata_pfn.yaml`). The
-earlier PFN (pooled events, contexts of 16 to 4096, random encoder, fp32)
-is the checkpoint of epoch 38. The new ones train 48 epochs of 1000 steps
-(2.3 h on one A100 each, bf16) with 2 summary tokens, alternating
-attention, the per-task momentum scale, the masked-object encoder
-(`delphes_masked`), half synthetic tasks, contexts of 16 to 16384 and
-batches of 16384 events; each ablation changes one of them.
-
-| PFN | N = 100 | 1000 | 10000 | 100000 |
-|---|---|---|---|---|
-| earlier | 80.63 | 82.47 | 83.54 | 83.53 |
-| earlier, ensemble of 8 | 81.92 | 83.87 | 84.67 | 84.77 |
-| new | 79.12 | 80.39 | 83.79 | 84.18 |
-| new, ensemble of 8 | 79.67 | 82.26 | 84.46 | 84.77 |
-| new, no alternating attention | 77.14 | 81.21 | 84.43 | 84.68 |
-| new, random encoder | 76.47 | 81.76 | 83.42 | 83.56 |
-| new, no synthetic tasks | 81.28 | 83.14 | 84.72 | 85.45 |
-| new, no synthetic tasks, ensemble of 8 | **83.75** | **86.11** | **86.54** | **86.52** |
-
-Validation on held-out Open Data processes after 48 epochs (loss /
-accuracy): new 0.772 / 0.722, no alternating 0.800 / 0.712, random encoder
-0.809 / 0.711, no synthetic tasks 0.665 / 0.764; all still falling. One
-pass takes 1 to 3 s for the 50,000 queries at N = 1000 and 10000 and 35 to
-58 s at N = 100000; the ensemble of 8 takes 460 s at N = 100000.
-
-Findings:
-
-* Synthetic tasks hurt at every N (0.9 to 2.8 against the same PFN
-  without them; 1.8 to 4.1 ensembled) and in validation. Without them the
-  PFN beats the earlier one at every N, by 1.8 to 2.2 with both
-  ensembled.
-* The ensemble of 8 adds 1.1 to 3.0 to the PFN without synthetic tasks.
-* The masked-object encoder adds 2.7 at N = 100 and 0.6 at N = 100000.
-* Alternating attention adds 2.0 at N = 100 and costs 0.5 to 0.8 from
-  N = 1000; it makes training 20% slower.
-* Against the validation AUCs of the first study (on other events, so
-  only roughly comparable): the best PFN is above TabPFN-3.5 on rows
-  (80.41) and every supervised transformer at N = 100, below the edge
-  network from `Pretrained_GNN` (84.95) and TabPFN-3.5 with its embeddings
-  (87.32). At N = 1000 it equals the transformer from masked objects
-  (86.12) and is 0.9 and 1.4 below TabPFN-3.5 and `Pretrained_GNN`; at
-  10000 it is 2.2 to 3.4 below all three, and it gains nothing from 10000
-  to 100000.
-
-Not done: all 1.98M training events as context (out of memory; see the
-PFN options above), longer meta-training, and the held-out-processes
-prior.
-
-### TabPFN-3.5 options and a larger PFN on GamGam (2026-10-28)
-
-Both PFNs start from the best one above (no synthetic tasks) and add the
-TabPFN-3.5 options `pfn_retrieval`, `pfn_qk_norm`, `pfn_query_kv`,
-`pfn_fourier=16` and `pfn_ranks`, the cosine schedule (2 warm-up epochs,
-learning rate 1e-4) and the faster meta-training. "Small" keeps its size
-(2.9M parameters) and trains 40 epochs on one A100 (168 s each). "Large"
-has 4 summary tokens joined to one 512-wide token and 12 blocks after the
-encoder (42M parameters) and trains 26 epochs on three A100s under
-`torchrun` (280 s each, three times the tasks per step). Same test events
-and protocol as above, one seed.
-
-| PFN | N = 100 | 1000 | 10000 | 100000 |
-|---|---|---|---|---|
-| best above, ensemble of 8 | 83.75 | 86.11 | 86.54 | 86.52 |
-| small | 81.73 | 84.65 | 86.76 | 87.66 |
-| small, ensemble of 8 | 83.15 | 86.20 | 87.63 | 88.20 |
-| large | 82.90 | 85.83 | 87.72 | 88.43 |
-| large, ensemble of 8 | **84.06** | **86.80** | **88.24** | **88.62** |
-
-Validation on held-out Open Data processes in the last epoch (loss /
-accuracy): small 0.599 / 0.789, large 0.561 / 0.799 (best above 0.665 /
-0.764). At N = 100000 one pass takes 21 s (small) and 87 s (large), the
-ensemble of 8 169 s and 698 s; the large ensemble peaks at 13 GiB.
-
-Findings:
-
-* The options, schedule and faster training add 0.5 to 2.2 to one pass of
-  the best PFN above (81.28, 83.14, 84.72, 85.45), and the AUC now rises
-  from 10000 to 100000 (by 0.4 to 0.9) instead of staying flat.
-* The larger PFN adds 0.8 to 1.2 more. One pass of it is 0.9 and 0.3
-  below the earlier ensemble at N = 100 and 1000 and 1.2 and 1.9 above it
-  from 10000; ensembled, it gains 0.3 to 2.1 on the earlier ensemble.
-* The ensemble adds less than before: 0.2 to 1.6.
-* The options were not ablated one by one, and the two runs differ in
-  epochs and GPUs as well as size.
-* Against the validation AUCs of the first study (only roughly
-  comparable): at N = 100 the large ensemble is 3.7 above TabPFN-3.5 on
-  rows and 0.9 below `Pretrained_GNN`; at 1000 it is 0.3 and 0.7 below
-  them, at 10000 0.9 and 1.7 below.
-
-Not done: all 1.98M training events as context, one ablation per option,
-more seeds.
-
-## Performance
-
-Graphs differ in size, which suits neither `torch.compile` nor CUDA graphs
-out of the box:
-
-* `torch.compile(dynamic=True)` treats node and edge counts as symbolic
-  sizes, so one compiled graph serves all batches without padding.
-* CUDA graph replay needs fixed shapes. Batches are copied into buffers
-  rounded up to multiples of 512 nodes and 4096 edges; the remainder forms 64
-  padding graphs connected only to themselves, so the real graphs' outputs do
-  not change. A step is captured once per bucket, a few per epoch. The
-  optimizer runs in Adam's `capturable` form with device-tensor learning
-  rates, which the scheduler updates in place.
-* bfloat16 autocast covers the forward pass; logits, loss, graph sums (which
-  divide by edge counts) and the optimizer stay float32.
-
-GamGam training step, 1024-event batches (about 5,600 nodes and 30,000
-edges), one A100-SXM4-80GB, `OMP_NUM_THREADS=8`, four loader workers (step
-times measured before the move to flat modules, epoch times after):
-
-| Mode | ms per batch | Epoch with validation |
-| --- | ---: | ---: |
-| float32 eager | 19.2 | 37 s |
-| CUDA graphs + TF32 | 8.5 | |
-| `torch.compile` | 13.0 | |
-| compile + CUDA graphs + TF32 | 6.1 | |
-| compile + CUDA graphs + bfloat16 (`train.fast`) | 5.9 | 9 s (first: 59 s, compiling) |
-
-Eager training is bound by kernel launches (about a thousand small kernels
-per step), not by arithmetic; bfloat16 alone is slower because of the extra
-cast kernels, which `torch.compile` fuses. Results of the fast modes agree
-with float32 to rounding, not bitwise (validation AUC 0.895 against 0.897
-after three epochs). Preparing the 2.5M GamGam events with 8 workers takes
-about 9 s on a login node (40 s was seen on a compute node, likely a cold
-file cache).
-
-PyTorch sizes its CPU thread pool to every visible core; on a 128-core node
-that makes batch construction several times slower and CPU reductions
-non-reproducible, so set `OMP_NUM_THREADS`. Single-threaded CPU training is
-bitwise reproducible.
-
-## ATLAS Open Data GamGam versus Ho et al. (2026)
-
-The main benchmark reproduces the from-scratch baseline of Table 3 (5
-classes, 2.5M events, 80:10:10 split, five seeds). Choices the paper does not
-specify:
-
-* electrons and muons are split from `lep_*` by `lep_type` (11, 13);
-* b-tagging is `jet_MV2c10 > 0.8244273`, the 70% working point;
-* momenta and energies are scaled by 1e-4 (MeV to the pretraining unit of
-  10 GeV);
-* the fold is `eventNumber % 10`: folds 0-7 train, 8 validates, 9 tests;
-* training is unweighted (23% of ttH events have `mcWeight = -1`).
-
-The paper's stopping rule (5 epochs without training-loss improvement) never
-fires within 200 epochs on this task; validation AUC peaks earlier (epoch 67
-for seed 1) and `best.pt` keeps that epoch.
-
-| Test split | Accuracy | AUC | F1 | Precision | Recall |
-| --- | ---: | ---: | ---: | ---: | ---: |
-| Paper (5 seeds) | 71.8 ± 0.2 | 91.2 | 63.0 | 62.2 | 66.3 |
-| This package, seed 1 (CUDA graphs + TF32) | 71.1 | 91.3 | 62.6 | 62.4 | 66.7 |

docs/end_to_end_validation.md

@@ -0,0 +1,30 @@
+# End-to-end legacy/rewrite validation
+
+Task 21 compares staged event identity, labels, folds, weights, globals, node
+features, topology, edge features, batching, fixed-weight forward, loss and
+metrics, one optimizer step, short training, checkpoint reload, and inference.
+The first failing stage is retained in JSON and Markdown reports.
+
+Legacy and rewrite extraction may use separate environments. Each writes the
+version-1 artifact described in `validation/README.md`; the comparator never
+imports legacy modules or DGL graph objects. Record ROOT file/tree, selection,
+event count, file size, SHA-256, branches, seeds, versions, device and source
+identities in the artifact manifest. Do not use an absolute personal path as a
+scientific event identity.
+
+The captured runtime versions and seed policy are recorded in
+`validation/manifests/environments.json`; the HF composite fixture provenance
+is recorded in `validation/manifests/multiclass_fixture.json`.
+
+Strict parity applies to preprocessing, topology, fixed forward, loss, and
+one-step CPU updates. Multi-epoch, GPU and DDP comparisons are scientific:
+compare curves, metrics, distributions and event-level correlations. Named
+metadata, modern NPZ names and the version-1 checkpoint schema are expected
+interface differences and must be reported explicitly.
+
+Run the existing gates separately:
+
+```bash
+GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 uv run pytest
+uv run ruff check .
+```

docs/export.md

@@ -0,0 +1,36 @@
+# ROOT-GNN ONNX export
+
+ONNX export is an inference-only adapter for prepared ROOT-GNN graph batches.
+It does not read ROOT files or move feature construction into the deployment
+model. Install the optional dependencies with:
+
+```bash
+uv sync --extra root-gnn --extra onnx
+```
+
+Export a checkpoint with the project CLI:
+
+```bash
+uv run gnn4colliders export \
+  export.checkpoint=/path/to/checkpoint.pt \
+  export.output=model.onnx \
+  data.cache.path=/path/to/graph-cache.pt
+```
+
+The exported model accepts six tensor inputs: `node_features`,
+`edge_features`, `edge_src`, `edge_dst`, `node_batch`, and `global_features`.
+`node_batch` identifies the graph for each node; edge membership is derived
+from `node_batch[edge_src]`. Graph, node, edge, and batch dimensions are
+dynamic. The model returns raw `logits`; task postprocessing and metadata stay
+outside ONNX.
+
+The exporter uses ONNX opset 17, validates the structure with `onnx.checker`,
+and writes compact provenance/schema metadata to `model.onnx.json`. It
+supports `EdgeNetwork` multiclass models and `FineTunedEdgeNetwork` binary
+models. Empty graphs are invalid; single-node graphs (zero edges) are
+represented and pooled with a zero edge contribution.
+
+Direct DGL export was not retained: the active model uses DGL graph mutation
+and reductions (`apply_edges`, `update_all`, and graph pooling) that are not a
+portable ONNX contract. The adapter expresses those operations with standard
+tensor indexing, `index_add`, and per-graph means.

docs/migration.md

@@ -0,0 +1,295 @@
+# Incremental migration plan: `root_gnn_dgl`
+
+The target is `legacy/root_gnn_dgl/`. `legacy/physicsnemo/` is a prior rewrite
+attempt and may inspire abstractions, but it is not a parity target. Neither
+legacy tree should be modified during migration.
+
+Each phase should add focused unit tests, a deterministic fixture in
+`data/fixtures/`, and parity tests under `tests/parity/` before moving upward.
+Record intentional differences and checkpoint consequences here.
+
+## Phase 0 — freeze observations and fixtures
+
+Capture a small representative ROOT-equivalent fixture containing the seven
+active node features, three edge features, labels, fold values, weights, and
+globals. Record outputs of `node_features_from_tree`, `full_connected_graph`,
+`EdgeDataset.make_graph`, and `fold_selection`
+([`dataset.py:15-59`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
+[`dataset.py:471-482`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
+[`utils.py:121-143`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)). Preserve
+one `.bin`, one `model_epoch_N.pt`, one evaluation `.npz`, and one inference
+`.npz` fixture if available.
+
+Task 3 characterization records the active-path observations. Node rows are
+concatenated by object type in the configured order (jets, electrons, muons,
+photons, MET), and the seven columns are `[pt, eta, phi, energy, btag, charge,
+node_type]`. `CALC_E` is `pt*cosh(eta)` before the configured column scale is
+applied. The graph is directed and uses all ordered pairs except self-loops
+for graphs with more than one node; edge order is source-major. A one-node
+graph is a special case: the no-self-loop branch retains its sole self-loop.
+Edge columns are `[deta, dphi, dR]`, with `dphi` wrapped into `[-pi, pi]`.
+Dataset items expose `(graph, label, tracking, global_features)`; tracking
+column 0 is the fold identifier and column 1 is the event weight. These are
+compatibility observations, not proposed fixes.
+
+## Phase 1 — configuration boundary
+
+Implement a typed configuration layer that reads `Training`, `Model`,
+optional `Loss`, and `Datasets`. Initially retain a compatibility adapter for
+`module`/`class`/`args` and runtime injection of `sample_graph` and
+`sample_global`, matching `buildFromConfig`
+([`utils.py:10-43`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)). Keep
+dynamic imports isolated at this boundary rather than spreading reflection
+through new code.
+
+## Phase 2 — pure preprocessing parity
+
+Port and test, in isolation:
+
+- branch-to-node conversion, `CALC_E`, `NODE_TYPE`, constants, scaling, empty
+  objects, and dtypes (`node_features_from_tree`,
+  [`dataset.py:15-50`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py));
+- string/tuple selections and cutflow (`check_selection`, `selection_mask`,
+  `compute_cutflow`, [`dataset.py:75-158`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py));
+- fold masks and cache suffixes (`fold_selection`, `fold_selection_name`,
+  [`utils.py:121-143`](../legacy/root_gnn_dgl/root_gnn_base/utils.py));
+- deterministic chunk partitioning (`hash_partition`,
+  [`batched_dataset.py:27-31`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
+
+This is the highest-value parity layer: model parity is invalid if graph inputs
+differ.
+
+Task 4 implements the shared branch-to-node feature builder under
+`gnn4colliders.features`. It preserves the active seven-column schema,
+object-type ordering, explicit scales, derived `CALC_E`, node-type codes,
+float32 output, and supported empty vector collections. Selection, fold, and
+chunk helpers remain deferred to later data-infrastructure work.
+
+Task 6 implements the shared ROOT/Awkward ingestion boundary under
+`gnn4colliders.data`. `RootEventDataset` returns immutable, architecture-neutral
+`EventSample` values with selected branch data, labels, tracking, and globals;
+events are ordered by input file order with a global zero-based index. Fold
+filtering, caching, batching, and model-specific conversion remain deferred.
+
+## Phase 3 — graph construction and cache format
+
+Implement graph construction with tests for node/edge counts, directed edge
+ordering, self-loop policy, `[deta, dphi, dR]` order, metadata, and empty graphs.
+Preserve the dataset item contract `(graph, label, tracking, global_features)`
+from `RootDataset.__getitem__`
+([`dataset.py:465-469`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)).
+
+Then implement DGL `.bin` serialization, lazy chunk loading, pre-batching, and
+padding. Compare against `RootDataset.save/load`, `LazyDataset`, and
+`PreBatchedDataset` ([`dataset.py:396-469`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
+[`batched_dataset.py:129-174`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
+Treat `NONE`, `STEPS`, `FIXED`, and `NODE` as explicit features; do not hide
+the hardcoded fixed padding sizes.
+
+Task 7 establishes the metadata-aware orchestration boundary around this
+phase: `EventMetadata`, `GraphSample`, `GraphBatch`, fold-based split
+selection, deterministic batching, and a version-checked graph-sample cache.
+The cache is deliberately Level 2; normalized event caching remains a future
+extension so non-graph model families can reuse ROOT preprocessing.
+
+## Phase 4 — active model parity
+
+Task 8 adds the active `EdgeNetwork` and `FineTunedEdgeNetwork` under
+`gnn4colliders.models.root_gnn`. The update order and MLP ordering follow the
+legacy active path. The rewrite uses an explicit backbone/classifier boundary,
+local DGL graph scope, and does not mutate global RNG state in constructors.
+Model parity now covers fixed-weight pretraining and transfer paths, including
+historical checkpoint prefixes. The legacy transfer implementation has an
+active bug when nonempty globals are supplied (`Pretrained_Output` ignores its
+argument); parity therefore characterizes its supported no-global path, while
+the rewritten model supports both global and fallback modes.
+
+Task 22 adds `TensorEdgeNetwork`, a native PyTorch realization of the same
+message-passing equations. It receives explicit node/edge index tensors from
+the batching boundary, retains parameter names for legacy checkpoint loading,
+and has fixed-weight CPU parity coverage against the DGL backend. It is the
+default training backend for CPU, CUDA, and Apple MPS; DGL remains available as
+a cache and compatibility adapter during the transition.
+
+Port `models.GCN.Edge_Network` first. Preserve constructor parameters,
+`forward(graph, global_feats)`, feature keys, processor order, MLP LayerNorm
+placement, and logits shape. Compare intermediate and final tensors on fixed
+graphs using the legacy architecture
+([`GCN.py:18-35`](../legacy/root_gnn_dgl/models/GCN.py),
+[`GCN.py:182-251`](../legacy/root_gnn_dgl/models/GCN.py)).
+
+Next port `Transferred_Learning_Finetuning`, including pretrained
+`model_state_dict` loading, removal of the final classifier, and new classifier
+initialization ([`GCN.py:884-997`](../legacy/root_gnn_dgl/models/GCN.py)). Test
+both frozen and unfrozen modes. Defer other model classes until an active
+config or consumer proves they are needed.
+
+## Phase 5 — objectives and metrics
+
+Implement the default objective exactly: elementwise configured loss,
+tracking-column weights, per-unique-label normalization, and averaging across
+labels ([`training_script.py:320-359`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+Add parity cases for positive, zero, and negative weights and binary versus
+multiclass shapes.
+
+Port metric behavior from
+[`training_script.py:438-510`](../legacy/root_gnn_dgl/scripts/training_script.py):
+sigmoid threshold 0.5, argmax, weight masking, weighted ROC AUC, one-vs-rest
+multiclass AUC, and NaN behavior when AUC is undefined. Add `models/loss.py`
+classes only with dedicated tests; do not substitute their reductions.
+
+## Phase 6 — checkpoint and lifecycle
+
+Task 10 implemented the in-memory single-process training lifecycle before the
+checkpoint portion of this phase: `Trainer`, explicit optimizer/scheduler
+builders, `EarlyStopping`, reproducibility seeding, `GraphBatch.to`, and
+epoch/history result types. Checkpoint persistence/resume and the Python
+inference/evaluation and named NPZ/ROOT output layers are now implemented.
+Distributed execution and CLI wiring were completed in the later phases.
+
+Task 10 also establishes corrected split semantics: validation is evaluated
+every epoch and is the only split used for model selection or early stopping;
+the test split remains held out and is evaluated separately after fitting. The
+legacy loader naming inversion (`test` used for selection and `val` held out)
+is not carried into the rewrite.
+
+Create a checkpoint adapter preserving `model_epoch_<epoch>.pt` and keys
+`epoch`, `model_state_dict`, `optimizer_state_dict`, and `early_stop`
+([`training_script.py:565-604`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+Support legacy DDP/compiled prefixes (`module.` and `_orig_mod.`) as exercised
+by checkpoint lookup and inference
+([`utils.py:145-248`](../legacy/root_gnn_dgl/root_gnn_base/utils.py),
+[`inference.py:274-290`](../legacy/root_gnn_dgl/scripts/inference.py)). Port
+`EarlyStop` state and log parsing separately
+([`utils.py:325-390`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)). Verify
+resume, restart, early termination, and `.npz` fields before distributed work.
+
+## Phase 7 — CLI, inference, and export
+
+Task 12 implemented ordered prediction/evaluation, task-owned score
+semantics, checkpoint weight-only loading, named metadata retention, NPZ
+output, and explicit ROOT entry alignment. The semantic CLI and the validated
+ROOT-GNN ONNX export adapter are implemented.
+
+Task 13 adds Hydra composition and a single-process CLI around those existing
+APIs. The current application data boundary is a versioned
+`GraphSampleCache`; ROOT preparation converts events through the shared
+feature and graph builders before writing that cache.
+
+Build thin new applications around tested library interfaces in this order:
+
+1. preprocessing/cache generation (`scripts/prep_data.py`);
+2. training/evaluation (`scripts/training_script.py`);
+3. inference to `.npz` and ROOT (`scripts/inference.py`);
+4. ONNX export after PyTorch parity (`gnn4colliders export`).
+
+Use subprocess integration tests with tiny fixtures. Preserve CLI options only
+where they serve an active workflow; document removed diagnostic/cluster-only
+options.
+
+## Phase 8 — reproducibility and deployment
+
+Task 14 adds the initial deployment boundary: CPU/GPU DDP through standard
+`torchrun` variables, rank-local graph-sample sharding, global metric/output
+gathering, rank-0 checkpoint/config writing, and Perlmutter-oriented Slurm
+examples. Evaluation deliberately avoids sampler padding duplicates. The
+remaining follow-up is a streaming or sharded output path for very large
+distributed inference jobs.
+
+The seed policy remains explicit: the configured seed is offset by rank for
+process-local randomness, while distributed sample assignment is derived from
+the configured seed, world size, and epoch. GPU kernel nondeterminism and
+exact per-rank RNG checkpoint replay remain environment-dependent. Slurm/NCCL,
+Podman-HPC, ROOT, and Hugging Face integrations stay in launcher/adapters
+rather than package code.
+
+## Checkpoint compatibility checklist
+
+- [x] Load a checked-in or generated multiclass pretrained checkpoint.
+- [x] Load a legacy fine-tuning checkpoint after prefix normalization.
+- [x] Resume optimizer and early-stop state.
+- [x] Produce equivalent logits on a deterministic graph fixture.
+- [x] Produce equivalent `.npz` score, label, and metadata fields.
+- [x] Preserve ROOT scalar/vector score branch conventions in the Python adapter.
+
+Known risks are documented in [`architecture.md`](architecture.md): edge order,
+self-loops, weight semantics, validation/test naming, padding, dynamic
+selection evaluation, reproducibility, and the experimental model/loss surface.
+
+## Migration closure status
+
+### Task 18 compatibility closure
+
+The compatibility boundary is now explicit in `gnn4colliders.compat`.
+Production ingestion stores named `EventMetadata`; legacy two-column tracking
+is converted only at the compatibility boundary. Checkpoint prefix cleanup and
+the historical ROOT-GNN `classify` to `classifier` mapping have one canonical
+implementation. The new checkpoint schema and named NPZ output remain
+canonical. See [`compatibility.md`](compatibility.md) for the supported and
+intentionally unsupported historical artifacts.
+
+The following matrix describes the supported new stack, rather than every
+class that exists in `legacy/`:
+
+| Legacy area | New-stack status | Notes |
+| --- | --- | --- |
+| ROOT/Awkward ingestion | migrated | `RootEventDataset` returns `EventSample` in file/event order |
+| node features | migrated + parity-tested | seven-column schema, `CALC_E`, ordering, scales, float32 |
+| edge construction | migrated + parity-tested | directed source-major topology and `[deta,dphi,dR]` |
+| graph cache | migrated | versioned `GraphSampleCache`; graph-level cache only |
+| folds and weights | migrated | named `EventMetadata.fold` and `.weight` |
+| batching | migrated | deterministic local loader and DDP sharding |
+| legacy padding modes | deferred | no active new-stack consumer |
+| `Edge_Network` | migrated + parity-tested | `EdgeNetwork`, raw logits |
+| transfer/fine-tuning | migrated + parity-tested | frozen or trainable backbone |
+| loss and metrics | migrated + parity-tested | task-owned weighted reductions and full-split AUC |
+| training lifecycle | migrated | `Trainer`, validation semantics, scheduler, early stopping |
+| checkpoints/resume | migrated | schema v1; historical weight/prefix adapter |
+| inference/NPZ | migrated | named output fields and ordered accumulation |
+| ROOT score output | compatibility adapter | Python API supported; CLI currently NPZ-only |
+| DDP | migrated | torchrun boundary, rank-0 artifacts, gathered metrics |
+| Slurm/Perlmutter | launcher examples | site policy remains outside package code |
+| ONNX export | migrated for ROOT-GNN | tensor-only adapter, ONNX Runtime validation, and `export` CLI; raw graph tensors are the input contract |
+
+### Intentional redesigns
+
+These are deliberate new-stack contracts, not accidental parity failures:
+
+* `tracking[:, 0]` and `tracking[:, 1]` become named `metadata.fold` and
+  `metadata.weight`; public consumers do not depend on positional columns.
+* Dynamic legacy YAML `module`/`class` construction becomes allow-listed
+  semantic Hydra configuration.
+* The monolithic training script becomes `Task` + `Trainer` + checkpoint and
+  inference adapters.
+* Graph state is scoped to the forward pass rather than relying on persistent
+  mutation of shared graph state.
+* Model constructors do not mutate global RNG state; seeding is explicit in
+  the training/application boundary.
+* Validation is the selection/early-stopping split and test is held out. This
+  corrects the legacy loader-name inversion.
+
+Compatibility preserves externally observable scientific behavior where it is
+validated; it does not promise to preserve every legacy implementation bug.
+The characterized legacy transfer path had a nonempty-global handling bug;
+the rewrite supports named globals. Negative weights, rare empty graphs,
+historical checkpoint variants, and legacy padding edge cases remain areas to
+audit when a supported consumer requires them.
+
+## ROOT-GNN v1 completion checklist
+
+- [x] active ROOT data path and graph cache
+- [x] validated feature, graph, model, task, and metric behavior
+- [x] train from scratch and fine-tune a pretrained backbone
+- [x] resume new-stack checkpoints and load supported historical weights
+- [x] evaluate and predict named outputs
+- [x] single-process and DDP application boundaries
+- [x] Perlmutter/Slurm launcher examples and profiling guidance
+- [x] ROOT-GNN ONNX export and CPU Runtime parity
+- [ ] streaming/sharded large-scale prediction output
+- [ ] removal of frozen legacy reference
+- [ ] ROOT-Transformer representation/model
+
+ROOT-GNN v1 is complete when the checked-in new stack can prepare active data,
+reproduce validated legacy behavior, train, transfer, resume, evaluate,
+predict, and run single-process or DDP workflows. The remaining unchecked
+items are intentionally deferred rather than undocumented promises.